diff --git a/.bazelrc b/.bazelrc index 17285afdb381018d0054e771475327b1f7ed9866..174109124256707aff9b8b6de58f6de22e0614fe 100644 --- a/.bazelrc +++ b/.bazelrc @@ -67,6 +67,7 @@ build:sycl_trisycl --define=using_sycl=true --define=using_trisycl=true build:gdr --define=with_gdr_support=true build:ngraph --define=with_ngraph_support=true build:verbs --define=with_verbs_support=true +build:numa --define=with_numa_support=true # Options to disable default on features build:noaws --define=no_aws_support=true diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b978f89f9e1d79dd4f7481711a59c2b94e8bf01b..73782143a3d4b1742f33bb96845ed300eedb6f50 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -55,24 +55,28 @@ TensorFlow coding style. #### General guidelines and philosophy for contribution -* Include unit tests when you contribute new features, as they help to - a) prove that your code works correctly, and b) guard against future breaking - changes to lower the maintenance cost. -* Bug fixes also generally require unit tests, because the presence of bugs - usually indicates insufficient test coverage. -* Keep API compatibility in mind when you change code in core TensorFlow, - e.g., code in [tensorflow/core](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core) and [tensorflow/python](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/python). - TensorFlow has reached version 1 and hence cannot make - non-backward-compatible API changes without a major release. Reviewers of your - pull request will comment on any API compatibility issues. -* When you contribute a new feature to TensorFlow, the maintenance burden is (by - default) transferred to the TensorFlow team. This means that benefit of the - contribution must be compared against the cost of maintaining the feature. -* Full new features (e.g., a new op implementing a cutting-edge algorithm) - typically will live in - [tensorflow/contrib](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib) - to get some airtime before decision is made regarding whether they are to be - migrated to the core. +* Include unit tests when you contribute new features, as they help to a) + prove that your code works correctly, and b) guard against future breaking + changes to lower the maintenance cost. +* Bug fixes also generally require unit tests, because the presence of bugs + usually indicates insufficient test coverage. +* Keep API compatibility in mind when you change code in core TensorFlow, + e.g., code in + [tensorflow/core](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core) + and + [tensorflow/python](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/python). + TensorFlow has reached version 1 and hence cannot make + non-backward-compatible API changes without a major release. Reviewers of + your pull request will comment on any API compatibility issues. +* When you contribute a new feature to TensorFlow, the maintenance burden is + (by default) transferred to the TensorFlow team. This means that benefit of + the contribution must be compared against the cost of maintaining the + feature. +* Full new features (e.g., a new op implementing a cutting-edge algorithm) + typically will live in + [tensorflow/addons](https://github.com/tensorflow/addons) to get some + airtime before decision is made regarding whether they are to be migrated to + the core. #### License diff --git a/configure.py b/configure.py index df7341357b274b11b162498f532206c79e676af5..68097b3bfafb5cd299f50bac28e1c35bcc71a836 100644 --- a/configure.py +++ b/configure.py @@ -1495,6 +1495,34 @@ def set_other_mpi_vars(environ_cp): 'Cannot find the MPI library file in %s/lib or %s/lib64 or %s/lib32' % (mpi_home, mpi_home, mpi_home)) +def system_specific_test_config(env): + """Add default test flags required for TF tests to bazelrc.""" + write_to_bazelrc('test --flaky_test_attempts=3') + write_to_bazelrc('test --test_size_filters=small,medium') + write_to_bazelrc( + 'test --test_tag_filters=-benchmark-test,-no_oss,-oss_serial') + write_to_bazelrc('test --build_tag_filters=-benchmark-test,-no_oss') + if is_windows(): + if env.get('TF_NEED_CUDA', None) == 1: + write_to_bazelrc( + 'test --test_tag_filters=-no_windows,-no_windows_gpu,-no_gpu') + write_to_bazelrc( + 'test --build_tag_filters=-no_windows,-no_windows_gpu,-no_gpu') + else: + write_to_bazelrc('test --test_tag_filters=-no_windows,-gpu') + write_to_bazelrc('test --build_tag_filters=-no_windows,-gpu') + elif is_macos(): + write_to_bazelrc('test --test_tag_filters=-gpu,-nomac,-no_mac') + write_to_bazelrc('test --build_tag_filters=-gpu,-nomac,-no_mac') + elif is_linux(): + if env.get('TF_NEED_CUDA', None) == 1: + write_to_bazelrc('test --test_tag_filters=-no_gpu') + write_to_bazelrc('test --build_tag_filters=-no_gpu') + write_to_bazelrc('test --test_env=LD_LIBRARY_PATH') + else: + write_to_bazelrc('test --test_tag_filters=-gpu') + write_to_bazelrc('test --build_tag_filters=-gpu') + def set_system_libs_flag(environ_cp): syslibs = environ_cp.get('TF_SYSTEM_LIBS', '') @@ -1519,6 +1547,9 @@ def set_windows_build_flags(environ_cp): write_to_bazelrc('build --config monolithic') # Suppress warning messages write_to_bazelrc('build --copt=-w --host_copt=-w') + # Fix winsock2.h conflicts + write_to_bazelrc( + 'build --copt=-DWIN32_LEAN_AND_MEAN --host_copt=-DWIN32_LEAN_AND_MEAN') # Output more verbose information when something goes wrong write_to_bazelrc('build --verbose_failures') # The host and target platforms are the same in Windows build. So we don't @@ -1705,6 +1736,8 @@ def main(): create_android_ndk_rule(environ_cp) create_android_sdk_rule(environ_cp) + system_specific_test_config(os.environ) + if get_var( environ_cp, 'TF_CONFIGURE_APPLE_BAZEL_RULES', 'Configure Bazel rules for Apple platforms', False, @@ -1721,6 +1754,7 @@ def main(): config_info_line('gdr', 'Build with GDR support.') config_info_line('verbs', 'Build with libverbs support.') config_info_line('ngraph', 'Build with Intel nGraph support.') + config_info_line('numa', 'Build with NUMA support.') config_info_line( 'dynamic_kernels', '(Experimental) Build kernels into separate shared objects.') diff --git a/tensorflow/BUILD b/tensorflow/BUILD index f53982f1efc9885cc12dcc672ad819c762aca378..db51979fbfec010b861db9d40bef2d94a8e9e024 100644 --- a/tensorflow/BUILD +++ b/tensorflow/BUILD @@ -16,6 +16,8 @@ exports_files([ ]) load("//tensorflow:tensorflow.bzl", "tf_cc_shared_object") +load("//tensorflow:tensorflow.bzl", "tf_custom_op_library_additional_deps_impl") +load("//tensorflow:tensorflow.bzl", "tf_native_cc_binary") load( "//tensorflow/core:platform/default/build_config.bzl", "tf_additional_binary_deps", @@ -304,6 +306,12 @@ config_setting( visibility = ["//visibility:public"], ) +config_setting( + name = "with_numa_support", + define_values = {"with_numa_support": "true"}, + visibility = ["//visibility:public"], +) + # Crosses between framework_shared_object and a bunch of other configurations # due to limitations in nested select() statements. config_setting( @@ -493,19 +501,26 @@ tf_cc_shared_object( # symbols in object files. tf_cc_shared_object( - name = "libtensorflow.so", + name = "tensorflow", linkopts = select({ "//tensorflow:darwin": [ "-Wl,-exported_symbols_list,$(location //tensorflow/c:exported_symbols.lds)", "-Wl,-install_name,@rpath/libtensorflow.so", ], - "//tensorflow:windows": [], + "//tensorflow:windows": [ + ], "//conditions:default": [ "-z defs", "-Wl,--version-script,$(location //tensorflow/c:version_script.lds)", ], }), + per_os_targets = True, visibility = ["//visibility:public"], + # add win_def_file for tensorflow + win_def_file = select({ + # We need this DEF file to properly export symbols on Windows "//tensorflow:windows": ":tensorflow_filtered_def_file", + "//conditions:default": None, + }), deps = [ "//tensorflow/c:c_api", "//tensorflow/c:c_api_experimental", @@ -517,7 +532,7 @@ tf_cc_shared_object( ) tf_cc_shared_object( - name = "libtensorflow_cc.so", + name = "tensorflow_cc", linkopts = select({ "//tensorflow:darwin": [ "-Wl,-exported_symbols_list,$(location //tensorflow:tf_exported_symbols.lds)", @@ -528,7 +543,13 @@ tf_cc_shared_object( "-Wl,--version-script,$(location //tensorflow:tf_version_script.lds)", ], }), + per_os_targets = True, visibility = ["//visibility:public"], + # add win_def_file for tensorflow_cc + win_def_file = select({ + # We need this DEF file to properly export symbols on Windows "//tensorflow:windows": ":tensorflow_filtered_def_file", + "//conditions:default": None, + }), deps = [ "//tensorflow:tf_exported_symbols.lds", "//tensorflow:tf_version_script.lds", @@ -542,6 +563,92 @@ tf_cc_shared_object( ] + if_ngraph(["@ngraph_tf//:ngraph_tf"]), ) +# ** Targets for Windows build (start) ** + +# Build a shared library (DLL) by cc_binary from tf_custom_op_library_additional_deps_impl, +# it contains all object code from its dependencies. +# This target is only used for parsing the symbols to be exported in tensorflow.dll. +# Do NOT depend on it. +tf_native_cc_binary( + name = "tf_custom_op_library_additional_deps.dll", + linkshared = 1, + linkstatic = 1, + deps = tf_custom_op_library_additional_deps_impl(), +) + +# Get a DEF file generated by parsing all object files +# of tf_custom_op_library_additional_deps.so +filegroup( + name = "tensorflow_def_file", + srcs = [":tf_custom_op_library_additional_deps.dll"], + output_group = "def_file", +) + +# Filter the DEF file to reduce the number of symbols to 64K or less. +# Note that we also write the name of the pyd file into DEF file so that +# the dynamic libraries of custom ops can find it at runtime. +genrule( + name = "tensorflow_filtered_def_file", + srcs = [":tensorflow_def_file"], + outs = ["tensorflow_filtered_def_file.def"], + cmd = select({ + "//tensorflow:windows": """ + $(location @local_config_def_file_filter//:def_file_filter) \\ + --input $(location :tensorflow_def_file) \\ + --output $@ + """, + "//conditions:default": "touch $@", # Just a placeholder for Unix platforms + }), + tools = ["@local_config_def_file_filter//:def_file_filter"], + visibility = ["//visibility:public"], +) + +# The interface library (tensorflow.dll.if.lib) for linking tensorflow DLL library (tensorflow.dll) on Windows. +# To learn more about import library (called interface library in Bazel): +# https://docs.microsoft.com/en-us/cpp/build/linking-an-executable-to-a-dll?view=vs-2017#linking-implicitly +filegroup( + name = "get_tensorflow_dll_import_lib", + srcs = ["//tensorflow:tensorflow.dll"], + output_group = "interface_library", + visibility = ["//visibility:public"], +) + +# Rename the import library for tensorflow.dll from tensorflow.dll.if.lib to tensorflow.lib +genrule( + name = "tensorflow_dll_import_lib", + srcs = [":get_tensorflow_dll_import_lib"], + outs = ["tensorflow.lib"], + cmd = select({ + "//tensorflow:windows": "cp -f $< $@", + "//conditions:default": "touch $@", # Just a placeholder for Unix platforms + }), + visibility = ["//visibility:public"], +) + +# The interface library (tensorflow_cc.dll.if.lib) for linking tensorflow DLL library (tensorflow_cc.dll) on Windows. +# To learn more about import library (called interface library in Bazel): +# https://docs.microsoft.com/en-us/cpp/build/linking-an-executable-to-a-dll?view=vs-2017#linking-implicitly +filegroup( + name = "get_tensorflow_cc_dll_import_lib", + srcs = ["//tensorflow:tensorflow_cc.dll"], + output_group = "interface_library", + visibility = ["//visibility:public"], +) + +# Rename the import library for tensorflow.dll from tensorflow_cc.dll.if.lib to tensorflow.lib +genrule( + name = "tensorflow_cc_dll_import_lib", + srcs = [":get_tensorflow_cc_dll_import_lib"], + outs = ["tensorflow_cc.lib"], + cmd = select({ + "//tensorflow:windows": "cp -f $< $@", + "//conditions:default": "touch $@", # Just a placeholder for Unix platforms + }), + visibility = ["//visibility:public"], +) + +# ** Targets for Windows build (end) ** + exports_files( [ "tf_version_script.lds", diff --git a/tensorflow/api_template.__init__.py b/tensorflow/api_template.__init__.py index ddcacfcbe2d4d8b089f10f1a771384dc8c4fd199..7bd6b7223989cddfea935f0ed2bcf7536015feea 100644 --- a/tensorflow/api_template.__init__.py +++ b/tensorflow/api_template.__init__.py @@ -122,5 +122,6 @@ if hasattr(_current_module, 'keras'): losses = keras.losses metrics = keras.metrics optimizers = keras.optimizers + initializers = keras.initializers # pylint: enable=undefined-variable diff --git a/tensorflow/c/BUILD b/tensorflow/c/BUILD index ef7863dc0d5cbd57da30baa6e04278c2a0354b25..22e283ee32fd6c66d19d6e52602a7cb7f02545c0 100644 --- a/tensorflow/c/BUILD +++ b/tensorflow/c/BUILD @@ -59,6 +59,7 @@ tf_cuda_library( "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:lib", + "//tensorflow/core:lib_platform", "//tensorflow/core:op_gen_lib", "//tensorflow/core/distributed_runtime:server_lib", ], diff --git a/tensorflow/c/c_api.cc b/tensorflow/c/c_api.cc index 245d7ba2b186895532953aa61ebfc3fc6bf635a7..af93d91b94c8b1a4bed9d3947b0fcc007529e2be 100644 --- a/tensorflow/c/c_api.cc +++ b/tensorflow/c/c_api.cc @@ -20,15 +20,19 @@ limitations under the License. #include #include -#ifndef __ANDROID__ +// Required for IS_MOBILE_PLATFORM +#include "tensorflow/core/platform/platform.h" // NOLINT + +#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) #include "tensorflow/cc/framework/gradients.h" #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/framework/scope_internal.h" #include "tensorflow/cc/ops/while_loop.h" #include "tensorflow/cc/saved_model/loader.h" +#include "tensorflow/core/distributed_runtime/server_lib.h" #include "tensorflow/core/framework/op_gen_lib.h" #include "tensorflow/core/kernels/logging_ops.h" -#endif +#endif // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) #include "tensorflow/c/c_api_internal.h" #include "tensorflow/core/common_runtime/device_mgr.h" #include "tensorflow/core/common_runtime/eval_const_tensor.h" @@ -1079,7 +1083,7 @@ TensorId ToTensorId(const TF_Output& output) { return TensorId(output.oper->node.name(), output.index); } -#ifndef __ANDROID__ +#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) std::vector OutputsFromTFOutputs(TF_Output* tf_outputs, int n) { std::vector outputs(n); @@ -1097,7 +1101,7 @@ void TFOutputsFromOutputs(const std::vector& outputs, tf_outputs[i].index = outputs[i].index(); } } -#endif // __ANDROID__ +#endif // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) } // namespace @@ -2219,7 +2223,7 @@ void TF_GraphImportGraphDef(TF_Graph* graph, const TF_Buffer* graph_def, namespace { -#ifndef __ANDROID__ +#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) // Creates a placeholder representing an input to the cond or body graph. // TODO(skyewm): remove these from final graph @@ -2313,7 +2317,7 @@ bool ValidateInputWhileParams(const TF_WhileParams& params, TF_Status* s) { return true; } -#endif // __ANDROID__ +#endif // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) void FreeWhileResources(const TF_WhileParams* params) { TF_DeleteGraph(params->cond_graph); @@ -2332,9 +2336,9 @@ TF_WhileParams EmptyWhileParams() { TF_WhileParams TF_NewWhile(TF_Graph* g, TF_Output* inputs, int ninputs, TF_Status* status) { -#ifdef __ANDROID__ +#if defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) status->status = tensorflow::errors::Unimplemented( - "Creating while loops is not supported in Android. File a bug at " + "Creating while loops is not supported on mobile. File a bug at " "https://github.com/tensorflow/tensorflow/issues if this feature is " "important to you"); return EmptyWhileParams(); @@ -2379,10 +2383,10 @@ TF_WhileParams TF_NewWhile(TF_Graph* g, TF_Output* inputs, int ninputs, return EmptyWhileParams(); } return params; -#endif // __ANDROID__ +#endif // defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) } -#ifndef __ANDROID__ +#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) namespace { // TODO(skyewm): make nodes in while loop unfetchable like in Python version @@ -2457,13 +2461,13 @@ void TF_FinishWhileHelper(const TF_WhileParams* params, TF_Status* status, } } // namespace -#endif // __ANDROID__ +#endif // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) void TF_FinishWhile(const TF_WhileParams* params, TF_Status* status, TF_Output* outputs) { -#ifdef __ANDROID__ +#if defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) status->status = tensorflow::errors::Unimplemented( - "Creating while loops is not supported in Android. File a bug at " + "Creating while loops is not supported on mobile. File a bug at " "https://github.com/tensorflow/tensorflow/issues if this feature is " "important to you"); #else @@ -2471,7 +2475,7 @@ void TF_FinishWhile(const TF_WhileParams* params, TF_Status* status, if (!ValidateConstWhileParams(*params, status)) return; TF_FinishWhileHelper(params, status, outputs); FreeWhileResources(params); -#endif // __ANDROID__ +#endif // defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) } void TF_AbortWhile(const TF_WhileParams* params) { FreeWhileResources(params); } @@ -2484,9 +2488,9 @@ void TF_AddGradients(TF_Graph* g, TF_Output* y, int ny, TF_Output* x, int nx, void TF_AddGradientsWithPrefix(TF_Graph* g, const char* prefix, TF_Output* y, int ny, TF_Output* x, int nx, TF_Output* dx, TF_Status* status, TF_Output* dy) { -#ifdef __ANDROID__ +#if defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) status->status = tensorflow::errors::Unimplemented( - "Adding gradients is not supported in Android. File a bug at " + "Adding gradients is not supported on mobile. File a bug at " "https://github.com/tensorflow/tensorflow/issues if this feature is " "important to you"); #else @@ -2566,7 +2570,7 @@ void TF_AddGradientsWithPrefix(TF_Graph* g, const char* prefix, TF_Output* y, // Unpack the results from grad_outputs_arg. TFOutputsFromOutputs(dy_arg, dy); -#endif // __ANDROID__ +#endif // defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) } // TF_Session functions ---------------------------------------------- @@ -2595,11 +2599,11 @@ TF_Session* TF_LoadSessionFromSavedModel( const TF_SessionOptions* session_options, const TF_Buffer* run_options, const char* export_dir, const char* const* tags, int tags_len, TF_Graph* graph, TF_Buffer* meta_graph_def, TF_Status* status) { -// TODO(ashankar): Remove the __ANDROID__ guard. This will require ensuring that -// the tensorflow/cc/saved_model:loader build target is Android friendly. -#ifdef __ANDROID__ +// TODO(sjr): Remove the IS_MOBILE_PLATFORM guard. This will require ensuring +// that the tensorflow/cc/saved_model:loader build target is mobile friendly. +#if defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) status->status = tensorflow::errors::Unimplemented( - "Loading a SavedModel is not supported in Android. File a bug at " + "Loading a SavedModel is not supported on mobile. File a bug at " "https://github.com/tensorflow/tensorflow/issues if this feature is " "important to you"); return nullptr; @@ -2651,7 +2655,7 @@ TF_Session* TF_LoadSessionFromSavedModel( graph->sessions[session] = ""; session->last_num_graph_nodes = graph->graph.num_node_ids(); return session; -#endif // __ANDROID__ +#endif // defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) } void TF_CloseSession(TF_Session* s, TF_Status* status) { @@ -2826,9 +2830,9 @@ void TF_DeleteApiDefMap(TF_ApiDefMap* apimap) { delete apimap; } void TF_ApiDefMapPut(TF_ApiDefMap* api_def_map, const char* text, size_t text_len, TF_Status* status) { -#ifdef __ANDROID__ +#if defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) status->status = tensorflow::errors::Unimplemented( - "ApiDefMap is not supported in Android."); + "ApiDefMap is not supported on mobile."); #else mutex_lock l(api_def_map->lock); if (api_def_map->update_docs_called) { @@ -2839,14 +2843,14 @@ void TF_ApiDefMapPut(TF_ApiDefMap* api_def_map, const char* text, } string api_def_text(text, text_len); status->status = api_def_map->api_def_map.LoadApiDef(api_def_text); -#endif // __ANDROID__ +#endif // defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) } TF_Buffer* TF_ApiDefMapGet(TF_ApiDefMap* api_def_map, const char* name, size_t name_len, TF_Status* status) { -#ifdef __ANDROID__ +#if defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) status->status = tensorflow::errors::Unimplemented( - "ApiDefMap is not supported in Android."); + "ApiDefMap is not supported on mobile."); return nullptr; #else mutex_lock l(api_def_map->lock); @@ -2867,7 +2871,7 @@ TF_Buffer* TF_ApiDefMapGet(TF_ApiDefMap* api_def_map, const char* name, return nullptr; } return ret; -#endif // __ANDROID__ +#endif // defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) } TF_Buffer* TF_GetAllRegisteredKernels(TF_Status* status) { @@ -2895,16 +2899,16 @@ TF_Buffer* TF_GetRegisteredKernelsForOp(const char* name, TF_Status* status) { // TF_Server functions ---------------------------------------------- -#ifndef __ANDROID__ +#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) TF_Server::TF_Server(std::unique_ptr server) : target(server->target()), server(std::move(server)) {} -#endif // __ANDROID__ +#endif // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) TF_Server* TF_NewServer(const void* proto, size_t proto_len, TF_Status* status) { -#ifdef __ANDROID__ +#if defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) status->status = tensorflow::errors::Unimplemented( - "Server functionality is not supported in Android"); + "Server functionality is not supported on mobile"); return nullptr; #else tensorflow::ServerDef server_def; @@ -2919,38 +2923,38 @@ TF_Server* TF_NewServer(const void* proto, size_t proto_len, if (!status->status.ok()) return nullptr; return new TF_Server(std::move(out_server)); -#endif +#endif // defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) } void TF_ServerStart(TF_Server* server, TF_Status* status) { -#ifdef __ANDROID__ +#if defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) status->status = tensorflow::errors::Unimplemented( - "Server functionality is not supported in Android"); + "Server functionality is not supported on mobile"); #else status->status = server->server->Start(); -#endif +#endif // defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) } void TF_ServerStop(TF_Server* server, TF_Status* status) { -#ifdef __ANDROID__ +#if defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) status->status = tensorflow::errors::Unimplemented( - "Server functionality is not supported in Android"); + "Server functionality is not supported on mobile"); #else status->status = server->server->Stop(); -#endif +#endif // defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) } void TF_ServerJoin(TF_Server* server, TF_Status* status) { -#ifdef __ANDROID__ +#if defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) status->status = tensorflow::errors::Unimplemented( - "Server functionality is not supported in Android"); + "Server functionality is not supported on mobile"); #else status->status = server->server->Join(); -#endif +#endif // defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) } const char* TF_ServerTarget(TF_Server* server) { -#ifdef __ANDROID__ +#if defined(IS_MOBILE_PLATFORM) || defined(IS_SLIM_BUILD) return nullptr; #else return server->target.c_str(); @@ -2958,15 +2962,15 @@ const char* TF_ServerTarget(TF_Server* server) { } void TF_DeleteServer(TF_Server* server) { -#ifndef __ANDROID__ +#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) delete server; -#endif +#endif // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) } void TF_RegisterLogListener(void (*listener)(const char*)) { -#ifndef __ANDROID__ +#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) tensorflow::logging::RegisterListener(listener); -#endif +#endif // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) } } // end extern "C" diff --git a/tensorflow/c/c_api_experimental_test.cc b/tensorflow/c/c_api_experimental_test.cc index c54021a7517ebbdd00405cbfa9cee8f3f6616cca..2c92e38f03a9d01d285f475b1a8996c44475c5c2 100644 --- a/tensorflow/c/c_api_experimental_test.cc +++ b/tensorflow/c/c_api_experimental_test.cc @@ -447,7 +447,7 @@ TEST_F(AddEagerOpToGraphTest, ListAttributesArePreserved) { } TEST_F(AddEagerOpToGraphTest, ListInputsAreAddedCorrectly) { - TFE_TensorHandle* scalar = TestScalarTensorHandle(); + TFE_TensorHandle* scalar = TestScalarTensorHandle(static_cast(1)); TFE_Op* identityn = TFE_NewOp(eager_ctx_, "IdentityN", status_); CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_); constexpr size_t kNumInputs = 3; diff --git a/tensorflow/c/c_api_internal.h b/tensorflow/c/c_api_internal.h index d520b6b76849e562def6abd8be0510d3b4797e8c..9a69c58718b3514287256124629f59443f38fd39 100644 --- a/tensorflow/c/c_api_internal.h +++ b/tensorflow/c/c_api_internal.h @@ -24,10 +24,12 @@ limitations under the License. #include #include -#ifndef __ANDROID__ -#include "tensorflow/core/distributed_runtime/server_lib.h" +// Required for IS_MOBILE_PLATFORM +#include "tensorflow/core/platform/platform.h" // NO_LINT + +#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) #include "tensorflow/core/framework/op_gen_lib.h" -#endif +#endif // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) #include "tensorflow/core/common_runtime/shape_refiner.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" @@ -42,6 +44,7 @@ limitations under the License. namespace tensorflow { class Device; class DeviceMgr; +class ServerInterface; } // namespace tensorflow // Internal structures used by the C API. These are likely to change and should @@ -167,27 +170,27 @@ struct TF_Function { struct TF_ApiDefMap { explicit TF_ApiDefMap(const tensorflow::OpList& op_list) : -#ifndef __ANDROID__ +#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) api_def_map(op_list), -#endif +#endif // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) update_docs_called(false) { } -#ifndef __ANDROID__ +#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) tensorflow::ApiDefMap api_def_map GUARDED_BY(lock); -#endif +#endif // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) bool update_docs_called GUARDED_BY(lock); tensorflow::mutex lock; }; -#ifndef __ANDROID__ +#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) struct TF_Server { TF_Server(std::unique_ptr server); const tensorflow::string target; std::unique_ptr server; }; -#endif +#endif // !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD) namespace tensorflow { diff --git a/tensorflow/c/eager/BUILD b/tensorflow/c/eager/BUILD index c1fd41b9f42292c3bb20d0943b70774e59582bbc..445b2cd25812e1d73fdd85b61f2a234150b880a6 100644 --- a/tensorflow/c/eager/BUILD +++ b/tensorflow/c/eager/BUILD @@ -147,6 +147,7 @@ tf_cuda_cc_test( ], deps = [ ":c_api", + ":c_api_internal", ":c_api_test_util", "//tensorflow/c:c_test_util", "//tensorflow/core:lib", @@ -211,6 +212,7 @@ tf_cuda_library( "//tensorflow/core/distributed_runtime:server_lib", "//tensorflow/core/distributed_runtime:worker_env", "//tensorflow/core/profiler/rpc:profiler_server", + "//tensorflow/core/profiler/rpc/client:capture_profile", "//tensorflow/core:gpu_runtime", ], ) diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc index 45701c7fcf02d5e6ec464ae10d4d20f20ba1d9f0..9509135e239d1feccd518ea5190b657b05eae6cf 100755 --- a/tensorflow/c/eager/c_api.cc +++ b/tensorflow/c/eager/c_api.cc @@ -234,6 +234,78 @@ tensorflow::Status UpdateTFE_ContextWithServerDef( return tensorflow::Status::OK(); #undef LOG_AND_RETURN_IF_ERROR } + +tensorflow::Status OpInferSingleInputAttrs(TFE_Op* op, + TFE_TensorHandle* input) { + TFE_OpInferenceContext* ictx = op->inference_ctx.get(); + const auto& input_def = ictx->op_def->input_arg(ictx->input_arg_idx++); + if (!input_def.number_attr().empty() || !input_def.type_list_attr().empty()) { + // Some clients that are still setting their input attributes manually are + // adding input list to their op by calling `TFE_OpAddInput` for each of + // its elements instead of calling `TFE_OpAddInputList`. When this happens, + // we cannot detect the end of such list, thus lose track of the input + // arguments in the op definition. To guarantee backward compatibility with + // those clients, disable automatic inference in this case. + op->inference_ctx.reset(nullptr); + return tensorflow::Status::OK(); + } + const std::string& type_attr = input_def.type_attr(); + if (!type_attr.empty() && ictx->attrs.find(type_attr) == ictx->attrs.end()) { + op->operation.MutableAttrs()->Set(type_attr, input->handle->dtype); + ictx->attrs.insert(type_attr); + } + return tensorflow::Status::OK(); +} + +void OpInferSingleTypeInputListAttrs(TFE_Op* op, + const tensorflow::OpDef::ArgDef& input_def, + TFE_TensorHandle** inputs, + int num_inputs) { + TFE_OpInferenceContext* ictx = op->inference_ctx.get(); + if (ictx->attrs.find(input_def.number_attr()) == ictx->attrs.end()) { + op->operation.MutableAttrs()->Set(input_def.number_attr(), num_inputs); + ictx->attrs.insert(input_def.number_attr()); + } + if (ictx->attrs.find(input_def.type_attr()) == ictx->attrs.end()) { + op->operation.MutableAttrs()->Set(input_def.type_attr(), + inputs[0]->handle->dtype); + ictx->attrs.insert(input_def.type_attr()); + } +} + +void OpInferMixedTypeInputListAttrs(TFE_Op* op, + const tensorflow::OpDef::ArgDef& input_def, + TFE_TensorHandle** inputs, int num_inputs) { + TFE_OpInferenceContext* ictx = op->inference_ctx.get(); + if (ictx->attrs.find(input_def.type_list_attr()) == ictx->attrs.end()) { + std::unique_ptr dtypes( + new tensorflow::DataType[num_inputs]); + for (int i = 0; i < num_inputs; ++i) { + dtypes[i] = inputs[i]->handle->dtype; + } + op->operation.MutableAttrs()->Set( + input_def.type_list_attr(), + tensorflow::gtl::ArraySlice(dtypes.get(), + num_inputs)); + ictx->attrs.insert(input_def.type_list_attr()); + } +} + +tensorflow::Status OpInferInputListAttrs(TFE_Op* op, TFE_TensorHandle** inputs, + int num_inputs) { + TFE_OpInferenceContext* ictx = op->inference_ctx.get(); + const auto& input_def = ictx->op_def->input_arg(ictx->input_arg_idx++); + if (!input_def.type_list_attr().empty()) { + OpInferMixedTypeInputListAttrs(op, input_def, inputs, num_inputs); + } else if (!input_def.type_attr().empty() && + !input_def.number_attr().empty()) { + OpInferSingleTypeInputListAttrs(op, input_def, inputs, num_inputs); + } else { + return tensorflow::errors::InvalidArgument("Invalid input list definition"); + } + return tensorflow::Status::OK(); +} + } // namespace extern "C" { @@ -249,6 +321,7 @@ void TFE_ContextOptionsSetAsync(TFE_ContextOptions* options, unsigned char enable) { options->async = enable; } + void TFE_ContextOptionsSetDevicePlacementPolicy( TFE_ContextOptions* options, TFE_ContextDevicePlacementPolicy policy) { options->policy = policy; @@ -492,20 +565,29 @@ TFE_Op* TFE_NewOp(TFE_Context* ctx, const char* op_or_function_name, const tensorflow::AttrTypeMap* types; bool is_function = false; status->status = tensorflow::AttrTypeMapForOp(name, &types, &is_function); - if (status->status.ok()) { - if (is_function && !ctx->context.FindFunctionByName(name)) { - status->status = tensorflow::errors::NotFound( - "'", name, - "' is neither a type of a primitive operation nor a name " - "of a function registered in binary running on ", - tensorflow::port::Hostname(), - ". Make sure the operation or function is " - "registered in the binary running in this process."); + if (!status->status.ok()) { + return nullptr; + } + if (!is_function) { + const tensorflow::OpDef* op_def; + status->status = tensorflow::OpDefForOp(op_or_function_name, &op_def); + if (!status->status.ok()) { return nullptr; } - return new TFE_Op(ctx, name, is_function, types); + return new TFE_Op(ctx, name, false, types, + new TFE_OpInferenceContext(op_def)); } - return nullptr; + if (!ctx->context.FindFunctionByName(name)) { + status->status = tensorflow::errors::NotFound( + "'", name, + "' is neither a type of a primitive operation nor a name " + "of a function registered in binary running on ", + tensorflow::port::Hostname(), + ". Make sure the operation or function is " + "registered in the binary running in this process."); + return nullptr; + } + return new TFE_Op(ctx, name, true, types, nullptr); } void TFE_DeleteOp(TFE_Op* op) { delete op; } @@ -529,8 +611,21 @@ void TFE_OpSetXLACompilation(TFE_Op* op, unsigned char enable) { #endif // TENSORFLOW_EAGER_USE_XLA } -void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status) { - op->operation.AddInput(h->handle); +void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* input, TF_Status* status) { + op->operation.AddInput(input->handle); + if (op->inference_ctx) { + status->status = OpInferSingleInputAttrs(op, input); + } +} + +void TFE_OpAddInputList(TFE_Op* op, TFE_TensorHandle** inputs, int num_inputs, + TF_Status* status) { + for (int i = 0; i < num_inputs; ++i) { + op->operation.AddInput(inputs[i]->handle); + } + if (op->inference_ctx) { + status->status = OpInferInputListAttrs(op, inputs, num_inputs); + } } TF_AttrType TFE_OpGetAttrType(TFE_Op* op, const char* attr_name, diff --git a/tensorflow/c/eager/c_api.h b/tensorflow/c/eager/c_api.h index 044dfb7415b027b707af05a197fdb41fe1f6d2e5..ce3da7f91893864c90039c40fdcb86d0cbe58f0d 100755 --- a/tensorflow/c/eager/c_api.h +++ b/tensorflow/c/eager/c_api.h @@ -282,9 +282,14 @@ TF_CAPI_EXPORT extern const char* TFE_OpGetDevice(TFE_Op* op, TF_CAPI_EXPORT extern void TFE_OpSetXLACompilation(TFE_Op* op, unsigned char enable); -TF_CAPI_EXPORT extern void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* h, +TF_CAPI_EXPORT extern void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* input, TF_Status* status); +TF_CAPI_EXPORT extern void TFE_OpAddInputList(TFE_Op* op, + TFE_TensorHandle** inputs, + int num_inputs, + TF_Status* status); + TF_CAPI_EXPORT extern TF_AttrType TFE_OpGetAttrType(TFE_Op* op, const char* attr_name, unsigned char* is_list, diff --git a/tensorflow/c/eager/c_api_debug_test.cc b/tensorflow/c/eager/c_api_debug_test.cc index cddb9f6e00e9d639026f4bbe061d58f76771c0a9..4e987c745ecabd85c89d039468eb94ed51b4d00f 100644 --- a/tensorflow/c/eager/c_api_debug_test.cc +++ b/tensorflow/c/eager/c_api_debug_test.cc @@ -21,7 +21,7 @@ limitations under the License. #include "tensorflow/core/platform/test.h" TEST(CApiDebug, ScalarCPU) { - TFE_TensorHandle* h = TestScalarTensorHandle(); + TFE_TensorHandle* h = TestScalarTensorHandle(1.0f); TF_Status* status = TF_NewStatus(); TFE_TensorDebugInfo* debug_info = TFE_TensorHandleTensorDebugInfo(h, status); CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); diff --git a/tensorflow/c/eager/c_api_experimental.cc b/tensorflow/c/eager/c_api_experimental.cc index af7f1bbf8aa5636d78c222f5ba95624054273c47..c6a12247ef1aeea674b2b18db5986ac4b9de8e9c 100644 --- a/tensorflow/c/eager/c_api_experimental.cc +++ b/tensorflow/c/eager/c_api_experimental.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/c/c_api.h" #include "tensorflow/c/eager/c_api_internal.h" +#include "tensorflow/core/profiler/rpc/client/capture_profile.h" #include "tensorflow/core/profiler/rpc/profiler_server.h" using tensorflow::string; @@ -76,3 +77,18 @@ void TFE_ContextEnableGraphCollection(TFE_Context* ctx) { void TFE_ContextDisableGraphCollection(TFE_Context* ctx) { ctx->context.SetShouldStoreGraphs(false); } + +bool TFE_ProfilerClientStartTracing(const char* service_addr, + const char* logdir, const char* worker_list, + bool include_dataset_ops, int duration_ms, + int num_tracing_attempts) { + tensorflow::Status s = + tensorflow::profiler::client::ValidateHostPortPair(service_addr); + if (!s.ok()) { + return false; + } + s = tensorflow::profiler::client::StartTracing( + service_addr, logdir, worker_list, include_dataset_ops, duration_ms, + num_tracing_attempts); + return s.ok(); +} diff --git a/tensorflow/c/eager/c_api_experimental.h b/tensorflow/c/eager/c_api_experimental.h index eb57077e6834354005fbf7913cf5f51db3087b07..219b9f40720a4fc212bd6e191b5bb441cf2abeb8 100644 --- a/tensorflow/c/eager/c_api_experimental.h +++ b/tensorflow/c/eager/c_api_experimental.h @@ -75,6 +75,18 @@ TF_CAPI_EXPORT extern void TFE_ContextEnableGraphCollection(TFE_Context* ctx); // this context. TF_CAPI_EXPORT extern void TFE_ContextDisableGraphCollection(TFE_Context* ctx); +// Send a grpc request to profiler server (service_addr) to perform on-demand +// profiling and save the result into logdir which can be visualized by +// TensorBoard. worker_list is the list of worker TPUs separated by ','. Set +// include_dataset_opts to false to profile longer traces. It will block the +// caller thread until receives tracing result. +// This API is designed for TensorBoard, for end user, please use +// tensorflow/contrib/tpu/profiler/capture_tpu_profile instead following +// https://cloud.google.com/tpu/docs/cloud-tpu-tools#capture_trace. +TF_CAPI_EXPORT extern bool TFE_ProfilerClientStartTracing( + const char* service_addr, const char* logdir, const char* worker_list, + bool include_dataset_ops, int duration_ms, int num_tracing_attempts); + #ifdef __cplusplus } /* end extern "C" */ #endif diff --git a/tensorflow/c/eager/c_api_internal.h b/tensorflow/c/eager/c_api_internal.h index a563e4b8f50f2a90497736f4cb9ca234400bfa04..35dafb9a7f14bfe1fad21bda35685598164f3895 100644 --- a/tensorflow/c/eager/c_api_internal.h +++ b/tensorflow/c/eager/c_api_internal.h @@ -99,12 +99,24 @@ struct TFE_TensorDebugInfo { std::vector dev_dims; }; +struct TFE_OpInferenceContext { + explicit TFE_OpInferenceContext(const tensorflow::OpDef* op_def) + : op_def(op_def) {} + + const tensorflow::OpDef* op_def; // op definition from protobuf + int input_arg_idx = 0; // arg definition index for the next input to be added + tensorflow::gtl::FlatSet attrs; // attributes inferred so far +}; + struct TFE_Op { TFE_Op(TFE_Context* ctx, const char* op, bool is_function, - const tensorflow::AttrTypeMap* t) - : operation(&ctx->context, op, is_function, t) {} + const tensorflow::AttrTypeMap* t, + TFE_OpInferenceContext* inference_ctx) + : operation(&ctx->context, op, is_function, t), + inference_ctx(inference_ctx) {} tensorflow::EagerOperation operation; + std::unique_ptr inference_ctx; }; struct TFE_ProfilerContext { diff --git a/tensorflow/c/eager/c_api_test.cc b/tensorflow/c/eager/c_api_test.cc index 3d1ca4fb4b561a03ea9d879b1876fb1fd08a3139..b5e55420016bc9015ab71a515299838be953f5f4 100644 --- a/tensorflow/c/eager/c_api_test.cc +++ b/tensorflow/c/eager/c_api_test.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/c/eager/c_api.h" +#include "tensorflow/c/eager/c_api_internal.h" #include #include "absl/strings/match.h" @@ -1626,4 +1627,158 @@ TEST(CAPI, TestTFE_TensorHandleCopySharingUnderlyingTensorHandle) { TFE_DeleteTensorHandle(h); TFE_DeleteTensorHandle(h_shares_tensor); } + +TEST(CAPI, TestTFE_OpInferSingleInputAttrs) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + TFE_TensorHandle* input = TestMatrixTensorHandle(); + TFE_TensorHandle* axis = TestAxisTensorHandle(); + TFE_Op* minOp = TFE_NewOp(ctx, "Min", status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpAddInput(minOp, input, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpAddInput(minOp, axis, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + tensorflow::AttrValueMap attr_values; + minOp->operation.Attrs().FillAttrValueMap(&attr_values); + tensorflow::AttrValueMap::const_iterator attr_found = attr_values.find("T"); + EXPECT_NE(attr_found, attr_values.cend()); + EXPECT_EQ(attr_found->second.type(), tensorflow::DataType::DT_FLOAT); + attr_found = attr_values.find("Tidx"); + EXPECT_NE(attr_found, attr_values.cend()); + EXPECT_EQ(attr_found->second.type(), tensorflow::DataType::DT_INT32); + + TFE_TensorHandle* retvals[1] = {nullptr}; + int num_retvals = 1; + TFE_Execute(minOp, &retvals[0], &num_retvals, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + TF_DeleteStatus(status); + TFE_DeleteOp(minOp); + TFE_DeleteTensorHandle(input); + TFE_DeleteTensorHandle(axis); + TFE_DeleteTensorHandle(retvals[0]); + TFE_DeleteContext(ctx); +} + +TEST(CAPI, TestTFE_OpInferSingleTypeInputListAttrs) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + TFE_TensorHandle* input1 = TestMatrixTensorHandle(); + TFE_TensorHandle* input2 = TestMatrixTensorHandle(); + TFE_TensorHandle* dim = TestScalarTensorHandle(0); + TFE_Op* concatOp = TFE_NewOp(ctx, "Concat", status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_TensorHandle* inputs[] = {input1, input2}; + TFE_OpAddInput(concatOp, dim, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpAddInputList(concatOp, inputs, 2, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + tensorflow::AttrValueMap attr_values; + concatOp->operation.Attrs().FillAttrValueMap(&attr_values); + tensorflow::AttrValueMap::const_iterator attr_found = attr_values.find("T"); + EXPECT_NE(attr_found, attr_values.cend()); + EXPECT_EQ(attr_found->second.type(), tensorflow::DataType::DT_FLOAT); + attr_found = attr_values.find("N"); + EXPECT_NE(attr_found, attr_values.cend()); + EXPECT_EQ(attr_found->second.i(), 2); + + TFE_TensorHandle* retvals[1] = {nullptr}; + int num_retvals = 1; + TFE_Execute(concatOp, &retvals[0], &num_retvals, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + TF_DeleteStatus(status); + TFE_DeleteOp(concatOp); + TFE_DeleteTensorHandle(input1); + TFE_DeleteTensorHandle(input2); + TFE_DeleteTensorHandle(retvals[0]); + TFE_DeleteTensorHandle(dim); + TFE_DeleteContext(ctx); +} + +TEST(CAPI, TestTFE_OpInferMixedTypeInputListAttrs) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + TFE_TensorHandle* condition = TestScalarTensorHandle(true); + TFE_TensorHandle* t1 = TestMatrixTensorHandle(); + TFE_TensorHandle* t2 = TestAxisTensorHandle(); + TFE_Op* assertOp = TFE_NewOp(ctx, "Assert", status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpAddInput(assertOp, condition, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_TensorHandle* data[] = {condition, t1, t2}; + TFE_OpAddInputList(assertOp, data, 3, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + tensorflow::AttrValueMap attr_values; + assertOp->operation.Attrs().FillAttrValueMap(&attr_values); + tensorflow::AttrValueMap::const_iterator attr_found = attr_values.find("T"); + EXPECT_NE(attr_found, attr_values.cend()); + EXPECT_EQ(attr_found->second.list().type(0), tensorflow::DataType::DT_BOOL); + EXPECT_EQ(attr_found->second.list().type(1), tensorflow::DataType::DT_FLOAT); + EXPECT_EQ(attr_found->second.list().type(2), tensorflow::DataType::DT_INT32); + + TFE_TensorHandle* retvals[1] = {nullptr}; + int num_retvals = 1; + TFE_Execute(assertOp, &retvals[0], &num_retvals, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + TF_DeleteStatus(status); + TFE_DeleteOp(assertOp); + TFE_DeleteTensorHandle(condition); + TFE_DeleteTensorHandle(t1); + TFE_DeleteTensorHandle(t2); + TFE_DeleteTensorHandle(retvals[0]); + TFE_DeleteContext(ctx); +} + +TEST(CAPI, TestTFE_OpAttrsInferenceDisabledWhenNotCallingOpAddInputList) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + TFE_TensorHandle* input1 = TestMatrixTensorHandle(); + TFE_TensorHandle* input2 = TestMatrixTensorHandle(); + TFE_TensorHandle* dim = TestScalarTensorHandle(0); + TFE_Op* concatOp = TFE_NewOp(ctx, "Concat", status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_TensorHandle* inputs[] = {input1, input2}; + TFE_OpAddInput(concatOp, dim, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + CHECK(concatOp->inference_ctx); + TFE_OpAddInput(concatOp, inputs[0], status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + EXPECT_FALSE(concatOp->inference_ctx) << "Inference context is still present"; + TFE_OpAddInput(concatOp, inputs[1], status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + tensorflow::AttrValueMap attr_values; + concatOp->operation.Attrs().FillAttrValueMap(&attr_values); + EXPECT_EQ(attr_values.find("T"), attr_values.end()); + EXPECT_EQ(attr_values.find("N"), attr_values.end()); + + TF_DeleteStatus(status); + TFE_DeleteOp(concatOp); + TFE_DeleteTensorHandle(input1); + TFE_DeleteTensorHandle(input2); + TFE_DeleteTensorHandle(dim); + TFE_DeleteContext(ctx); +} } // namespace diff --git a/tensorflow/c/eager/c_api_test_util.cc b/tensorflow/c/eager/c_api_test_util.cc index bd38127d50c171af801dd1b937acefdba491b4a6..17d17c0b7f7909e8dc1aaea61ade2cce1c466a3f 100644 --- a/tensorflow/c/eager/c_api_test_util.cc +++ b/tensorflow/c/eager/c_api_test_util.cc @@ -21,8 +21,8 @@ limitations under the License. using tensorflow::string; -TFE_TensorHandle* TestScalarTensorHandle() { - float data[] = {1.0f}; +TFE_TensorHandle* TestScalarTensorHandle(float value) { + float data[] = {value}; TF_Tensor* t = TF_AllocateTensor(TF_FLOAT, nullptr, 0, sizeof(float)); memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t)); TF_Status* status = TF_NewStatus(); @@ -33,6 +33,30 @@ TFE_TensorHandle* TestScalarTensorHandle() { return th; } +TFE_TensorHandle* TestScalarTensorHandle(int value) { + int data[] = {value}; + TF_Tensor* t = TF_AllocateTensor(TF_INT32, nullptr, 0, sizeof(int)); + memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t)); + TF_Status* status = TF_NewStatus(); + TFE_TensorHandle* th = TFE_NewTensorHandle(t, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteTensor(t); + TF_DeleteStatus(status); + return th; +} + +TFE_TensorHandle* TestScalarTensorHandle(bool value) { + bool data[] = {value}; + TF_Tensor* t = TF_AllocateTensor(TF_BOOL, nullptr, 0, sizeof(bool)); + memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t)); + TF_Status* status = TF_NewStatus(); + TFE_TensorHandle* th = TFE_NewTensorHandle(t, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteTensor(t); + TF_DeleteStatus(status); + return th; +} + TFE_TensorHandle* DoubleTestMatrixTensorHandle() { int64_t dims[] = {2, 2}; double data[] = {1.0, 2.0, 3.0, 4.0}; diff --git a/tensorflow/c/eager/c_api_test_util.h b/tensorflow/c/eager/c_api_test_util.h index 75ef9459e93b4f2ed471c423a34565594efc1714..4ff3ff4301f63c001bec1eb23fb9e098b78c6a5e 100644 --- a/tensorflow/c/eager/c_api_test_util.h +++ b/tensorflow/c/eager/c_api_test_util.h @@ -20,7 +20,13 @@ limitations under the License. #include "tensorflow/core/platform/types.h" // Return a tensor handle containing a float scalar -TFE_TensorHandle* TestScalarTensorHandle(); +TFE_TensorHandle* TestScalarTensorHandle(float value); + +// Return a tensor handle containing a int scalar +TFE_TensorHandle* TestScalarTensorHandle(int value); + +// Return a tensor handle containing a bool scalar +TFE_TensorHandle* TestScalarTensorHandle(bool value); // Return a tensor handle containing a 2x2 matrix of doubles TFE_TensorHandle* DoubleTestMatrixTensorHandle(); diff --git a/tensorflow/c/python_api.cc b/tensorflow/c/python_api.cc index 98d8393332269ae349cf8aa5c0b612c6f17172e6..6449e7f44f72e9475d3c4eaabeb9e506bfd794f0 100644 --- a/tensorflow/c/python_api.cc +++ b/tensorflow/c/python_api.cc @@ -41,6 +41,15 @@ void SetAttr(TF_Graph* graph, TF_Operation* op, const char* attr_name, RecordMutation(graph, *op, "setting attribute"); } +void ClearAttr(TF_Graph* graph, TF_Operation* op, const char* attr_name, + TF_Status* status) { + AttrValue attr_val; + + mutex_lock l(graph->mu); + op->node.ClearAttr(attr_name); + RecordMutation(graph, *op, "clearing attribute"); +} + void SetRequestedDevice(TF_Graph* graph, TF_Operation* op, const char* device) { mutex_lock l(graph->mu); op->node.set_requested_device(device); diff --git a/tensorflow/c/python_api.h b/tensorflow/c/python_api.h index 44779ca656165dd65590cb5e9ea3ccf71165ed63..f26c0cb2ae2f6e00a247660a02525901e87920cd 100644 --- a/tensorflow/c/python_api.h +++ b/tensorflow/c/python_api.h @@ -32,6 +32,11 @@ void AddControlInput(TF_Graph* graph, TF_Operation* op, TF_Operation* input); void SetAttr(TF_Graph* graph, TF_Operation* op, const char* attr_name, TF_Buffer* attr_value_proto, TF_Status* status); +// Clears the attr in the node_def Protocol Buffer and sets a status upon +// completion. +void ClearAttr(TF_Graph* graph, TF_Operation* op, const char* attr_name, + TF_Status* status); + void SetRequestedDevice(TF_Graph* graph, TF_Operation* op, const char* device); // Updates 'dst' to consume 'new_src'. diff --git a/tensorflow/cc/framework/gradients.cc b/tensorflow/cc/framework/gradients.cc index d5ba56ea4ab502d2746c0e3625320b097e62ff2d..303fdf64ec723864848096009a57dabda2fc93e4 100644 --- a/tensorflow/cc/framework/gradients.cc +++ b/tensorflow/cc/framework/gradients.cc @@ -167,7 +167,6 @@ Status SymbolicGradientBuilder::BackpropAlongEdge(const Output& dst_grad, std::vector SymbolicGradientBuilder::GetReachableNodes() { std::vector reachable_nodes(scope_.graph()->num_node_ids(), false); std::deque queue; - std::vector visited(scope_.graph()->num_node_ids(), false); for (const Output& out : outputs_) { if (!reachable_nodes[out.node()->id()]) { queue.push_back(out.node()); @@ -180,10 +179,10 @@ std::vector SymbolicGradientBuilder::GetReachableNodes() { queue.pop_front(); for (const Edge* e : n->in_edges()) { if (e->IsControlEdge()) continue; - if (visited[e->src()->id()]) continue; - queue.push_back(e->src()); - reachable_nodes[e->src()->id()] = true; - visited[e->src()->id()] = true; + if (!reachable_nodes[e->src()->id()]) { + queue.push_back(e->src()); + reachable_nodes[e->src()->id()] = true; + } } } return reachable_nodes; @@ -201,9 +200,9 @@ std::unordered_set SymbolicGradientBuilder::GetStopBackpropNodes( // `output_` node was encountered, pair.second will be nullptr. std::deque> queue; for (const Output& nout : inputs_) { - if (visited.find(nout.node()) == visited.end()) { + auto const& pair = visited.insert(nout.node()); + if (pair.second) { queue.push_back(std::make_pair(nout.node(), static_cast(nullptr))); - visited.insert(nout.node()); } } // BFS from nodes in 'inputs_' along out edges for the entire graph. Internal @@ -217,22 +216,23 @@ std::unordered_set SymbolicGradientBuilder::GetStopBackpropNodes( for (const Edge* e : n->out_edges()) { // If a node is not reachable from outputs_, we can stop. if (e->IsControlEdge() || !reachable_nodes[e->dst()->id()]) continue; - if (visited.find(e->dst()) != visited.end()) continue; - - int node_id = e->dst()->id(); - Node* last_output_node = p.second; - if (output_nodes.find(node_id) != output_nodes.end()) { - // We reached an output node. - if (last_output_node != nullptr) { - // If we had already found an output node on this path so we mark - // it as an internal output. - internal_outputs.insert(last_output_node->id()); + + auto const& pair = visited.insert(e->dst()); + if (pair.second) { + int node_id = e->dst()->id(); + Node* last_output_node = p.second; + if (output_nodes.find(node_id) != output_nodes.end()) { + // We reached an output node. + if (last_output_node != nullptr) { + // If we had already found an output node on this path so we mark + // it as an internal output. + internal_outputs.insert(last_output_node->id()); + } + // Mark this newly found output node to insert in the queue. + last_output_node = e->dst(); } - // Mark this newly found output node to insert in the queue. - last_output_node = e->dst(); + queue.push_back(std::make_pair(e->dst(), last_output_node)); } - queue.push_back(std::make_pair(e->dst(), last_output_node)); - visited.insert(e->dst()); } } // Finally, we set stop_backprop_nodes to all output_nodes that aren't also @@ -286,9 +286,9 @@ Status SymbolicGradientBuilder::Initialize() { std::unordered_set visited; std::deque queue; for (const Output& nout : inputs_) { - if (visited.find(nout.node()) == visited.end()) { + auto const& pair = visited.insert(nout.node()); + if (pair.second) { queue.push_back(nout.node()); - visited.insert(nout.node()); } } @@ -309,9 +309,9 @@ Status SymbolicGradientBuilder::Initialize() { // we don't expect it to receive a backpropagated gradient. // It will not be counted in num_expected_backprops. if (e->IsControlEdge() || !reachable_nodes[e->dst()->id()]) continue; - if (visited.find(e->dst()) == visited.end()) { + auto const& pair = visited.insert(e->dst()); + if (pair.second) { queue.push_back(e->dst()); - visited.insert(e->dst()); } ++num_expected_backprops; } diff --git a/tensorflow/cc/gradients/image_grad.cc b/tensorflow/cc/gradients/image_grad.cc index 05c287bdc62cdb8be7208ce3975f280aaa816766..7d0f63efbcc9c217b5e86fafa4573908f803c8fb 100644 --- a/tensorflow/cc/gradients/image_grad.cc +++ b/tensorflow/cc/gradients/image_grad.cc @@ -29,13 +29,17 @@ Status ResizeNearestNeighborGradHelper(const Scope& scope, const Operation& op, bool align_corners; TF_RETURN_IF_ERROR( GetNodeAttr(op.node()->attrs(), "align_corners", &align_corners)); + bool half_pixel_centers; + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "half_pixel_centers", + &half_pixel_centers)); // The internal gradient implementation needs the shape of the input image. // x_shape = shape(x)[1:3] // = slice(shape(x), {1}, {3 - 1}) auto x_shape = Slice(scope, Shape(scope, op.input(0)), {1}, {2}); grad_outputs->push_back(internal::ResizeNearestNeighborGrad( scope, grad_inputs[0], x_shape, - internal::ResizeNearestNeighborGrad::AlignCorners(align_corners))); + internal::ResizeNearestNeighborGrad::AlignCorners(align_corners) + .HalfPixelCenters(half_pixel_centers))); grad_outputs->push_back(NoGradient()); return scope.status(); } @@ -47,9 +51,13 @@ Status ResizeBilinearGradHelper(const Scope& scope, const Operation& op, bool align_corners; TF_RETURN_IF_ERROR( GetNodeAttr(op.node()->attrs(), "align_corners", &align_corners)); + bool half_pixel_centers; + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "half_pixel_centers", + &half_pixel_centers)); grad_outputs->push_back(internal::ResizeBilinearGrad( scope, grad_inputs[0], op.input(0), - internal::ResizeBilinearGrad::AlignCorners(align_corners))); + internal::ResizeBilinearGrad::AlignCorners(align_corners) + .HalfPixelCenters(half_pixel_centers))); grad_outputs->push_back(NoGradient()); return scope.status(); } @@ -61,9 +69,14 @@ Status ResizeBicubicGradHelper(const Scope& scope, const Operation& op, bool align_corners; TF_RETURN_IF_ERROR( GetNodeAttr(op.node()->attrs(), "align_corners", &align_corners)); + bool half_pixel_centers; + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "half_pixel_centers", + &half_pixel_centers)); + grad_outputs->push_back(internal::ResizeBicubicGrad( scope, grad_inputs[0], op.input(0), - internal::ResizeBicubicGrad::AlignCorners(align_corners))); + internal::ResizeBicubicGrad::AlignCorners(align_corners) + .HalfPixelCenters(half_pixel_centers))); grad_outputs->push_back(NoGradient()); return scope.status(); } @@ -86,6 +99,25 @@ Status ScaleAndTranslateGradHelper(const Scope& scope, const Operation& op, } REGISTER_GRADIENT_OP("ScaleAndTranslate", ScaleAndTranslateGradHelper); +Status CropAndResizeGradHelper(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + DataType input_type; + string method; + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "method", &method)); + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "T", &input_type)); + auto image_shape = Shape(scope, op.input(0)); + grad_outputs->push_back(CropAndResizeGradImage( + scope, grad_inputs[0], op.input(1), op.input(2), image_shape, input_type, + CropAndResizeGradImage::Method(method))); + grad_outputs->push_back(CropAndResizeGradBoxes( + scope, grad_inputs[0], op.input(0), op.input(1), op.input(2))); + grad_outputs->push_back(NoGradient()); + grad_outputs->push_back(NoGradient()); + return scope.status(); +} + +REGISTER_GRADIENT_OP("CropAndResize", CropAndResizeGradHelper); } // anonymous namespace } // namespace ops } // namespace tensorflow diff --git a/tensorflow/cc/gradients/image_grad_test.cc b/tensorflow/cc/gradients/image_grad_test.cc index 1d150226538093467e092e02f38090a327f9c9b6..3bd52c80bd9a59c5bda977c206e6822f87c8bbbd 100644 --- a/tensorflow/cc/gradients/image_grad_test.cc +++ b/tensorflow/cc/gradients/image_grad_test.cc @@ -27,6 +27,7 @@ namespace tensorflow { namespace { using ops::Const; +using ops::CropAndResize; using ops::ResizeBicubic; using ops::ResizeBilinear; using ops::ResizeNearestNeighbor; @@ -51,7 +52,8 @@ class ImageGradTest : public ::testing::Test { template void MakeOp(const OpType op_type, const Tensor& x_data, const Input& y_shape, - const bool align_corners, Output* x, Output* y) { + const bool align_corners, const bool half_pixel_centers, + Output* x, Output* y) { *x = Const(scope_, x_data); switch (op_type) { case RESIZE_NEAREST: @@ -61,22 +63,26 @@ class ImageGradTest : public ::testing::Test { return; case RESIZE_BILINEAR: *y = ResizeBilinear(scope_, *x, y_shape, - ResizeBilinear::AlignCorners(align_corners)); + ResizeBilinear::AlignCorners(align_corners) + .HalfPixelCenters(half_pixel_centers)); return; case RESIZE_BICUBIC: *y = ResizeBicubic(scope_, *x, y_shape, - ResizeBicubic::AlignCorners(align_corners)); + ResizeBicubic::AlignCorners(align_corners) + .HalfPixelCenters(half_pixel_centers)); return; } assert(false); } template - void TestResizedShapeForType(const OpType op_type, const bool align_corners) { + void TestResizedShapeForType(const OpType op_type, const bool align_corners, + const bool half_pixel_centers) { TensorShape x_shape({1, 2, 2, 1}); Tensor x_data = MakeData(x_shape); Output x, y; - MakeOp(op_type, x_data, {4, 6}, align_corners, &x, &y); + MakeOp(op_type, x_data, {4, 6}, align_corners, half_pixel_centers, &x, + &y); ClientSession session(scope_); std::vector outputs; @@ -86,44 +92,64 @@ class ImageGradTest : public ::testing::Test { } void TestResizedShape(OpType op_type) { - for (const bool align_corners : {true, false}) { - TestResizedShapeForType(op_type, align_corners); - TestResizedShapeForType(op_type, align_corners); - TestResizedShapeForType(op_type, align_corners); + for (const bool half_pixel_centers : {true, false}) { + for (const bool align_corners : {true, false}) { + if (half_pixel_centers && align_corners) { + continue; + } + TestResizedShapeForType(op_type, align_corners, + half_pixel_centers); + TestResizedShapeForType(op_type, align_corners, + half_pixel_centers); + TestResizedShapeForType(op_type, align_corners, + half_pixel_centers); + } } } template void TestResizeToSmallerAndAlign(const OpType op_type, - const bool align_corners) { + const bool align_corners, + const bool half_pixel_centers) { TensorShape x_shape({1, 4, 6, 1}); Tensor x_data = MakeData(x_shape); Output x, y; - MakeOp(op_type, x_data, {2, 3}, align_corners, &x, &y); + MakeOp(op_type, x_data, {2, 3}, align_corners, half_pixel_centers, &x, + &y); JAC_T max_error; TF_ASSERT_OK((ComputeGradientError( scope_, x, x_data, y, {1, 2, 3, 1}, &max_error))); - EXPECT_LT(max_error, 1e-3); + EXPECT_LT(max_error, 1.5e-3); } template void TestResizeToLargerAndAlign(const OpType op_type, - const bool align_corners) { + const bool align_corners, + const bool half_pixel_centers) { TensorShape x_shape({1, 2, 3, 1}); Tensor x_data = MakeData(x_shape); Output x, y; - MakeOp(op_type, x_data, {4, 6}, align_corners, &x, &y); + MakeOp(op_type, x_data, {4, 6}, align_corners, half_pixel_centers, &x, + &y); JAC_T max_error; TF_ASSERT_OK((ComputeGradientError( scope_, x, x_data, y, {1, 4, 6, 1}, &max_error))); - EXPECT_LT(max_error, 1e-3); + EXPECT_LT(max_error, 1.5e-3); } template void TestResize(OpType op_type) { - for (const bool align_corners : {true, false}) { - TestResizeToSmallerAndAlign(op_type, align_corners); - TestResizeToLargerAndAlign(op_type, align_corners); + for (const bool half_pixel_centers : {true, false}) { + for (const bool align_corners : {true, false}) { + // if (!half_pixel_centers) continue; + if (half_pixel_centers && align_corners) { + continue; + } + TestResizeToSmallerAndAlign(op_type, align_corners, + half_pixel_centers); + TestResizeToLargerAndAlign(op_type, align_corners, + half_pixel_centers); + } } } @@ -194,5 +220,50 @@ class ScaleAndTranslateGradTest : public ::testing::Test { TEST_F(ScaleAndTranslateGradTest, Works) { TestResize(); } +class CropAndResizeGradTest : public ::testing::Test { + protected: + CropAndResizeGradTest() : scope_(Scope::NewRootScope()) {} + + template + Tensor MakeData(const TensorShape& data_shape) { + DataType data_type = DataTypeToEnum::v(); + Tensor data(data_type, data_shape); + auto data_flat = data.flat(); + for (int i = 0; i < data_flat.size(); ++i) { + data_flat(i) = T(i); + } + return data; + } + + template + void MakeOp(const Tensor& x_data, const Input& boxes, const Input& box_ind, + const Input& crop_szie, Output* x, Output* y) { + *x = Const(scope_, x_data); + *y = CropAndResize(scope_, *x, boxes, box_ind, crop_szie, + CropAndResize::Method("bilinear")); + TF_ASSERT_OK(scope_.status()); + } + + template + void TestCropAndResize() { + TensorShape x_shape({1, 4, 2, 1}); + Tensor x_data = MakeData(x_shape); + TensorShape box_shape({1, 4}); + Tensor boxes = MakeData(box_shape); + Output x, y; + MakeOp(x_data, boxes, {0}, {1, 1}, &x, &y); + JAC_T max_error; + TF_ASSERT_OK((ComputeGradientError( + scope_, x, x_data, y, {1, 1, 1, 1}, &max_error))); + EXPECT_LT(max_error, 1e-3); + } + + Scope scope_; +}; + +TEST_F(CropAndResizeGradTest, TestCrop) { + TestCropAndResize(); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compat_template.__init__.py b/tensorflow/compat_template.__init__.py index 2cf68c9cd8396987899b4f34f21b994b4722ead4..49cb74f19ef325c5861b124e458dd7e3b7f436e9 100644 --- a/tensorflow/compat_template.__init__.py +++ b/tensorflow/compat_template.__init__.py @@ -54,3 +54,4 @@ if hasattr(_current_module, 'keras'): losses = keras.losses metrics = keras.metrics optimizers = keras.optimizers + initializers = keras.initializers diff --git a/tensorflow/compiler/aot/codegen.cc b/tensorflow/compiler/aot/codegen.cc index da0598736a7d6b7f55458d76ca30fa6ad46a74f9..2355fad8802a490fafb702f53d88312611f9ebf4 100644 --- a/tensorflow/compiler/aot/codegen.cc +++ b/tensorflow/compiler/aot/codegen.cc @@ -256,7 +256,7 @@ Status GenVariableMethods(const tf2xla::Config& config, TF_RETURN_IF_ERROR( AddRewritesForShape(i, xla::Shape(ps.parameters(i)), &rewrites)); const string code = R"( - void set_var_{{NAME}}_input_data({{TYPE}}* data) { + void set_var_{{NAME}}_data({{TYPE}}* data) { set_arg_data({{I}}, data); } )"; @@ -270,17 +270,17 @@ Status GenVariableMethods(const tf2xla::Config& config, TF_RETURN_IF_ERROR(AddRewritesForShape( i, xla::Shape(ps.result().tuple_shapes(i)), &rewrites)); string code = R"( - {{TYPE}}* var_{{NAME}}_result_data() { + {{TYPE}}* var_{{NAME}}_data() { return static_cast<{{TYPE}}*>(result_data({{I}})); } - {{TYPE}}& var_{{NAME}}_result({{DIM_VARS}}) { + {{TYPE}}& var_{{NAME}}({{DIM_VARS}}) { return (*static_cast<{{TYPE}}(*){{DIM_SIZES}}>( result_data({{I}}))){{INDICES}}; } - const {{TYPE}}* var_{{NAME}}_result_data() const { + const {{TYPE}}* var_{{NAME}}_data() const { return static_cast(result_data({{I}})); } - const {{TYPE}}& var_{{NAME}}_result({{DIM_VARS}}) const { + const {{TYPE}}& var_{{NAME}}({{DIM_VARS}}) const { return (*static_cast( result_data({{I}}))){{INDICES}}; } diff --git a/tensorflow/compiler/aot/codegen_test_h.golden b/tensorflow/compiler/aot/codegen_test_h.golden index b5f33d690d492489e9090786cd341e035ae7ca15..8591df538779e3bc0f6e55607180a6d49009735e 100644 --- a/tensorflow/compiler/aot/codegen_test_h.golden +++ b/tensorflow/compiler/aot/codegen_test_h.golden @@ -228,40 +228,40 @@ class MyClass final : public tensorflow::XlaCompiledCpuFunction { // with dim indices specifying which value. No bounds checking is performed // on dim indices. - void set_var_myvar_input_data(float* data) { + void set_var_myvar_data(float* data) { set_arg_data(2, data); } - void set_var_myvar2_input_data(tensorflow::int32* data) { + void set_var_myvar2_data(tensorflow::int32* data) { set_arg_data(3, data); } - float* var_myvar_result_data() { + float* var_myvar_data() { return static_cast(result_data(1)); } - float& var_myvar_result() { + float& var_myvar() { return (*static_cast( result_data(1)))[0]; } - const float* var_myvar_result_data() const { + const float* var_myvar_data() const { return static_cast(result_data(1)); } - const float& var_myvar_result() const { + const float& var_myvar() const { return (*static_cast( result_data(1)))[0]; } - tensorflow::int32* var_myvar2_result_data() { + tensorflow::int32* var_myvar2_data() { return static_cast(result_data(2)); } - tensorflow::int32& var_myvar2_result(size_t dim0) { + tensorflow::int32& var_myvar2(size_t dim0) { return (*static_cast( result_data(2)))[dim0]; } - const tensorflow::int32* var_myvar2_result_data() const { + const tensorflow::int32* var_myvar2_data() const { return static_cast(result_data(2)); } - const tensorflow::int32& var_myvar2_result(size_t dim0) const { + const tensorflow::int32& var_myvar2(size_t dim0) const { return (*static_cast( result_data(2)))[dim0]; } diff --git a/tensorflow/compiler/aot/tests/BUILD b/tensorflow/compiler/aot/tests/BUILD index 444264ba6e1f59c33551796025ba845c62c02d43..ce8dae4262913c975ca69dedd0420f1457e11ee9 100644 --- a/tensorflow/compiler/aot/tests/BUILD +++ b/tensorflow/compiler/aot/tests/BUILD @@ -26,6 +26,8 @@ test_suite( ":test_graph_tfmatmulandadd_test", ":test_graph_tfsplits_test", ":test_graph_tftop_k_test", + ":test_graph_tfvariable_sequential_updates_test", + ":test_graph_tfvariable_test", ":tfcompile_test", ], ) @@ -70,6 +72,7 @@ genrule( "test_graph_tfsplits.pb", "test_graph_tftop_k.pb", "test_graph_tfvariable.pb", + "test_graph_tfvariable_sequential_updates.pb", ], # Set CUDA_VISIBLE_DEVICES='' to prevent the code we launch from using any # GPUs which might be present. This is important because builds may run @@ -234,6 +237,17 @@ tf_library( ], ) +tf_library( + name = "test_graph_tfvariable_sequential_updates", + testonly = 1, + config = "test_graph_tfvariable_sequential_updates.config.pbtxt", + cpp_class = "VariableSequentialUpdatesComp", + graph = "test_graph_tfvariable_sequential_updates.pb", + tags = [ + "manual", + ], +) + tf_cc_test( name = "tfcompile_test", srcs = ["tfcompile_test.cc"], @@ -254,6 +268,7 @@ tf_cc_test( ":test_graph_tfsplits", ":test_graph_tftop_k", ":test_graph_tfvariable", + ":test_graph_tfvariable_sequential_updates", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:xla_data_proto", diff --git a/tensorflow/compiler/aot/tests/make_test_graphs.py b/tensorflow/compiler/aot/tests/make_test_graphs.py index 42f8812def0503824416d92daa2db71a64c3db88..7f5e907e26365c0d9ec65e6f00d410a87f452241 100644 --- a/tensorflow/compiler/aot/tests/make_test_graphs.py +++ b/tensorflow/compiler/aot/tests/make_test_graphs.py @@ -157,6 +157,17 @@ def tfvariable(_): array_ops.stack([old_x, new_x], name='result') +def tfvariable_sequential_updates(_): + x = variables.Variable(1.0, name='x') + updates = control_flow_ops.no_op() + for _ in range(3): + with ops.control_dependencies([updates]): + x_val = x.read_value() + 1.0 + updates = x.assign_sub(0.1 * x_val) + + array_ops.identity(updates, name='result') + + def write_graph(build_graph, out_dir): """Build a graph using build_graph and write it out.""" g = ops.Graph() @@ -180,6 +191,7 @@ def main(_): write_graph(tfsplits, FLAGS.out_dir) write_graph(tftop_k, FLAGS.out_dir) write_graph(tfvariable, FLAGS.out_dir) + write_graph(tfvariable_sequential_updates, FLAGS.out_dir) if __name__ == '__main__': diff --git a/tensorflow/compiler/aot/tests/test_graph_tfvariable_sequential_updates.config.pbtxt b/tensorflow/compiler/aot/tests/test_graph_tfvariable_sequential_updates.config.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..7312c40baf6957c273fc389efa11d08ed9f7a0dd --- /dev/null +++ b/tensorflow/compiler/aot/tests/test_graph_tfvariable_sequential_updates.config.pbtxt @@ -0,0 +1,9 @@ +# Text form of tensorflow.tf2xla.Config proto. +fetch { + id { node_name: "result" } +} + +variable { + node_name: "x" + type: DT_FLOAT +} diff --git a/tensorflow/compiler/aot/tests/tfcompile_test.cc b/tensorflow/compiler/aot/tests/tfcompile_test.cc index 5f9316f3933713e12fc5960b9adfecc6e9bd99b5..5bee7f2540a4177a9c4e726bb739d7b92a4dacfc 100644 --- a/tensorflow/compiler/aot/tests/tfcompile_test.cc +++ b/tensorflow/compiler/aot/tests/tfcompile_test.cc @@ -31,6 +31,7 @@ limitations under the License. #include "tensorflow/compiler/aot/tests/test_graph_tfsplits.h" #include "tensorflow/compiler/aot/tests/test_graph_tftop_k.h" #include "tensorflow/compiler/aot/tests/test_graph_tfvariable.h" +#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_sequential_updates.h" #include "tensorflow/compiler/xla/service/hlo_profile_printer.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" @@ -480,20 +481,41 @@ TEST(TFCompileTest, Variable) { VariableComp fn; float x = 23; - fn.set_var_x_input_data(&x); + fn.set_var_x_data(&x); fn.set_thread_pool(&device); fn.Run(); EXPECT_EQ(fn.result0(0, 0), 23); EXPECT_EQ(fn.result0(1, 0), 65); - EXPECT_EQ(fn.var_x_result(), 65); + EXPECT_EQ(fn.var_x(), 65); - EXPECT_EQ(x, 23); - x = fn.var_x_result(); + EXPECT_EQ(fn.var_x_data(), &x); + EXPECT_EQ(x, 65); fn.Run(); EXPECT_EQ(fn.result0(0, 0), 65); EXPECT_EQ(fn.result0(1, 0), 107); - EXPECT_EQ(fn.var_x_result(), 107); + EXPECT_EQ(fn.var_x(), 107); +} + +TEST(TFCompileTest, VariableSequentialUpdates) { + Eigen::ThreadPool tp(1); + Eigen::ThreadPoolDevice device(&tp, tp.NumThreads()); + + // This implements the recursion: + // x[0] = 1.0 + // x[n+1] = x[n] - 0.1*(x[n-1] + 1.0) + VariableSequentialUpdatesComp fn; + float x = 1; + fn.set_var_x_data(&x); + + fn.set_thread_pool(&device); + // First calculate x[3] + fn.Run(); + EXPECT_NEAR(x, 0.458f, 1e-6); + + // Then calculate x[6] + fn.Run(); + EXPECT_NEAR(x, 0.062882f, 1e-6); } TEST(TFCompileTest, AssertEqAndReturnDiff) { diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index 121de401cefb2b56b984944dde769f226590dc67..4253cec5fc4cbace9cd70e06117a97cbe3a8f554 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -264,11 +264,11 @@ cc_library( "//tensorflow/compiler/xla/service:device_memory_allocator", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", + "//tensorflow/core:framework_internal", "//tensorflow/core:gpu_runtime", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", - "//tensorflow/core/kernels:variable_ops", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/memory", @@ -292,7 +292,6 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", - "//tensorflow/core/kernels:variable_ops", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", @@ -535,11 +534,11 @@ cc_library( "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", + "//tensorflow/core:framework_bounds_check", "//tensorflow/core:graph", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", - "//tensorflow/core/kernels:bounds_check", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", @@ -558,9 +557,9 @@ cc_library( ":resource_operation_safety_analysis", "//tensorflow/compiler/jit/graphcycles", "//tensorflow/core:framework", + "//tensorflow/core:framework_bounds_check", "//tensorflow/core:graph", "//tensorflow/core:protos_all_cc", - "//tensorflow/core/kernels:bounds_check", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", ], diff --git a/tensorflow/compiler/jit/create_xla_launch_op.cc b/tensorflow/compiler/jit/create_xla_launch_op.cc index 6f1ff85f24a4c1fd3e6d54fcff9f8868aee6f750..7e4c8466d88afa508e3d695c2529007e4b14cce9 100644 --- a/tensorflow/compiler/jit/create_xla_launch_op.cc +++ b/tensorflow/compiler/jit/create_xla_launch_op.cc @@ -126,8 +126,9 @@ Status GetBodyAndConstantsAndResources(FunctionLibraryRuntime* flr, const DataTypeVector& arg_types = (*fbody)->arg_types; std::vector const_args(arg_types.size()); // If we can't analyze the const args. Bail out. - TF_RETURN_IF_ERROR(BackwardsConstAnalysis( - *((*fbody)->graph), &const_args, /*compile_time_const_nodes=*/nullptr)); + TF_RETURN_IF_ERROR( + BackwardsConstAnalysis(*((*fbody)->graph), &const_args, + /*compile_time_const_nodes=*/nullptr, flr)); for (int i = 0; i < const_args.size(); ++i) { if (const_args[i]) { diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc index d0d7a3f3785469acd79a83b6897668f94fc6ea2e..8b6bffa267d638f7b4e2119da598c0eaa9b53fe5 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc @@ -2327,7 +2327,8 @@ Status Encapsulator::MakePrunedGraphCopyAndInline( return library->LookUpOpDef(op, sig); }, &fbody)); - InlineFunctionBody(*library, pruned_graph->get(), node, fbody); + TF_RETURN_IF_ERROR( + InlineFunctionBody(*library, pruned_graph->get(), node, fbody)); delete fbody; } @@ -2575,8 +2576,9 @@ Status EncapsulateSubgraphsPass::Run( const int num_args = input_permutation->size(); std::vector const_args(num_args); - TF_RETURN_IF_ERROR(BackwardsConstAnalysis( - **subgraph, &const_args, /*compile_time_const_nodes=*/nullptr)); + TF_RETURN_IF_ERROR( + BackwardsConstAnalysis(**subgraph, &const_args, + /*compile_time_const_nodes=*/nullptr, flr)); DataTypeVector arg_types(num_args); TF_RETURN_IF_ERROR(GetArgTypes(**subgraph, &arg_types)); diff --git a/tensorflow/compiler/jit/kernels/BUILD b/tensorflow/compiler/jit/kernels/BUILD index d0fa2c40be9d6b13ec736a9d6483dae0b4f0f45e..bd816e749c2044302242e832da39f1163a6b5c81 100644 --- a/tensorflow/compiler/jit/kernels/BUILD +++ b/tensorflow/compiler/jit/kernels/BUILD @@ -28,7 +28,6 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:state_ops_op_lib", "//tensorflow/core:stream_executor_no_cuda", - "//tensorflow/core/kernels:variable_ops", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", ], diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc index 997ef6e14bb9bd16ddac13eaf67368966818b29e..a5f5bc68c15adcbb80d93878ec093ca01f44cffd 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.cc +++ b/tensorflow/compiler/jit/kernels/xla_ops.cc @@ -35,7 +35,6 @@ limitations under the License. #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/variable_ops.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/env.h" diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index 20c2cd7e0561f92a01486102c4d2c572fd80c957..fd778490c320d4e344f3c8fcbddc3dc02df63673 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -441,7 +441,7 @@ Status FindCompilationCandidates( std::vector compile_time_const_nodes(graph.num_node_ids(), false); TF_RETURN_IF_ERROR( BackwardsConstAnalysis(graph, /*compile_time_const_arg_indices=*/nullptr, - &compile_time_const_nodes)); + &compile_time_const_nodes, lib_runtime)); int64& fuel = GetMarkForCompilationPassFlags()->tf_xla_clustering_fuel; @@ -1176,6 +1176,8 @@ Status MarkForCompilationPass::RunImpl( if (absl::optional cluster_name = GetXlaClusterForNode(*n)) { n->set_name(absl::StrCat(*cluster_name, "/", n->name())); + } else if (n->type_string() == "VarHandleOp") { + n->set_name(absl::StrCat("varhandle/", n->name())); } else { // There is room for improvement here. In particular, it may help to // split these unclustered nodes into classes where every node in a diff --git a/tensorflow/compiler/jit/partially_decluster_pass.cc b/tensorflow/compiler/jit/partially_decluster_pass.cc index e1fd2aaee2822daeffb415d053c9c4f56002a856..ffc5d0edbcc7668d5ee137c3c8bbe74167e37a1a 100644 --- a/tensorflow/compiler/jit/partially_decluster_pass.cc +++ b/tensorflow/compiler/jit/partially_decluster_pass.cc @@ -20,9 +20,13 @@ limitations under the License. #include "tensorflow/compiler/jit/xla_cluster_util.h" #include "tensorflow/compiler/tf2xla/const_analysis.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/core/common_runtime/function.h" +#include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/memory_types.h" #include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/public/version.h" namespace tensorflow { namespace { @@ -272,12 +276,20 @@ Status MustCompileNode(const Node* n, bool* must_compile) { // We assume here that the extra repeated (repeated compared to a clustered f // where it will always be constant folded) host-side computation of f does not // regress performance in any significant manner. We will have to revisit this -// algorith with a more complex cost model if this assumption turns out to be +// algorithm with a more complex cost model if this assumption turns out to be // incorrect. -Status PartiallyDeclusterGraph(Graph* graph) { +Status PartiallyDeclusterGraph(Graph* graph, + const FunctionLibraryDefinition* flib_def, + Env* env) { std::vector compile_time_const_nodes(graph->num_node_ids()); - TF_RETURN_IF_ERROR(BackwardsConstAnalysis( - *graph, nullptr, &compile_time_const_nodes, IsIntraClusterEdge)); + OptimizerOptions opts; + auto pflr = absl::make_unique( + nullptr, env, TF_GRAPH_DEF_VERSION, flib_def, opts); + FunctionLibraryRuntime* lib_runtime = + pflr->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice); + TF_RETURN_IF_ERROR(BackwardsConstAnalysis(*graph, nullptr, + &compile_time_const_nodes, + lib_runtime, IsIntraClusterEdge)); std::vector rpo; GetReversePostOrder(*graph, &rpo, /*stable_comparator=*/NodeComparatorName(), @@ -341,7 +353,19 @@ Status PartiallyDeclusterPass::Run( TF_RETURN_IF_ERROR( reduce_device_to_host_copies::PartiallyDeclusterGraph(graph)); - TF_RETURN_IF_ERROR(reduce_recompilation::PartiallyDeclusterGraph(graph)); + if (options.flib_def == nullptr) { + return errors::InvalidArgument( + "GraphOptimizationPassOptions::flib_def must be set for " + "PartiallyDeclusterPass."); + } + if (options.session_options == nullptr || + options.session_options->env == nullptr) { + return errors::InvalidArgument( + "GraphOptimizationPassOptions::session_options::env must be set for " + "PartiallyDeclusterPass."); + } + TF_RETURN_IF_ERROR(reduce_recompilation::PartiallyDeclusterGraph( + graph, options.flib_def, options.session_options->env)); return Status::OK(); } diff --git a/tensorflow/compiler/jit/partially_decluster_pass_test.cc b/tensorflow/compiler/jit/partially_decluster_pass_test.cc index 1d81a8f4fcbf050663626b1f7660afd71f4027bc..3494d0ee7efb51a5620f68bc1772e111db493c8d 100644 --- a/tensorflow/compiler/jit/partially_decluster_pass_test.cc +++ b/tensorflow/compiler/jit/partially_decluster_pass_test.cc @@ -27,6 +27,8 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/cc/ops/xla_ops.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/core/framework/function.h" +#include "tensorflow/core/framework/function.pb.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/op.h" #include "tensorflow/core/graph/algorithm.h" @@ -90,6 +92,12 @@ Status PartiallyDecluster(std::unique_ptr* graph) { GraphOptimizationPassOptions opt_options; opt_options.graph = graph; + FunctionDefLibrary fdef_lib; + FunctionLibraryDefinition flib_def(OpRegistry::Global(), fdef_lib); + opt_options.flib_def = &flib_def; + SessionOptions session_options; + session_options.env = Env::Default(); + opt_options.session_options = &session_options; PartiallyDeclusterPass pass; return pass.Run(opt_options); } diff --git a/tensorflow/compiler/jit/xla_compilation_cache.cc b/tensorflow/compiler/jit/xla_compilation_cache.cc index 611515cf33bc1abe21e06eb7f1513800276e095b..9f958463076765f2f9895e7860f79dcbbd4d342e 100644 --- a/tensorflow/compiler/jit/xla_compilation_cache.cc +++ b/tensorflow/compiler/jit/xla_compilation_cache.cc @@ -31,7 +31,6 @@ limitations under the License. #include "tensorflow/core/framework/types.h" #include "tensorflow/core/graph/graph_constructor.h" #include "tensorflow/core/graph/node_builder.h" -#include "tensorflow/core/kernels/variable_ops.h" #include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/logging.h" diff --git a/tensorflow/compiler/jit/xla_cpu_device.cc b/tensorflow/compiler/jit/xla_cpu_device.cc index 94dc61d55fb047c0ea81d98fde24cb55387c27d7..345e87a5735ae15ee637874b3bf9dcf4ff71537d 100644 --- a/tensorflow/compiler/jit/xla_cpu_device.cc +++ b/tensorflow/compiler/jit/xla_cpu_device.cc @@ -83,9 +83,10 @@ REGISTER_LOCAL_DEVICE_FACTORY(DEVICE_XLA_CPU, XlaCpuDeviceFactory); // Kernel registrations -constexpr std::array kAllXlaCpuTypes = { +constexpr std::array kAllXlaCpuTypes = { {DT_UINT8, DT_QUINT8, DT_INT8, DT_QINT8, DT_INT32, DT_QINT32, DT_INT64, - DT_HALF, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64, DT_COMPLEX128, DT_BOOL}}; + DT_HALF, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64, DT_COMPLEX128, DT_BOOL, + DT_BFLOAT16}}; REGISTER_XLA_LAUNCH_KERNEL(DEVICE_XLA_CPU, XlaLocalLaunchOp, kAllXlaCpuTypes); REGISTER_XLA_COMPILE_KERNEL(DEVICE_XLA_CPU, XlaCompileOp, kAllXlaCpuTypes); diff --git a/tensorflow/compiler/jit/xla_device.cc b/tensorflow/compiler/jit/xla_device.cc index eceb47f167f46784dc935a1d8b6fb4e5fe469367..56c4220f12b54be09821eca4590df52e8e71850b 100644 --- a/tensorflow/compiler/jit/xla_device.cc +++ b/tensorflow/compiler/jit/xla_device.cc @@ -480,6 +480,23 @@ bool XlaDevice::AllowsSyncOnCompletion() const { return sync_on_completion_; } +void XlaDevice::SetHandleDeviceErrorCallback(std::function callback) { + mutex_lock lock(mu_); + device_error_callback_ = callback; +} + +Status XlaDevice::HandleDeviceError() { + std::function local_device_error_callback; + { + mutex_lock lock(mu_); + local_device_error_callback = device_error_callback_; + } + if (local_device_error_callback != nullptr) { + return local_device_error_callback(); + } + return Status::OK(); +} + Status XlaDevice::RefreshStatus() { std::shared_ptr stream; { @@ -489,8 +506,14 @@ Status XlaDevice::RefreshStatus() { if (!stream) { return Status::OK(); } - // Stream status is XlaDevice status, no extra operations needed. - return stream->RefreshStatus(); + Status status = stream->RefreshStatus(); + if (!status.ok()) { + // Ignore errors from HandleDeviceError, since by definition the status is + // already non-ok, so there's nothing extra to report if HandleDeviceError + // itself returns an error. + HandleDeviceError().IgnoreError(); + } + return status; } XlaDeviceOpRegistrations* RegisterXlaDeviceKernels(const char* device, diff --git a/tensorflow/compiler/jit/xla_device.h b/tensorflow/compiler/jit/xla_device.h index 5fe1290fa03f2b1f9d90e36dbc5769b3c2728c8d..977f5f5cf151d979d025c2966012445af04fc502 100644 --- a/tensorflow/compiler/jit/xla_device.h +++ b/tensorflow/compiler/jit/xla_device.h @@ -171,6 +171,9 @@ class XlaDevice : public LocalDevice { void SetAllowsSyncOnCompletion(bool sync_on_completion) LOCKS_EXCLUDED(mu_); bool AllowsSyncOnCompletion() const override LOCKS_EXCLUDED(mu_); + // Installs an error handling callback when RefreshStatus sees !status.ok(). + void SetHandleDeviceErrorCallback(std::function callback); + Status RefreshStatus() override LOCKS_EXCLUDED(mu_); private: @@ -187,6 +190,9 @@ class XlaDevice : public LocalDevice { static Status GetMetadataFromDevice(DeviceBase* device, const XlaDevice::Metadata** metadata); + // Handles error when RefreshStatus sees !status.ok(). + Status HandleDeviceError(); + mutable mutex mu_; // The metadata of this XlaDevice. const Metadata xla_metadata_; @@ -237,6 +243,9 @@ class XlaDevice : public LocalDevice { // regardless of status. bool sync_on_completion_ GUARDED_BY(mu_) = true; + // A callback that will be invoked when RefreshStatus sees a status error. + std::function device_error_callback_ GUARDED_BY(mu_); + // Set of devices to use. This controls which of the devices on the given // platform will have resources allocated. For GPUs this will be // filled from visible_gpu_devices list from session configuration. diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h index 554227f09de0ab4d9e07f199b957657f3121ff06..c915b7118d09abe467ebf0b1d74a1efab94fd724 100644 --- a/tensorflow/compiler/jit/xla_launch_util.h +++ b/tensorflow/compiler/jit/xla_launch_util.h @@ -26,9 +26,9 @@ limitations under the License. #include "tensorflow/compiler/xla/service/device_memory_allocator.h" #include "tensorflow/compiler/xla/service/owning_device_memory.h" #include "tensorflow/core/framework/allocation_description.pb.h" +#include "tensorflow/core/framework/resource_var.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/variable_ops.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/gtl/array_slice.h" diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index 7c1e0daf0b7b418530367cb80fbd18b93e8e5f5e..c4fdbd67bf89f29a45e89551fb76368f554bb160 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -170,13 +170,6 @@ tf_xla_py_test( name = "argminmax_test", size = "small", srcs = ["argminmax_test.py"], - # ArgMax needs CustomCall on CPU, which is not available in normal - # (not precompiled) TensorFlow. The flag below excludes the CPU - # backend. - disabled_backends = [ - "cpu", - "cpu_ondemand", - ], deps = [ ":xla_test", "//tensorflow/python:array_ops", @@ -1083,6 +1076,7 @@ cuda_py_test( "//tensorflow/python:math_ops", "//tensorflow/python:nn_ops", ], + shard_count = 5, ) cuda_py_test( @@ -1216,10 +1210,6 @@ tf_xla_py_test( name = "quantized_ops_test", size = "medium", srcs = ["quantized_ops_test.py"], - disabled_backends = [ - "cpu", - "cpu_ondemand", - ], deps = [ ":xla_test", "//tensorflow/compiler/tf2xla/python:xla", diff --git a/tensorflow/compiler/tests/matrix_band_part_test.py b/tensorflow/compiler/tests/matrix_band_part_test.py index c61965b97fc142ce452cf28def8c937f692d2f84..0eec070a906670ff36c772edda22f8291b5b734a 100644 --- a/tensorflow/compiler/tests/matrix_band_part_test.py +++ b/tensorflow/compiler/tests/matrix_band_part_test.py @@ -167,6 +167,11 @@ class MatrixBandPartTest(xla_test.XLATestCase, parameterized.TestCase): }, ) def testMatrixBandPart(self, batch_shape, rows, cols): + # TODO(b/125505881): Disabled due to LLVM backend crash. + if self.device == 'XLA_CPU' and cols == 7 and rows == 1 and batch_shape == [ + 1, 3, 2 + ]: + pass for dtype in self.float_types: with self.cached_session(): mat = np.ones(batch_shape + [rows, cols]).astype(dtype) diff --git a/tensorflow/compiler/tests/while_test.py b/tensorflow/compiler/tests/while_test.py index 4ee144beb7f3243be069d59ee4a613484fe183b3..55d1f85370000c4050b0eaadfe5b3fcaae7143f9 100644 --- a/tensorflow/compiler/tests/while_test.py +++ b/tensorflow/compiler/tests/while_test.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import os import numpy as np from tensorflow.compiler.tests import xla_test @@ -25,7 +26,10 @@ from tensorflow.compiler.tf2xla.python import xla from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import function +from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import gradients_impl from tensorflow.python.platform import test @@ -125,6 +129,105 @@ class WhileTest(xla_test.XLATestCase): result = sess.run(loop_outputs, {init_index: 0}) self.assertAllClose(result, [10, 7], rtol=1e-3) - -if __name__ == '__main__': + def _testMaxItersSimple(self): + if is_compile_on_demand(): + self.skipTest("list_ops are not supported in cpu_ondemand") + with self.cached_session() as sess, self.test_scope(): + xla_context = control_flow_ops.XLAControlFlowContext() + xla_context.Enter() + v = constant_op.constant(1.0) + p = array_ops.placeholder(dtype=dtypes.int32) + + def create_while_loop(): + iterations = array_ops.size(p, name="iterations") + r = control_flow_ops.while_loop( + lambda *_: True, + lambda i, x: (i + 1, v * x), (0, 1.0), + maximum_iterations=iterations, + name="outer") + return array_ops.identity(r[1]) + + output = create_while_loop() + output = gradients_impl.gradients(output, v)[0] + + result = sess.run(output, feed_dict={p: [0, 0, 0]}) + print(result) + xla_context.Exit() + + def testMaxItersSimple(self): + self.skipTest("Fails with v1 control flow") + # This fails with old control. + # self._testMaxItersSimple() + + @test_util.enable_control_flow_v2 + def testMaxItersSimpleV2(self): + self._testMaxItersSimple() + + def _testNestedWhileLoopWithMaxItersFromOuterContext(self): + if is_compile_on_demand(): + self.skipTest("list_ops are not supported in cpu_ondemand") + with self.cached_session() as sess, self.test_scope(): + xla_context = control_flow_ops.XLAControlFlowContext() + xla_context.Enter() + v = constant_op.constant(1.0) + p = array_ops.placeholder(dtype=dtypes.int32) + + def mid_body_builder(iterations): + + def mid_body(i, x): + r = control_flow_ops.while_loop( + lambda *_: True, + lambda i, x: (i + 1, v * x), (0, x), + maximum_iterations=iterations, + name="inner") + return (i + 1, gradients_impl.gradients(x + r[1], v)[0]) + + return mid_body + + def outer_body(i, x): + iterations = array_ops.size(p, name="iterations") + return (i + 1, x + control_flow_ops.while_loop( + lambda *_: True, + mid_body_builder(iterations), (0, x), + maximum_iterations=iterations, + name="mid")[1]) + + def create_while_loop(): + r = control_flow_ops.while_loop( + lambda *_: True, + outer_body, (0, 1.0), + maximum_iterations=5, + name="outer") + return array_ops.identity(r[1]) + + # p:placeholder + # j = 0 + # i, x = 0, 1. + # while j++ < 5: + # i1, x1 = 0, x + # while i1++ < len(p): + # i2, x2 = 0, x1 + # while i2++ < len(p): + # x2 = v * x2 + # x1 = grad(x1 + x2, v) + # x = x1 + # output = x + output = create_while_loop() + sess.run(output, feed_dict={p: [0, 0, 0]}) + xla_context.Exit() + + def testNestedWhileLoopWithMaxItersFromOuterContext(self): + self._testNestedWhileLoopWithMaxItersFromOuterContext() + + @test_util.enable_control_flow_v2 + def testNestedWhileLoopWithMaxItersFromOuterContextV2(self): + self._testNestedWhileLoopWithMaxItersFromOuterContext() + + +def is_compile_on_demand(): + return ("TF_XLA_FLAGS" in os.environ and + "tf_xla_compile_on_demand" in os.environ["TF_XLA_FLAGS"]) + + +if __name__ == "__main__": test.main() diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc index 1f3cae3fda0cd7be296882b7b17ea47554edace8..9ab21d6a35dabda29a260f9015dfbc104e2136e8 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc @@ -86,95 +86,7 @@ TrtCandidateSelector::TrtCandidateSelector( TrtPrecisionMode precision_mode) : graph_properties_(graph_properties), precision_mode_(precision_mode) {} -Status TrtCandidateSelector::IsTensorRTCandidate(const tensorflow::Node* node) { - // TODO(laigd): move this set to TrtNodeValidator where it should belong. - // LINT.IfChange - static const auto* candidate_ops = new std::set{ - "Abs", - "Acos", - "Acosh", - "Add", - "Asin", - "Asinh", - "Atan", - "Atanh", - "AvgPool", - "BatchMatMul", - "BiasAdd", - "Ceil", - "ConcatV2", - "Const", - "Conv2D", - "Conv2DBackpropInput", - "Cos", - "Cosh", - "DepthwiseConv2dNative", - "Div", - "Exp", - "ExpandDims", - "Floor", - "FusedBatchNorm", - "FusedBatchNormV2", - "GatherV2", - "Identity", - "LeakyRelu", - "Log", - "MatMul", - "Max", - "Maximum", - "MaxPool", - "Mean", - "Min", - "Minimum", - "Mul", - "Neg", - "Pad", - "Prod", - "RealDiv", - "Reciprocal", - "Relu", - "Relu6", - "Reshape", - "Rsqrt", - "Sigmoid", - "Sin", - "Sinh", - "Slice", - "Snapshot", - "Softmax", - "Sqrt", - "Square", - "Squeeze", - "StridedSlice", - "Sub", - "Sum", - "Tan", - "Tanh", - "TopKV2", - "Transpose", - }; - bool is_supported_op_type = - (candidate_ops->count(node->type_string()) || - PluginFactoryTensorRT::GetInstance()->IsPlugin(node->type_string())); - static const auto* quantize_ops = new std::set{ - "QuantizeAndDequantizeV2", - "QuantizeAndDequantizeV3", - "FakeQuantWithMinMaxVars", - "FakeQuantWithMinMaxArgs", - }; - // In INT8 mode, we will always apply the quantization ranges provided by - // these ops to the relevant tensors. This happens regardless of the value of - // use_calibration. - if (precision_mode_ == TrtPrecisionMode::INT8 && - quantize_ops->count(node->type_string())) { - is_supported_op_type = true; - } - // LINT.ThenChange(//tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc) - if (!is_supported_op_type) { - return errors::Unimplemented("Op type ", node->type_string(), - " is not supported"); - } - +Status TrtCandidateSelector::IsTensorRTCandidate(const Node* node) { std::vector input_edges; TF_RETURN_IF_ERROR(node->input_edges(&input_edges)); std::vector> input_node_and_ports; @@ -184,34 +96,32 @@ Status TrtCandidateSelector::IsTensorRTCandidate(const tensorflow::Node* node) { input_edge->src_output()); } return validator_.ValidateNode(node->def(), input_node_and_ports, - graph_properties_); + precision_mode_, graph_properties_); } namespace { -tensorflow::Status BuildNodeMap( - const tensorflow::Graph& graph, - std::unordered_map* node_map) { +Status BuildNodeMap(const Graph& graph, + std::unordered_map* node_map) { for (auto* node : graph.op_nodes()) { if (!node_map->insert({node->name(), node}).second) { - return tensorflow::errors::AlreadyExists( - "Node name is not unique in graph: " + node->name()); + return errors::AlreadyExists("Node name is not unique in graph: " + + node->name()); } } - return tensorflow::Status::OK(); + return Status::OK(); } } // namespace -tensorflow::Status ConvertGraphDefToTensorRT( - const tensorflow::GraphDef& graph_def, - const std::vector& output_names, size_t max_batch_size, - size_t max_workspace_size_bytes, tensorflow::GraphDef* new_graph_def, - TrtPrecisionMode precision_mode, int minimum_segment_size, bool is_dyn_op, - int max_cached_engines, std::vector cached_engine_batches, - bool use_calibration) { +Status ConvertGraphDefToTensorRT( + const GraphDef& graph_def, const std::vector& output_names, + size_t max_batch_size, size_t max_workspace_size_bytes, + GraphDef* new_graph_def, TrtPrecisionMode precision_mode, + int minimum_segment_size, bool is_dyn_op, int max_cached_engines, + std::vector cached_engine_batches, bool use_calibration) { // Create GrapplerItem. - tensorflow::grappler::GrapplerItem item; + grappler::GrapplerItem item; item.fetch = output_names; item.graph = graph_def; @@ -225,13 +135,13 @@ tensorflow::Status ConvertGraphDefToTensorRT( // Create single machine cluster. Note that this will create a session and // initialize the gpu devices. const int num_cpu_cores = - tensorflow::grappler::GetNumAvailableLogicalCPUCores(); - const int num_gpus = tensorflow::grappler::GetNumAvailableGPUs(); + grappler::GetNumAvailableLogicalCPUCores(); + const int num_gpus = grappler::GetNumAvailableGPUs(); VLOG(2) << "cpu_cores: " << num_cpu_cores; VLOG(2) << "gpus: " << num_gpus; const int timeout_s = 60 * 10; - std::unique_ptr cluster( - new tensorflow::grappler::SingleMachine( + std::unique_ptr cluster( + new grappler::SingleMachine( timeout_s, num_cpu_cores, num_gpus)); // These settings are the defaults in tensorflow/python/grappler/cluster.py. cluster->DisableDetailedStats(true); @@ -242,18 +152,17 @@ tensorflow::Status ConvertGraphDefToTensorRT( // Create virtual cluster. Grappler requires a virtual cluster with a proper // GPU device in order to calculate flops>0 or fails with FATAL in dbg mode. // We add numbers from a Pascal card here to have flops>0. - tensorflow::DeviceProperties device_properties; + DeviceProperties device_properties; device_properties.set_type("GPU"); device_properties.mutable_environment()->insert({"architecture", "6"}); device_properties.set_num_cores(3584); device_properties.set_frequency(1531); - std::unique_ptr cluster( - new tensorflow::grappler::VirtualCluster( - {{"/GPU:0", device_properties}})); + std::unique_ptr cluster( + new grappler::VirtualCluster({{"/GPU:0", device_properties}})); #endif // Create RewriterConfig. - tensorflow::ConfigProto config_proto; + ConfigProto config_proto; auto& rw_cfg = *config_proto.mutable_graph_options()->mutable_rewrite_options(); // TODO(aaroey): use only const folding and layout for the time being since @@ -279,7 +188,7 @@ tensorflow::Status ConvertGraphDefToTensorRT( parameters["use_calibration"].set_b(use_calibration); // Run optimizer. - tensorflow::grappler::MetaOptimizer meta_opt(nullptr, config_proto); + grappler::MetaOptimizer meta_opt(nullptr, config_proto); TF_RETURN_IF_ERROR(meta_opt.Optimize(cluster.get(), item, new_graph_def)); if (VLOG_IS_ON(5)) { @@ -293,20 +202,18 @@ tensorflow::Status ConvertGraphDefToTensorRT( } struct EdgePtrCompare { - bool operator()(const tensorflow::Edge* lhs, - const tensorflow::Edge* rhs) const { + bool operator()(const Edge* lhs, const Edge* rhs) const { return lhs->id() < rhs->id(); } }; // Function to get subsegment information structure. -tensorflow::Status GetEngineInfo( - const tensorflow::Graph* g, - const tensorflow::grappler::GraphProperties& graph_properties, - const std::set& segment_nodes, - const std::unordered_map& node_map, - const std::vector& reverse_topo_order, - EngineInfo* info) { +Status GetEngineInfo(const Graph* g, + const grappler::GraphProperties& graph_properties, + const std::set& segment_nodes, + const std::unordered_map& node_map, + const std::vector& reverse_topo_order, + EngineInfo* info) { std::vector subgraph_nodes; // Topologically sorted nodes. std::set added_const_nodes; // Used to prevent double insertion. std::set segment_devices; @@ -353,8 +260,8 @@ tensorflow::Status GetEngineInfo( // Create input connections. Sort edges first to make determnistic since // in_edges is a set of pointers. - std::vector in_edges(node->in_edges().begin(), - node->in_edges().end()); + std::vector in_edges(node->in_edges().begin(), + node->in_edges().end()); std::sort(in_edges.begin(), in_edges.end(), EdgePtrCompare()); for (const auto edge : in_edges) { auto input_node = edge->src(); @@ -405,8 +312,8 @@ tensorflow::Status GetEngineInfo( } // Create output connections. Sort edges first to make determnistic since // out_edges is a set of pointers. - std::vector out_edges(node->out_edges().begin(), - node->out_edges().end()); + std::vector out_edges(node->out_edges().begin(), + node->out_edges().end()); std::sort(out_edges.begin(), out_edges.end(), EdgePtrCompare()); for (const auto edge : out_edges) { auto output_node = edge->dst(); @@ -465,7 +372,7 @@ void UpdateToEngineNode(const std::vector& infos, const size_t my_engine_id, const std::vector& engine_nodes, const bool is_input_edge, const string& node_name, - tensorflow::Node** node, int* port) { + Node** node, int* port) { for (size_t t = 0; t < infos.size(); ++t) { if (t == my_engine_id) { continue; @@ -502,20 +409,20 @@ void UpdateToEngineNode(const std::vector& infos, // one). Connect to the pre-existing engine node instead. // 3. In this way, we ensure the graph is topologically sort-able after each // invocation of CreateTRTNode(). -tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, - int max_batch_size, tensorflow::Graph* graph, - nvinfer1::IGpuAllocator* alloc, - std::vector* engine_nodes) { +Status CreateTRTNode(const std::vector& infos, int pos, + int max_batch_size, Graph* graph, + nvinfer1::IGpuAllocator* alloc, + std::vector* engine_nodes) { const auto& info = infos.at(pos); TRT_RETURN_IF_TEST_VALUE(StrCat(info.engine_name, ":CreateTRTNode"), "fail"); - std::vector output_shape_protos; - std::vector input_shape_protos; - std::vector input_shapes; - std::vector inputs; - std::vector input_nodes; - std::vector control_input_nodes; + std::vector output_shape_protos; + std::vector input_shape_protos; + std::vector input_shapes; + std::vector inputs; + std::vector input_nodes; + std::vector control_input_nodes; std::unordered_set control_input_names; - std::vector out_types; + std::vector out_types; VLOG(1) << "Processing " << info.engine_name; // Collect needed info for creating the engine node in the graph @@ -527,8 +434,8 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, if (!conn.is_input_edge) continue; // Rewrire control input if it's not found in original graph. - tensorflow::Node* input_node = graph->FindNodeId(conn.outside_id); - int port = tensorflow::Graph::kControlSlot; + Node* input_node = graph->FindNodeId(conn.outside_id); + int port = Graph::kControlSlot; if (!input_node) { UpdateToEngineNode(infos, pos, *engine_nodes, /*is_input_edge=*/true, conn.outside_node_name, &input_node, &port); @@ -544,7 +451,7 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, // Data edges if (!conn.is_input_edge) { // Set the shapes and data types of output edge. - tensorflow::TensorShapeProto out_shape; + TensorShapeProto out_shape; // shape of the output node inside segment conn.inside_shape.AsProto(&out_shape); if (output_shape_protos.size() <= conn.port_number) { @@ -555,7 +462,7 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, out_types.at(conn.port_number) = conn.connection_type; } else { // Set the shapes and data types of input edge. - tensorflow::TensorShapeProto in_shape; + TensorShapeProto in_shape; conn.outside_shape.AsProto(&in_shape); if (input_shape_protos.size() <= conn.port_number) { input_shape_protos.resize(conn.port_number + 1); @@ -568,7 +475,7 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, if (info.engine_type == EngineInfo::EngineType::TRTStatic) { for (int i = 1; i < conn.outside_shape.dims(); i++) { if (conn.outside_shape.dim_size(i) <= 0) { - return tensorflow::errors::Internal( + return errors::Internal( "Input shapes must be fully defined when in static mode. " "Please try is_dynamic_op=True (shape was ", conn.outside_shape.DebugString(), ")"); @@ -577,7 +484,7 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, } // Rewrire data input if it's not found in original graph. - tensorflow::Node* input_node = graph->FindNodeId(conn.outside_id); + Node* input_node = graph->FindNodeId(conn.outside_id); int port = conn.outside_port; if (!input_node) { UpdateToEngineNode(infos, pos, *engine_nodes, /*is_input_edge=*/true, @@ -600,7 +507,7 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, // avoid crash later. Constant folding should've folded the ops that make up // these segments. if (inputs.empty()) { - return tensorflow::errors::Internal( + return errors::Internal( "Segment has no inputs (possible " "constfold failure)"); } @@ -638,7 +545,7 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, string prec_string; TF_RETURN_IF_ERROR(TrtPrecisionModeToName(info.precision_mode, &prec_string)); - tensorflow::NodeDefBuilder node_builder(info.engine_name, "TRTEngineOp"); + NodeDefBuilder node_builder(info.engine_name, "TRTEngineOp"); if (!info.device.empty()) node_builder.Device(info.device); if (VLOG_IS_ON(1)) { string ins = StrCat(info.engine_name, " inputs= "); @@ -656,8 +563,8 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, !info.cached_engine_batches.empty()) { LOG(WARNING) << "Cached engine batches are ignored for static engines"; } - tensorflow::NodeDef trt_node; - tensorflow::Status status = + NodeDef trt_node; + Status status = node_builder.Attr("input_shapes", input_shape_protos) .Attr("output_shapes", output_shape_protos) .Attr("static_engine", @@ -682,7 +589,7 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, // here, this segment will be skipped // TODO(aaroey): let it return proper error status for the following logic // instead of checking fail. - tensorflow::Node* engine_node = graph->AddNode(trt_node, &status); + Node* engine_node = graph->AddNode(trt_node, &status); (*engine_nodes)[pos] = engine_node; if (!status.ok()) { LOG(ERROR) << "Adding node failed " << status; @@ -709,7 +616,7 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, if (conn.is_input_edge) { continue; } - tensorflow::Node* output_node = graph->FindNodeId(conn.outside_id); + Node* output_node = graph->FindNodeId(conn.outside_id); int port = conn.outside_port; if (!output_node) { UpdateToEngineNode(infos, pos, *engine_nodes, /*is_input_edge=*/false, @@ -732,20 +639,19 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, } // Function to construct a funcdef from the segment and add it to the graph. -tensorflow::Status RegisterSegmentFunctionToFunctionLibrary( - tensorflow::Graph* graph, const tensorflow::GraphDef& segment, - const string& engine_name) { - tensorflow::Graph sgraph(graph->flib_def()); - tensorflow::GraphConstructorOptions gcopts; - TF_RETURN_IF_ERROR( - tensorflow::ConvertGraphDefToGraph(gcopts, segment, &sgraph)); - std::map io_nodes; +Status RegisterSegmentFunctionToFunctionLibrary(Graph* graph, + const GraphDef& segment, + const string& engine_name) { + Graph sgraph(graph->flib_def()); + GraphConstructorOptions gcopts; + TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(gcopts, segment, &sgraph)); + std::map io_nodes; int num_inputs = 0; for (auto n : sgraph.op_nodes()) { - if (tensorflow::str_util::StartsWith(n->name(), kInputPHName)) { + if (str_util::StartsWith(n->name(), kInputPHName)) { num_inputs++; io_nodes.insert({n->name(), n}); - } else if (tensorflow::str_util::StartsWith(n->name(), kOutputPHName)) { + } else if (str_util::StartsWith(n->name(), kOutputPHName)) { io_nodes.insert({n->name(), n}); } } @@ -753,14 +659,14 @@ tensorflow::Status RegisterSegmentFunctionToFunctionLibrary( for (int i = 0; i < num_inputs; ++i) { auto name = StrCat(kInputPHName, i); auto node = io_nodes[name]; - tensorflow::NodeDef nd; - tensorflow::NodeDefBuilder node_builder( - StrCat(name, "_Arg"), tensorflow::FunctionLibraryDefinition::kArgOp); + NodeDef nd; + NodeDefBuilder node_builder(StrCat(name, "_Arg"), + FunctionLibraryDefinition::kArgOp); VLOG(1) << "Adding " << StrCat(name, "_Arg"); TF_RETURN_IF_ERROR(node_builder.Attr("T", node->output_type(0)) .Attr("index", i) .Finalize(&nd)); - tensorflow::Status s; + Status s; auto node_arg = sgraph.AddNode(nd, &s); if (!s.ok()) { LOG(ERROR) << "Couldn't add _Arg node for " << name; @@ -780,15 +686,14 @@ tensorflow::Status RegisterSegmentFunctionToFunctionLibrary( for (int i = 0; i < io_nodes.size() - num_inputs; ++i) { auto name = StrCat(kOutputPHName, i); auto node = io_nodes[name]; - tensorflow::NodeDef nd; - tensorflow::NodeDefBuilder node_builder( - StrCat(name, "_Ret"), tensorflow::FunctionLibraryDefinition::kRetOp); + NodeDef nd; + NodeDefBuilder node_builder(StrCat(name, "_Ret"), + FunctionLibraryDefinition::kRetOp); auto edge = *(node->in_edges().begin()); - tensorflow::NodeDefBuilder::NodeOut nout( - edge->src()->name(), edge->src_output(), - edge->src()->output_type(edge->src_output())); + NodeDefBuilder::NodeOut nout(edge->src()->name(), edge->src_output(), + edge->src()->output_type(edge->src_output())); VLOG(1) << " input " << nout.node << ":" << nout.index - << " dtype=" << tensorflow::DataTypeString(nout.data_type); + << " dtype=" << DataTypeString(nout.data_type); // nvcc complains that Input() is // ambiguous, so do not use Input({nout}). node_builder.Input(nout); @@ -798,7 +703,7 @@ tensorflow::Status RegisterSegmentFunctionToFunctionLibrary( if (VLOG_IS_ON(3)) { VLOG(3) << nd.DebugString(); } - tensorflow::Status s; + Status s; auto node_ret = sgraph.AddNode(nd, &s); if (!s.ok()) { LOG(ERROR) << "Couldn't add _Ret node for " << name; @@ -814,9 +719,9 @@ tensorflow::Status RegisterSegmentFunctionToFunctionLibrary( } sgraph.RemoveNode(node); } - tensorflow::FunctionDefLibrary fdeflib; + FunctionDefLibrary fdeflib; auto native_segment = fdeflib.add_function(); - TF_RETURN_IF_ERROR(tensorflow::GraphToFunctionDef( + TF_RETURN_IF_ERROR(GraphToFunctionDef( sgraph, StrCat(engine_name, "_native_segment"), native_segment)); // Set kIntsonDeviceAttr to true so that all TRTEngineOp outputs are always on // a GPU device as expected. Otherwise, some of the tensors of type DT_INT32 @@ -830,13 +735,13 @@ tensorflow::Status RegisterSegmentFunctionToFunctionLibrary( } VLOG(1) << "Adding funcdef to graphlib"; TF_RETURN_IF_ERROR(graph->AddFunctionLibrary(fdeflib)); - return tensorflow::Status::OK(); + return Status::OK(); } -std::pair GetDeviceAndAllocator( - const ConversionParams& params, const EngineInfo& engine) { +std::pair GetDeviceAndAllocator(const ConversionParams& params, + const EngineInfo& engine) { int cuda_device_id = -1; - tensorflow::Allocator* dev_allocator = nullptr; + Allocator* dev_allocator = nullptr; if (params.cluster == nullptr || params.cluster->GetDeviceSet() == nullptr || engine.device.empty()) { // If device is not set, use the first found GPU device for the conversion. @@ -864,7 +769,7 @@ std::pair GetDeviceAndAllocator( // Use the device requested by the engine. auto device_set = params.cluster->GetDeviceSet(); - std::vector devices; + std::vector devices; DeviceNameUtils::ParsedName parsed_name; if (DeviceNameUtils::ParseFullName(engine.device, &parsed_name) && parsed_name.has_id) { @@ -878,7 +783,7 @@ std::pair GetDeviceAndAllocator( StrAppend(&msg, ". Will get the allocator from first one."); LOG(WARNING) << msg; } - tensorflow::AllocatorAttributes alloc_attr; + AllocatorAttributes alloc_attr; cuda_device_id = devices[0]->tensorflow_gpu_device_info()->gpu_id; dev_allocator = devices[0]->GetAllocator(alloc_attr); VLOG(1) << "Using allocator " << dev_allocator->Name() @@ -892,25 +797,25 @@ std::pair GetDeviceAndAllocator( // Entry function from optimization pass. // TODO(aaeory): parameter should use pointer type. -tensorflow::Status ConvertAfterShapes(ConversionParams& params) { +Status ConvertAfterShapes(ConversionParams& params) { // Convert graphdef to graph. - tensorflow::FunctionLibraryDefinition flib(tensorflow::OpRegistry::Global(), - params.input_graph_def->library()); - tensorflow::Graph graph(flib); - TF_RETURN_IF_ERROR(tensorflow::ConvertGraphDefToGraph( - tensorflow::GraphConstructorOptions(), *params.input_graph_def, &graph)); + FunctionLibraryDefinition flib(OpRegistry::Global(), + params.input_graph_def->library()); + Graph graph(flib); + TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(GraphConstructorOptions(), + *params.input_graph_def, &graph)); // Segment the graph into subgraphs that can be converted to TensorRT - tensorflow::tensorrt::segment::SegmentOptions segment_options; + segment::SegmentOptions segment_options; // TODO(ben,jie,sami): exclude output nodes (DISCUSS IT) for (auto node : *(params.output_names)) { segment_options.exclude_node_list.insert(node); } segment_options.minimum_segment_size = params.minimum_segment_size; - tensorflow::tensorrt::segment::SegmentNodesVector initial_segments; + segment::SegmentNodesVector initial_segments; TrtCandidateSelector candidate_selector(*params.graph_properties, params.precision_mode); - TF_RETURN_IF_ERROR(tensorrt::segment::SegmentGraph( + TF_RETURN_IF_ERROR(segment::SegmentGraph( &graph, std::bind(&TrtCandidateSelector::IsTensorRTCandidate, &candidate_selector, std::placeholders::_1), @@ -922,16 +827,16 @@ tensorflow::Status ConvertAfterShapes(ConversionParams& params) { << initial_segments.size(); // Get the EngineInfo for each segment. - std::unordered_map node_map; + std::unordered_map node_map; TF_RETURN_IF_ERROR(BuildNodeMap(graph, &node_map)); float total_num_nodes_in_segments = 0.; std::vector engine_segments; engine_segments.reserve(initial_segments.size()); - std::vector reverse_topo_order; - tensorflow::GetPostOrder(graph, &reverse_topo_order); + std::vector reverse_topo_order; + GetPostOrder(graph, &reverse_topo_order); size_t total_engine_bytes_size = 0; std::vector engine_bytes_size; - tensorflow::tensorrt::segment::SegmentNodesVector converted_segments; + segment::SegmentNodesVector converted_segments; converted_segments.reserve(initial_segments.size()); for (size_t t = 0; t < initial_segments.size(); t++) { auto& curr_segment = initial_segments.at(t); @@ -1044,7 +949,7 @@ tensorflow::Status ConvertAfterShapes(ConversionParams& params) { cudaSetDevice(old_cuda_device); graph.ToGraphDef(params.output_graph_def); VLOG(1) << "Returning from conversion"; - return tensorflow::Status::OK(); + return Status::OK(); } } // namespace convert diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h index 80f68d36a3ab894e97586687ee9ab93dddc73c50..f4ae3f8811beb1277ce14fa07f9371ce7bccbe33 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h @@ -40,7 +40,7 @@ class TrtCandidateSelector { // Returns OK iff 'node' is a TF-TRT conversion candidate, which will be added // to TRT subgraph and later converted into TRT engine. - Status IsTensorRTCandidate(const tensorflow::Node* node); + Status IsTensorRTCandidate(const Node* node); private: // The TF-TRT node converter used to verify whether individual node is @@ -69,15 +69,15 @@ struct ConversionParams { fixed_input_size(true), use_calibration(true), max_cached_engines(1) {} - const tensorflow::GraphDef* input_graph_def; + const GraphDef* input_graph_def; const std::vector* output_names; size_t max_batch_size; size_t max_workspace_size_bytes; - tensorflow::GraphDef* output_graph_def; + GraphDef* output_graph_def; TrtPrecisionMode precision_mode; int minimum_segment_size; - const tensorflow::grappler::GraphProperties* graph_properties; - const tensorflow::grappler::Cluster* cluster; + const grappler::GraphProperties* graph_properties; + const grappler::Cluster* cluster; bool is_dyn_op; // Whether to create engine on conversion or execution time bool fixed_input_size; // Assume non-batch ranks of input tensors are fixed int max_cached_engines; // maximum number of cached engines @@ -89,17 +89,17 @@ struct ConversionParams { // optimization targets inference run with max batch size. // - max_workspace_size_bytes: The upper bound of memory allowance for engine // building. -tensorflow::Status ConvertGraphDefToTensorRT( - const tensorflow::GraphDef& graph_def, - const std::vector& output_names, size_t max_batch_size, - size_t max_workspace_size_bytes, tensorflow::GraphDef* new_graph_def, +Status ConvertGraphDefToTensorRT( + const GraphDef& graph_def, const std::vector& output_names, + size_t max_batch_size, size_t max_workspace_size_bytes, + GraphDef* new_graph_def, TrtPrecisionMode precision_mode = TrtPrecisionMode::FP32, int minimum_segment_size = 3, bool is_dyn_op = false, int max_cached_engines = 1, std::vector cached_engine_batches = {}, bool use_calibration = true); // Method to call from optimization pass -tensorflow::Status ConvertAfterShapes(ConversionParams& params); +Status ConvertAfterShapes(ConversionParams& params); // Return compile time TensorRT library version information. std::vector GetLinkedTensorRTVersion(); @@ -108,8 +108,8 @@ std::vector GetLinkedTensorRTVersion(); std::vector GetLoadedTensorRTVersion(); // Helper method for the conversion, expose for testing. -std::pair GetDeviceAndAllocator( - const ConversionParams& params, const EngineInfo& engine); +std::pair GetDeviceAndAllocator(const ConversionParams& params, + const EngineInfo& engine); } // namespace convert } // namespace tensorrt diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc index 059aadad11aea56bf35ce11a95c141bd7eeecfdc..81c7dd6149f696016805ab3fd3df3c3254ff5555 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc @@ -57,10 +57,10 @@ limitations under the License. // would work! #define TFTRT_CHECK_EQ_TYPE(val1, val2) CHECK_EQ((int)val1, (int)val2) -#define TFTRT_INTERNAL_ERROR_AT_NODE(node) \ - do { \ - return tensorflow::errors::Internal( \ - "TFTRT::", __FUNCTION__, " failed to add TRT layer, at: ", node); \ +#define TFTRT_INTERNAL_ERROR_AT_NODE(node) \ + do { \ + return errors::Internal("TFTRT::", __FUNCTION__, \ + " failed to add TRT layer, at: ", node); \ } while (0) #define TFTRT_RETURN_ERROR_IF_FALSE(status, node) \ @@ -94,27 +94,26 @@ namespace convert { using absl::StrAppend; using absl::StrCat; -inline tensorflow::Status ConvertDType(tensorflow::DataType tf_dtype, - nvinfer1::DataType* trt_dtype) { +inline Status ConvertDType(DataType tf_dtype, nvinfer1::DataType* trt_dtype) { switch (tf_dtype) { - case tensorflow::DataType::DT_FLOAT: + case DataType::DT_FLOAT: *trt_dtype = nvinfer1::DataType::kFLOAT; break; // TODO(aaroey): this should be DT_QINT8 which is not a well supported type. - case tensorflow::DataType::DT_INT8: + case DataType::DT_INT8: *trt_dtype = nvinfer1::DataType::kINT8; break; - case tensorflow::DataType::DT_HALF: + case DataType::DT_HALF: *trt_dtype = nvinfer1::DataType::kHALF; break; - case tensorflow::DataType::DT_INT32: + case DataType::DT_INT32: *trt_dtype = nvinfer1::DataType::kINT32; break; default: - return tensorflow::errors::InvalidArgument( - "Unsupported data type ", tensorflow::DataTypeString(tf_dtype)); + return errors::InvalidArgument("Unsupported data type ", + DataTypeString(tf_dtype)); } - return tensorflow::Status::OK(); + return Status::OK(); } template @@ -135,13 +134,23 @@ Status TensorShapeArrayToTrtDims(const std::vector& shape, PartialTensorShape tensor_shape; TF_RETURN_IF_ERROR(TensorShapeUtils::MakeShape(shape, &tensor_shape)); *out = TensorShapeToTrtDims(tensor_shape, ignore_first_dim); - return tensorflow::Status::OK(); + return Status::OK(); +} + +// TODO(laigd): use this utility function in more places. +Status RemoveBatchDimension(nvinfer1::Dims* dims) { + if (dims->nbDims < 2) { + return errors::InvalidArgument( + "Dropping batch dimension requires dims with rank>=2."); + } + std::copy(dims->d + 1, dims->d + dims->nbDims, dims->d); + dims->nbDims--; + return Status::OK(); } void GetOutputProperties(const grappler::GraphProperties& graph_properties, const Node* node, const int out_port, - PartialTensorShape* shape, - tensorflow::DataType* dtype) { + PartialTensorShape* shape, DataType* dtype) { if (graph_properties.HasOutputProperties(node->name())) { auto output_params = graph_properties.GetOutputProperties(node->name()); auto out_shape = output_params.at(out_port); @@ -155,8 +164,7 @@ void GetOutputProperties(const grappler::GraphProperties& graph_properties, void GetInputProperties(const grappler::GraphProperties& graph_properties, const Node* node, const int in_port, - PartialTensorShape* shape, - tensorflow::DataType* dtype) { + PartialTensorShape* shape, DataType* dtype) { if (graph_properties.HasInputProperties(node->name())) { auto input_params = graph_properties.GetInputProperties(node->name()); auto in_shape = input_params.at(in_port); @@ -168,7 +176,7 @@ void GetInputProperties(const grappler::GraphProperties& graph_properties, } Status ValidateTensorProperties(const string& producer_node_type, - const tensorflow::DataType dtype, + const DataType dtype, const PartialTensorShape& shape, bool validation_only, nvinfer1::DataType* trt_dtype, @@ -367,16 +375,16 @@ nvinfer1::ITensor* Converter::CreateConstantLayer( return trt_tensor; } -tensorflow::Status CreateBroadcastableScalarConstant( - OpConverterParams* params, float value, const nvinfer1::Dims& dims, - const nvinfer1::ITensor** tensor) { +Status CreateBroadcastableScalarConstant(OpConverterParams* params, float value, + const nvinfer1::Dims& dims, + const nvinfer1::ITensor** tensor) { // In order to be broadcastable, the number of dims has to match. nvinfer1::Dims broadcastable_dims(dims); for (int i = 0; i < broadcastable_dims.nbDims; i++) { broadcastable_dims.d[i] = 1; } TRT_ShapedWeights weights = params->weight_store->GetTempWeights( - tensorflow::DataType::DT_FLOAT, broadcastable_dims); + DataType::DT_FLOAT, broadcastable_dims); auto weights_ptr = static_cast(const_cast(weights.GetValues())); weights_ptr[0] = value; @@ -391,12 +399,12 @@ tensorflow::Status CreateBroadcastableScalarConstant( // includes the batch dimension, while TRT does not. TF can also use negative // indices. // TODO(tmorris): Use this method in more ops. -tensorflow::Status ConvertAxis(int tf_axis, int trt_nb_dims, - absl::string_view node_name, int* trt_axis) { +Status ConvertAxis(int tf_axis, int trt_nb_dims, absl::string_view node_name, + int* trt_axis) { const int tf_nb_dims = trt_nb_dims + 1; // Check bounds. if (tf_axis < -tf_nb_dims || tf_axis >= tf_nb_dims) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Axis value of ", tf_axis, " is out of bounds, must be in range [", -tf_nb_dims, ", ", tf_nb_dims, "), at ", node_name); } @@ -404,7 +412,7 @@ tensorflow::Status ConvertAxis(int tf_axis, int trt_nb_dims, if (tf_axis < 0) tf_axis += tf_nb_dims; // Don't allow axis to be the batch dimension. if (tf_axis == 0) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "TensorRT does not allow manipulation of the batch dimension, at ", node_name); } @@ -435,7 +443,7 @@ bool AllLengthsEqual(const std::vector>& inputs) { return true; } -inline nvinfer1::Dims GetTrtDimsForTensor(const tensorflow::Tensor& tensor) { +inline nvinfer1::Dims GetTrtDimsForTensor(const Tensor& tensor) { nvinfer1::Dims dims; dims.nbDims = tensor.dims(); for (int i = 0; i < dims.nbDims; i++) { @@ -518,7 +526,7 @@ nvinfer1::Weights TRT_ShapedWeights::GetTrtWeights() const { } size_t TRT_ShapedWeights::size_bytes() const { - return this->count() * tensorflow::DataTypeSize(this->type_); + return this->count() * DataTypeSize(this->type_); } string TRT_ShapedWeights::DebugString() const { @@ -654,7 +662,7 @@ string TRT_TensorOrWeights::DebugString() const { class TFAttrs { public: - explicit TFAttrs(const tensorflow::NodeDef& tf_node) { + explicit TFAttrs(const NodeDef& tf_node) { for (const auto& attr : tf_node.attr()) { attrs_.insert({attr.first, &attr.second}); } @@ -662,7 +670,7 @@ class TFAttrs { bool count(const string& key) const { return attrs_.count(key); } - tensorflow::AttrValue const* at(const string& key) const { + AttrValue const* at(const string& key) const { if (!attrs_.count(key)) { LOG(FATAL) << "Attribute not found: " << key; } @@ -686,7 +694,7 @@ class TFAttrs { } private: - typedef std::map AttrMap; + typedef std::map AttrMap; AttrMap attrs_; }; @@ -696,9 +704,9 @@ string TFAttrs::get(const string& key) const { } template <> -std::vector TFAttrs::get>(const string& key) const { +std::vector TFAttrs::get>(const string& key) const { auto attr = this->at(key)->list().i(); - return std::vector(attr.begin(), attr.end()); + return std::vector(attr.begin(), attr.end()); } template <> @@ -715,8 +723,7 @@ nvinfer1::DataType TFAttrs::get(const string& key) const { } template <> -tensorflow::DataType TFAttrs::get( - const string& key) const { +DataType TFAttrs::get(const string& key) const { return this->at(key)->type(); } @@ -731,7 +738,7 @@ bool TFAttrs::get(const string& key) const { } template <> -int TFAttrs::get(const string& key) const { +int64 TFAttrs::get(const string& key) const { return this->at(key)->i(); } @@ -776,7 +783,7 @@ void ReorderCKtoKC(const TRT_ShapedWeights& iweights, const nvinfer1::DimsHW istrides = {1, k}; const nvinfer1::DimsHW ostrides = {c, 1}; switch (iweights.type_) { - case tensorflow::DataType::DT_FLOAT: { + case DataType::DT_FLOAT: { Reorder2({k, c}, static_cast(iweights.GetValues()), istrides, // TODO(aaroey): get rid of all the const_cast like this. @@ -784,7 +791,7 @@ void ReorderCKtoKC(const TRT_ShapedWeights& iweights, ostrides); break; } - case tensorflow::DataType::DT_HALF: { + case DataType::DT_HALF: { Reorder2( {k, c}, static_cast(iweights.GetValues()), istrides, @@ -820,14 +827,14 @@ void ReorderRSCKToKCRS(const TRT_ShapedWeights& iweights, const nvinfer1::DimsNCHW istrides = {1, k, s * k * c, c * k}; const nvinfer1::DimsNCHW ostrides = {c * r * s, r * s, s, 1}; switch (iweights.type_) { - case tensorflow::DataType::DT_FLOAT: { + case DataType::DT_FLOAT: { Reorder4({k, c, r, s}, static_cast(iweights.GetValues()), istrides, static_cast(const_cast(oweights->GetValues())), ostrides); break; } - case tensorflow::DataType::DT_HALF: { + case DataType::DT_HALF: { Reorder4( {k, c, r, s}, static_cast(iweights.GetValues()), istrides, @@ -842,7 +849,7 @@ void ReorderRSCKToKCRS(const TRT_ShapedWeights& iweights, } } -TRT_ShapedWeights TrtWeightStore::GetTempWeights(tensorflow::DataType type, +TRT_ShapedWeights TrtWeightStore::GetTempWeights(DataType type, const nvinfer1::Dims& dims) { TensorShape shape; // TODO(laigd): make it return a status. @@ -854,6 +861,13 @@ TRT_ShapedWeights TrtWeightStore::GetTempWeights(tensorflow::DataType type, return weights; } +const std::set* TrtNodeValidator::quantize_ops = new std::set{ + "QuantizeAndDequantizeV2", + "QuantizeAndDequantizeV3", + "FakeQuantWithMinMaxVars", + "FakeQuantWithMinMaxArgs", +}; + TrtNodeValidator::TrtNodeValidator() { RegisterOpValidators(); } Status TrtNodeValidator::ConvertToTensorOrWeights( @@ -899,9 +913,27 @@ Status TrtNodeValidator::ConvertToTensorOrWeights( } Status TrtNodeValidator::ValidateNode( - const tensorflow::NodeDef& node_def, + const NodeDef& node_def, const std::vector>& input_node_and_ports, + const TrtPrecisionMode precision_mode, const grappler::GraphProperties& graph_properties) { + const string& op = node_def.op(); + // It doesn't support validation of plugins. + if (PluginFactoryTensorRT::GetInstance()->IsPlugin(op)) return Status::OK(); + + // In INT8 mode, we will always apply the quantization ranges provided by + // these ops to the relevant tensors. This happens regardless of the value of + // use_calibration. + bool is_supported_op = false; + if (quantize_ops->count(op)) { + is_supported_op = (precision_mode == TrtPrecisionMode::INT8); + } else { + is_supported_op = op_validators_.count(node_def.op()); + } + if (!is_supported_op) { + return errors::Unimplemented("Op type ", op, " is not supported."); + } + // Convert input NodeDef and corresponding output ports to // TRT_TensorOrWeights. std::vector inputs; @@ -918,14 +950,7 @@ Status TrtNodeValidator::ValidateNode( inputs.push_back(tensor_or_weights); } - // Validate the node. - const auto iter = op_validators_.find(node_def.op()); - if (iter == op_validators_.end()) { - // If validator is not registered, it means no validation is needed. - return Status::OK(); - } - - OpConverter validator = iter->second; + OpConverter validator = op_validators_[node_def.op()]; OpConverterParams params( /*arg_converter=*/nullptr, node_def, inputs, /*arg_outputs=*/nullptr, /*arg_validation_only=*/true, &weight_store_); @@ -964,7 +989,7 @@ Status Converter::ConvertNode(const NodeDef& node_def) { TF_RETURN_IF_ERROR(plugin_converter_(¶ms)); } else { if (!op_registry_.count(op)) { - return errors::Unimplemented("No converter registered for op: " + op); + return errors::Unimplemented("No converter registered for op: ", op); } OpConverter op_converter = op_registry_.at(op); TF_RETURN_IF_ERROR(op_converter(¶ms)); @@ -1112,11 +1137,11 @@ Status Converter::TransposeTensor(nvinfer1::ITensor* input_tensor, const auto dims = input_tensor->getDimensions(); if (order_with_batch_dim.size() - 1 != size_t(dims.nbDims)) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Rank of perm for transpose does not match with that of the input."); } if (order_with_batch_dim[0] != 0) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "Transpose at batch dimension is not supported."); } @@ -1142,7 +1167,7 @@ Status Converter::TransposeTensor(nvinfer1::ITensor* input_tensor, layer->setReshapeDimensions(reshape_dims); *output_tensor = layer->getOutput(0); - return tensorflow::Status::OK(); + return Status::OK(); } Status Converter::GetWeightRange(const TRT_ShapedWeights& weights, @@ -1179,6 +1204,7 @@ Status Converter::GetWeightRange(const TRT_ShapedWeights& weights, Status Converter::PrepareTensorForShape(const TRT_TensorOrWeights& input, const nvinfer1::Dims& dims, + const bool validation_only, const nvinfer1::ITensor** tensor) { // If -1 is not used for one of the dims, we can check if the shapes are // compatible. @@ -1195,6 +1221,10 @@ Status Converter::PrepareTensorForShape(const TRT_TensorOrWeights& input, DebugString(input.GetTrtDims()), " vs ", DebugString(dims), ")"); } + if (validation_only) { + *tensor = nullptr; + return Status::OK(); + } if (input.is_tensor()) { if (DimsEqual(input.GetTrtDims(), dims)) { @@ -1230,7 +1260,7 @@ Status Converter::PrepareTensorForShape(const TRT_TensorOrWeights& input, min_range, max_range); } } - return tensorflow::Status::OK(); + return Status::OK(); } void Converter::MarkQuantizationRangesAsInferrable(nvinfer1::ITensor* input, @@ -1328,7 +1358,7 @@ void Converter::PropagateQuantizationRanges() { } } -Status Converter::GetInputs(const tensorflow::NodeDef& node_def, +Status Converter::GetInputs(const NodeDef& node_def, std::vector* inputs) const { for (auto const& input_name : node_def.input()) { /************************************************************************* @@ -1363,43 +1393,43 @@ Status Converter::GetInputs(const tensorflow::NodeDef& node_def, StrAppend(&msg, node_def.name(), " should have an input named '", name, "' but it is not available"); LOG(ERROR) << msg; - return tensorflow::errors::InvalidArgument(msg); + return errors::InvalidArgument(msg); } } - return tensorflow::Status::OK(); + return Status::OK(); } // Checks that the number of inputs match, and enforces that the inputs marked // as true are constant weights. true means that the input must be a weight, // while false means the input must be a tensor. In the future, false will mean // the input can be a tensor or weight. -tensorflow::Status CheckInputsWeights( +Status CheckInputsWeights( const OpConverterParams& params, const std::vector>& inputs_is_weight) { const auto& inputs = params.inputs; const auto& node_def = params.node_def; if (inputs.size() != inputs_is_weight.size()) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( node_def.op(), " got ", inputs.size(), " inputs but expected ", inputs_is_weight.size(), ", at ", node_def.name()); } for (int i = 0; i < inputs.size(); i++) { if (inputs_is_weight[i].second && inputs.at(i).is_tensor()) { - return tensorflow::errors::Unimplemented( - "The input \"", inputs_is_weight[i].first, "\" for ", node_def.op(), - " must be a constant, at ", node_def.name()); + return errors::Unimplemented("The input \"", inputs_is_weight[i].first, + "\" for ", node_def.op(), + " must be a constant, at ", node_def.name()); } // TODO(tmorris): Remove this check and provide a method to automatically // retrive an input as a tensor, converting via CreateConstantLayer if it // was originally a weight. We will want a caching mechanism to prevent many // duplicate constants from being created. if (!inputs_is_weight[i].second && inputs.at(i).is_weights()) { - return tensorflow::errors::Unimplemented( - "The input \"", inputs_is_weight[i].first, "\" for ", node_def.op(), - " must be a tensor, at ", node_def.name()); + return errors::Unimplemented("The input \"", inputs_is_weight[i].first, + "\" for ", node_def.op(), + " must be a tensor, at ", node_def.name()); } } - return tensorflow::Status::OK(); + return Status::OK(); } tensorflow::Status AllowDataTypes( @@ -1429,7 +1459,7 @@ tensorflow::Status AllowDataTypes( TRT_ShapedWeights ConvertFP32ToFP16(TrtWeightStore* store, const TRT_ShapedWeights& weights_src) { - auto dtype_new = tensorflow::DataType::DT_HALF; + auto dtype_new = DataType::DT_HALF; TRT_ShapedWeights weights = store->GetTempWeights(dtype_new, weights_src.shape_); const float* src = static_cast(weights_src.GetValues()); @@ -1488,18 +1518,17 @@ std::function LambdaFactory::unary() { } } -tensorflow::Status UnaryCompute(const TRT_ShapedWeights& iweights, - TRT_ShapedWeights* oweights, - LambdaFactory unary_op) { +Status UnaryCompute(const TRT_ShapedWeights& iweights, + TRT_ShapedWeights* oweights, LambdaFactory unary_op) { CHECK_EQ(iweights.type_, oweights->type_); switch (iweights.type_) { - case tensorflow::DataType::DT_FLOAT: { + case DataType::DT_FLOAT: { auto inp = static_cast(iweights.GetValues()); auto oup = static_cast(const_cast(oweights->GetValues())); std::transform(inp, inp + iweights.count(), oup, unary_op.unary()); break; } - case tensorflow::DataType::DT_HALF: { + case DataType::DT_HALF: { auto inp = static_cast(iweights.GetValues()); auto oup = static_cast(const_cast(oweights->GetValues())); @@ -1508,11 +1537,10 @@ tensorflow::Status UnaryCompute(const TRT_ShapedWeights& iweights, break; } default: - return tensorflow::errors::Unimplemented( - "Data type not supported: " + - tensorflow::DataTypeString(iweights.type_)); + return errors::Unimplemented("Data type not supported: " + + DataTypeString(iweights.type_)); } - return tensorflow::Status::OK(); + return Status::OK(); } // If swapped_inputs is false, 'tensor' is the left operand and 'weights' is the @@ -1709,11 +1737,11 @@ Status BinaryTensorOpWeight(OpConverterParams* params, // Pass the output params->outputs->push_back( TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, int group, - bool is_conv2d_backprop_input) { +Status ConvertConv2DHelper(OpConverterParams* params, int group, + bool is_conv2d_backprop_input) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TRT_TensorOrWeights backprop_output_size; @@ -1735,45 +1763,45 @@ tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, int group, tensorflow::DataType::DT_HALF})); TRT_ShapedWeights weights_rsck = inputs.at(1).weights(); if (weights_rsck.shape_.nbDims != 4) { - return tensorflow::errors::InvalidArgument( - "Conv2D expects kernel of dimension 4, at " + node_def.name()); + return errors::InvalidArgument("Conv2D expects kernel of dimension 4, at " + + node_def.name()); } TFAttrs attrs(node_def); auto data_format = attrs.get("data_format"); int c_index = (data_format == "NHWC") ? 3 : 1; int h_index = (data_format == "NHWC") ? 1 : 2; int w_index = (data_format == "NHWC") ? 2 : 3; - auto tf_dilations = attrs.get>("dilations"); + auto tf_dilations = attrs.get>("dilations"); if (tf_dilations.size() != 4) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Convolution dilations field must specify 4 dimensions, at ", node_def.name()); } if (tf_dilations[0] != 1 || tf_dilations[c_index] != 1) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "Dilation rate must be 1 for batch and channel dimensions, at ", node_def.name()); } const nvinfer1::DimsHW dilation(tf_dilations[h_index], tf_dilations[w_index]); if (is_conv2d_backprop_input && (dilation.d[0] != 1 || dilation.d[1] != 1)) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "Dilation with Conv2DBackpropInput (conv2d_transpose) is not supported", ", at ", node_def.name()); } - const auto tf_stride = attrs.get>("strides"); + const auto tf_stride = attrs.get>("strides"); if (tf_stride.size() != 4) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Convolution strides field must specify 4 dimensions, at ", node_def.name()); } if (tf_stride[0] != 1 || tf_stride[c_index] != 1) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "Stride must be 1 for batch and channel dimensions, at ", node_def.name()); } const nvinfer1::DimsHW stride(tf_stride[h_index], tf_stride[w_index]); - if (params->validation_only) return tensorflow::Status::OK(); + if (params->validation_only) return Status::OK(); // Transpose to NCHW (NCHW is required for IConvLayer). const bool need_transpose = (data_format == "NHWC"); @@ -1881,7 +1909,7 @@ tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, int group, } params->outputs->push_back( TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); + return Status::OK(); } Status BinaryTensorOpTensor(OpConverterParams* params, @@ -1923,10 +1951,10 @@ Status BinaryTensorOpTensor(OpConverterParams* params, const nvinfer1::ITensor* tensor_l = nullptr; const nvinfer1::ITensor* tensor_r = nullptr; status = params->converter->PrepareTensorForShape( - operand_l, broadcasted_dims_l, &tensor_l); + operand_l, broadcasted_dims_l, /*validation_only=*/false, &tensor_l); if (status.ok()) { status = params->converter->PrepareTensorForShape( - operand_r, broadcasted_dims_r, &tensor_r); + operand_r, broadcasted_dims_r, /*validation_only=*/false, &tensor_r); } if (!status.ok()) { return errors::Internal("Failed to convert binary op ", node_def.name(), @@ -1949,10 +1977,10 @@ Status BinaryTensorOpTensor(OpConverterParams* params, // Pass the output params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertPlugin(OpConverterParams* params) { +Status ConvertPlugin(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; // prepare input @@ -1977,7 +2005,7 @@ tensorflow::Status ConvertPlugin(OpConverterParams* params) { size_t size_data = data.size() * sizeof(float); if (!plugin->SetAttribute(attr_key, static_cast(data.data()), size_data)) { - return tensorflow::errors::InvalidArgument("plugin SetAttribute failed"); + return errors::InvalidArgument("plugin SetAttribute failed"); } } @@ -1988,10 +2016,10 @@ tensorflow::Status ConvertPlugin(OpConverterParams* params) { nvinfer1::ITensor* output_tensor = layer->getOutput(i); params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); } - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertTranspose(OpConverterParams* params) { +Status ConvertTranspose(OpConverterParams* params) { const auto& inputs = params->inputs; TF_RETURN_IF_ERROR( CheckInputsWeights(*params, {{"x", false}, {"perm", true}})); @@ -2024,10 +2052,10 @@ tensorflow::Status ConvertTranspose(OpConverterParams* params) { params->converter->TransposeTensor(input_tensor, perm, &output_tensor)); params->outputs->push_back( TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertReshape(OpConverterParams* params) { +Status ConvertReshape(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR( @@ -2038,8 +2066,8 @@ tensorflow::Status ConvertReshape(OpConverterParams* params) { TRT_TensorOrWeights input_tensor = inputs.at(0); TRT_ShapedWeights weights = inputs.at(1).weights(); if (weights.count() == 0) { - return tensorflow::errors::Unimplemented( - "Reshape to shape=[] is not supported, at ", node_def.name()); + return errors::Unimplemented("Reshape to shape=[] is not supported, at ", + node_def.name()); } const int* weights_ptr = @@ -2121,13 +2149,13 @@ tensorflow::Status ConvertReshape(OpConverterParams* params) { // Start conversion. const nvinfer1::ITensor* output_tensor = nullptr; TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( - input_tensor, reshape_dims, &output_tensor)); + input_tensor, reshape_dims, /*validation_only=*/false, &output_tensor)); params->outputs->push_back( TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertExpandDims(OpConverterParams* params) { +Status ConvertExpandDims(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR( @@ -2145,15 +2173,15 @@ tensorflow::Status ConvertExpandDims(OpConverterParams* params) { // Get axis to expand on. TRT_ShapedWeights weights = inputs.at(1).weights(); if (weights.count() != 1) { - return tensorflow::errors::InvalidArgument( - "ExpandDims axis must be a scalar, at ", node_def.name()); + return errors::InvalidArgument("ExpandDims axis must be a scalar, at ", + node_def.name()); } const int* weights_ptr = static_cast(const_cast(weights.GetValues())); int axis = weights_ptr[0]; // Make sure axis is valid. if ((axis < (-input_rank - 1)) || (axis > input_rank)) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Axis for ExpandDims is invalid, must be in the range " "[-rank(input) - 1, rank(input)], at ", node_def.name()); @@ -2161,7 +2189,7 @@ tensorflow::Status ConvertExpandDims(OpConverterParams* params) { // Convert negative axis to corresponding positive axis. if (axis < 0) axis += input_rank + 1; if (axis == 0) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "Modifying batch dimension is not supported for ExpandDims, at ", node_def.name()); } @@ -2175,13 +2203,13 @@ tensorflow::Status ConvertExpandDims(OpConverterParams* params) { /*ignore_first_dim=*/true)); const nvinfer1::ITensor* output_tensor = nullptr; TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( - input_tensor, new_dims, &output_tensor)); + input_tensor, new_dims, /*validation_only=*/false, &output_tensor)); params->outputs->push_back( TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertSqueeze(OpConverterParams* params) { +Status ConvertSqueeze(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"input", false}})); @@ -2197,15 +2225,15 @@ tensorflow::Status ConvertSqueeze(OpConverterParams* params) { const int input_rank = input_dims.size(); // Mark axes to remove by setting them to 0. TFAttrs attrs(node_def); - auto squeeze_dims = attrs.get>("squeeze_dims"); + auto squeeze_dims = attrs.get>("squeeze_dims"); if (squeeze_dims.empty()) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "Squeeze is only implemented for explicit dims, at ", node_def.name()); } for (int axis : squeeze_dims) { // Make sure axis is valid. if ((axis < -input_rank) || (axis >= input_rank)) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Axis for Squeeze is invalid, must be in the range " "[-rank(input), rank(input)), at ", node_def.name()); @@ -2214,12 +2242,12 @@ tensorflow::Status ConvertSqueeze(OpConverterParams* params) { if (axis < 0) axis += input_rank; // Don't squeeze batch dim. if (axis == 0) { - return tensorflow::errors::Unimplemented( - "Cannot squeeze batch dimension, at ", node_def.name()); + return errors::Unimplemented("Cannot squeeze batch dimension, at ", + node_def.name()); } // Make sure target dimension is size 1. if (input_dims[axis] != 1) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Cannot squeeze a dimension which isn't size 1, at ", node_def.name()); } @@ -2237,17 +2265,16 @@ tensorflow::Status ConvertSqueeze(OpConverterParams* params) { /*ignore_first_dim=*/true)); const nvinfer1::ITensor* output_tensor = nullptr; TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( - input_tensor, new_dims, &output_tensor)); + input_tensor, new_dims, /*validation_only=*/false, &output_tensor)); params->outputs->push_back( TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertStridedSliceHelper(OpConverterParams* params, - const TRT_TensorOrWeights& input, - std::vector begin, - std::vector size, - const std::vector& stride) { +Status ConvertStridedSliceHelper(OpConverterParams* params, + const TRT_TensorOrWeights& input, + std::vector begin, std::vector size, + const std::vector& stride) { const auto& node_def = params->node_def; // Get input dims. nvinfer1::Dims dims = input.GetTrtDims(); @@ -2257,20 +2284,19 @@ tensorflow::Status ConvertStridedSliceHelper(OpConverterParams* params, // Check bounds. for (int i = 1; i < input_dims.size(); i++) { if (begin[i] < 0 || begin[i] > input_dims[i]) { - return tensorflow::errors::InvalidArgument( - "\"begin\" for dimension ", std::to_string(i), " in ", node_def.op(), - " is out of range, at ", node_def.name()); + return errors::InvalidArgument("\"begin\" for dimension ", + std::to_string(i), " in ", node_def.op(), + " is out of range, at ", node_def.name()); } const int end = begin[i] + size[i]; if (end < 0 || end > input_dims[i]) { - return tensorflow::errors::InvalidArgument( - "\"begin\" + \"size\" for dimension ", std::to_string(i), " in ", - node_def.op(), " is out of range, at ", node_def.name()); + return errors::InvalidArgument("\"begin\" + \"size\" for dimension ", + std::to_string(i), " in ", node_def.op(), + " is out of range, at ", node_def.name()); } if (size[i] <= 0) { - return tensorflow::errors::InvalidArgument( - "\"size\" cannot be negative or zero for ", node_def.op(), ", at ", - node_def.name()); + return errors::InvalidArgument("\"size\" cannot be negative or zero for ", + node_def.op(), ", at ", node_def.name()); } } // TRT 5.1 adds a slice layer. For older versions, we attempt to use the @@ -2290,13 +2316,13 @@ tensorflow::Status ConvertStridedSliceHelper(OpConverterParams* params, *const_cast(input.tensor()), begin_dims, size_dims, stride_dims); params->outputs->push_back(TRT_TensorOrWeights(layer->getOutput(0))); - return tensorflow::Status::OK(); + return Status::OK(); #else // Use IPaddingLayer. // Strides must be 1 in this case. for (int x : stride) { if (x != 1) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "Strides other than 1 are not supported with this version of TRT, " "at ", node_def.name()); @@ -2304,11 +2330,11 @@ tensorflow::Status ConvertStridedSliceHelper(OpConverterParams* params, } // Rank must be 2, 3 or 4. if (input_dims.size() > 4) { - return tensorflow::errors::Unimplemented(node_def.op(), - " for tensors with rank > 4 is " - "not supported in this version of " - "TRT, at ", - node_def.name()); + return errors::Unimplemented(node_def.op(), + " for tensors with rank > 4 is " + "not supported in this version of " + "TRT, at ", + node_def.name()); } // Reshape if necessary to 4-D, since IPaddingLayer requires a 4-D input. const bool need_reshape = (input_dims.size() != 4); @@ -2343,7 +2369,7 @@ tensorflow::Status ConvertStridedSliceHelper(OpConverterParams* params, nvinfer1::IShuffleLayer* layer = params->converter->network()->addShuffle( *const_cast(input.tensor())); params->outputs->push_back(TRT_TensorOrWeights(layer->getOutput(0))); - return tensorflow::Status::OK(); + return Status::OK(); } else if (pad_dims.size() == 1) { // Only one dim is modified but we have to have 2, mark a second dim which // will have padding of 0. The dim we add is chosen to avoid an unecessary @@ -2354,7 +2380,7 @@ tensorflow::Status ConvertStridedSliceHelper(OpConverterParams* params, pad_dims.push_back(3); } } else if (pad_dims.size() > 2) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( node_def.op(), " can only modify up to 2 dimensions in this version of TRT, at ", node_def.name()); @@ -2391,7 +2417,7 @@ tensorflow::Status ConvertStridedSliceHelper(OpConverterParams* params, if (need_reshape) { const nvinfer1::ITensor* output_tensor = nullptr; TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( - input, reshape_dims, &output_tensor)); + input, reshape_dims, /*validation_only=*/false, &output_tensor)); tensor = const_cast(output_tensor); } if (need_transpose) { @@ -2425,8 +2451,8 @@ tensorflow::Status ConvertStridedSliceHelper(OpConverterParams* params, for (int i = 0; i < reshape_dims_added; i++) { int value = input_dims[1]; if (value != 1) { - return tensorflow::errors::Internal( - "StridedSlice error when reshaping, at ", node_def.name()); + return errors::Internal("StridedSlice error when reshaping, at ", + node_def.name()); } input_dims.erase(input_dims.begin() + 1); } @@ -2436,17 +2462,18 @@ tensorflow::Status ConvertStridedSliceHelper(OpConverterParams* params, /*ignore_first_dim=*/true)); const nvinfer1::ITensor* output_tensor = nullptr; TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( - TRT_TensorOrWeights(tensor), new_dims, &output_tensor)); + TRT_TensorOrWeights(tensor), new_dims, /*validation_only=*/false, + &output_tensor)); tensor = const_cast(output_tensor); } params->outputs->push_back( TRT_TensorOrWeights(const_cast(tensor))); - return tensorflow::Status::OK(); + return Status::OK(); #endif } -tensorflow::Status ConvertSlice(OpConverterParams* params) { +Status ConvertSlice(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights( @@ -2462,7 +2489,7 @@ tensorflow::Status ConvertSlice(OpConverterParams* params) { // Add batch dimension so that indexes line up properly. input_dims.insert(input_dims.begin(), inputs.at(0).batch_size()); if (!AllLengthsEqual({input_dims, begin, size})) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Length of begin and size arguments must equal rank of input for " "Slice, at ", node_def.name()); @@ -2477,7 +2504,7 @@ tensorflow::Status ConvertSlice(OpConverterParams* params) { size[0] != -1 && (!batch_size_is_defined || (batch_size_is_defined && size[0] != input_dims[0])); if (begin_is_modified || size_is_modified) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "TensorRT does not allow modifications to the batch dimension, at ", node_def.name()); } @@ -2492,7 +2519,7 @@ tensorflow::Status ConvertSlice(OpConverterParams* params) { return ConvertStridedSliceHelper(params, inputs.at(0), begin, size, stride); } -tensorflow::Status ConvertStridedSlice(OpConverterParams* params) { +Status ConvertStridedSlice(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights( @@ -2511,7 +2538,7 @@ tensorflow::Status ConvertStridedSlice(OpConverterParams* params) { std::vector end = inputs.at(2).weights().ToVector(); std::vector stride = inputs.at(3).weights().ToVector(); if (!AllLengthsEqual({input_dims, begin, end, stride})) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Length of begin, end, and stride arguments must equal rank of input " "for StridedSlice, at ", node_def.name()); @@ -2520,14 +2547,14 @@ tensorflow::Status ConvertStridedSlice(OpConverterParams* params) { TFAttrs attrs(node_def); for (const string& attr : {"ellipsis_mask", "new_axis_mask", "shrink_axis_mask"}) { - int attr_val = attrs.get(attr); + int attr_val = attrs.get(attr); if (attr_val != 0) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( attr, " is not supported for StridedSlice, at ", node_def.name()); } } - const int begin_mask = attrs.get("begin_mask"); - const int end_mask = attrs.get("end_mask"); + const int begin_mask = attrs.get("begin_mask"); + const int end_mask = attrs.get("end_mask"); // Check that batch dimension is unmodified. const bool begin_is_modified = !(begin_mask & 1) && begin[0] != 0; const bool stride_is_modified = stride[0] != 1; @@ -2539,7 +2566,7 @@ tensorflow::Status ConvertStridedSlice(OpConverterParams* params) { !(end_mask & 1) && (!batch_size_is_defined || (batch_size_is_defined && end[0] != input_dims[0])); if (begin_is_modified || stride_is_modified || end_is_modified) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "TensorRT does not allow modifications to the batch dimension, at ", node_def.name()); } @@ -2565,7 +2592,7 @@ tensorflow::Status ConvertStridedSlice(OpConverterParams* params) { // Negative or zero strides currently not supported. for (int i = 0; i < input_dims.size(); i++) { if (stride[i] <= 0) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "Negative or zero stride values are not supported for StridedSlice, " "at ", node_def.name()); @@ -2580,19 +2607,19 @@ tensorflow::Status ConvertStridedSlice(OpConverterParams* params) { return ConvertStridedSliceHelper(params, inputs.at(0), begin, size, stride); } -tensorflow::Status ConvertConv2D(OpConverterParams* params) { +Status ConvertConv2D(OpConverterParams* params) { return ConvertConv2DHelper(params, 1, /*is_conv2d_backprop_input=*/false); } -tensorflow::Status ConvertConv2DDepthwise(OpConverterParams* params) { +Status ConvertConv2DDepthwise(OpConverterParams* params) { return ConvertConv2DHelper(params, 0, /*is_conv2d_backprop_input=*/false); } -tensorflow::Status ConvertConv2DBackpropInput(OpConverterParams* params) { +Status ConvertConv2DBackpropInput(OpConverterParams* params) { return ConvertConv2DHelper(params, 1, /*is_conv2d_backprop_input=*/true); } -tensorflow::Status ConvertPool(OpConverterParams* params) { +Status ConvertPool(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"input", false}})); @@ -2604,14 +2631,14 @@ tensorflow::Status ConvertPool(OpConverterParams* params) { } else if (node_def.op() == "AvgPool") { type = nvinfer1::PoolingType::kAVERAGE; } else { - return tensorflow::errors::Unimplemented( - "Unsupported pooling type: ", node_def.op(), ", at ", node_def.name()); + return errors::Unimplemented("Unsupported pooling type: ", node_def.op(), + ", at ", node_def.name()); } TFAttrs attrs(node_def); const string padding_type = attrs.get("padding"); if ((padding_type != "SAME") && (padding_type != "VALID")) { - return tensorflow::errors::Unimplemented( - "Unsupported padding type: ", padding_type, ", at ", node_def.name()); + return errors::Unimplemented("Unsupported padding type: ", padding_type, + ", at ", node_def.name()); } if (params->validation_only) return Status::OK(); @@ -2626,10 +2653,10 @@ tensorflow::Status ConvertPool(OpConverterParams* params) { const_cast(tensor), {0, 3, 1, 2}, &tensor)); } - const auto tf_stride = attrs.get>("strides"); + const auto tf_stride = attrs.get>("strides"); const nvinfer1::DimsHW stride(tf_stride[h_index], tf_stride[w_index]); - const auto tf_kernel = attrs.get>("ksize"); + const auto tf_kernel = attrs.get>("ksize"); const nvinfer1::DimsHW ksize(tf_kernel[h_index], tf_kernel[w_index]); auto tensor_dim = tensor->getDimensions(); @@ -2681,25 +2708,26 @@ tensorflow::Status ConvertPool(OpConverterParams* params) { } params->outputs->push_back( TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); + return Status::OK(); } // TODO(tmorris): Use ActivationType::kLEAKY_RELU in TRT 5.1+ once perf // improves. -tensorflow::Status ConvertLeakyRelu(OpConverterParams* params) { +Status ConvertLeakyRelu(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"input", false}})); TF_RETURN_IF_ERROR(AllowDataTypes(*params, {tensorflow::DataType::DT_FLOAT, tensorflow::DataType::DT_HALF})); + TFAttrs attrs(node_def); const float alpha = attrs.get("alpha"); if (alpha < 0.0f || alpha > 1.0f) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "Alpha value for LeakyRelu must be between 0 and 1, at ", node_def.name()); } - if (params->validation_only) return tensorflow::Status::OK(); + if (params->validation_only) return Status::OK(); // Input Tensor const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); @@ -2729,7 +2757,7 @@ tensorflow::Status ConvertLeakyRelu(OpConverterParams* params) { return Status::OK(); } -tensorflow::Status ConvertActivation(OpConverterParams* params) { +Status ConvertActivation(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"input", false}})); @@ -2742,11 +2770,10 @@ tensorflow::Status ConvertActivation(OpConverterParams* params) { }; auto op_pair = ops.find(node_def.op()); if (op_pair == ops.end()) { - return tensorflow::errors::Unimplemented( - "Activation op: ", node_def.op(), - " not supported at: ", node_def.name()); + return errors::Unimplemented("Activation op: ", node_def.op(), + " not supported at: ", node_def.name()); } - if (params->validation_only) return tensorflow::Status::OK(); + if (params->validation_only) return Status::OK(); // Start conversion. const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); @@ -2762,7 +2789,7 @@ tensorflow::Status ConvertActivation(OpConverterParams* params) { params->converter->ProvideQuantizationRange(output_tensor, -1.0f, 1.0f); } params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); + return Status::OK(); } Status ConvertQuantize(OpConverterParams* params) { @@ -2828,7 +2855,7 @@ Status ConvertQuantize(OpConverterParams* params) { } // TODO(tmorris): Use ActivationType::kCLIP in TRT 5.1+ once perf improves. -tensorflow::Status ConvertRelu6(OpConverterParams* params) { +Status ConvertRelu6(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"input", false}})); @@ -2878,7 +2905,7 @@ tensorflow::Status ConvertRelu6(OpConverterParams* params) { return Status::OK(); } -tensorflow::Status ConvertBiasAdd(OpConverterParams* params) { +Status ConvertBiasAdd(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR( @@ -3052,7 +3079,7 @@ Status TfTensorToTrtWeights(const Tensor& tensor, TrtWeightStore* weight_store, // weights to params->outputs. We did this since TrtNodeValidator needs the // weights as input to other nodes, and use it to determine whether those nodes // are supported by TRT. -tensorflow::Status ConvertConst(OpConverterParams* params) { +Status ConvertConst(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; if (!inputs.empty()) { @@ -3063,14 +3090,14 @@ tensorflow::Status ConvertConst(OpConverterParams* params) { // Create shaped weights as output const auto& tensor_proto = node_def.attr().at("value").tensor(); - tensorflow::Tensor tensor; + Tensor tensor; if (!tensor.FromProto(tensor_proto)) { - return tensorflow::errors::Internal("Cannot parse weight tensor proto: ", - node_def.name()); + return errors::Internal("Cannot parse weight tensor proto: ", + node_def.name()); } TFAttrs attrs(node_def); - const DataType dtype = attrs.get("dtype"); + const DataType dtype = attrs.get("dtype"); if (dtype != tensor.dtype()) { return errors::InvalidArgument("DataType mismatch between attr (", DataTypeString(dtype), ") and tensor (", @@ -3087,12 +3114,13 @@ tensorflow::Status ConvertConst(OpConverterParams* params) { return Status::OK(); } -tensorflow::Status ConvertIdentity(OpConverterParams* params) { +Status ConvertIdentity(OpConverterParams* params) { // TODO(tmorris): TRT's Identity layer does not get optimized away as of TRT // 5.0, however once we know that it does it would be nice to use that // instead. + if (params->validation_only) return Status::OK(); params->outputs->push_back(params->inputs.at(0)); - return tensorflow::Status::OK(); + return Status::OK(); } Status ConvertBinary(OpConverterParams* params) { @@ -3102,9 +3130,9 @@ Status ConvertBinary(OpConverterParams* params) { // TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"x", false}, {"y", // false}})); if (inputs.size() != 2) { - return tensorflow::errors::InvalidArgument( - node_def.op(), " got ", inputs.size(), " inputs but expected 2, at ", - node_def.name()); + return errors::InvalidArgument(node_def.op(), " got ", inputs.size(), + " inputs but expected 2, at ", + node_def.name()); } TF_RETURN_IF_ERROR(AllowDataTypes(*params, {tensorflow::DataType::DT_FLOAT, tensorflow::DataType::DT_HALF})); @@ -3139,19 +3167,19 @@ Status ConvertBinary(OpConverterParams* params) { // If both input are tensors, or one of them is weights but the conversion // above failed, try the conversion using BinaryTensorOpTensor. if ((inputs.at(0).is_tensor() && inputs.at(1).is_tensor()) || !status.ok()) { - if (!status.ok()) VLOG(1) << status; + if (!status.ok()) VLOG(2) << status; status = BinaryTensorOpTensor(params, inputs.at(0), inputs.at(1)); } return status; } -tensorflow::Status ConvertRsqrt(OpConverterParams* params) { +Status ConvertRsqrt(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"x", false}})); TF_RETURN_IF_ERROR(AllowDataTypes(*params, {tensorflow::DataType::DT_FLOAT, tensorflow::DataType::DT_HALF})); - if (params->validation_only) return tensorflow::Status::OK(); + if (params->validation_only) return Status::OK(); // TODO(tmorris): params->converter is null during validation. Allow // precision_mode and use_calibration to be accessed during validation and @@ -3181,7 +3209,7 @@ tensorflow::Status ConvertRsqrt(OpConverterParams* params) { *sqrt_layer->getOutput(0), nvinfer1::UnaryOperation::kRECIP); TFTRT_RETURN_ERROR_IF_NULLPTR(recip_layer, node_def.name()); params->outputs->push_back(TRT_TensorOrWeights(recip_layer->getOutput(0))); - return tensorflow::Status::OK(); + return Status::OK(); } const std::unordered_map* @@ -3213,7 +3241,7 @@ UnaryOperationMap() { return m; } -tensorflow::Status ConvertUnary(OpConverterParams* params) { +Status ConvertUnary(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"x", false}})); @@ -3221,10 +3249,10 @@ tensorflow::Status ConvertUnary(OpConverterParams* params) { tensorflow::DataType::DT_HALF})); auto op_pair = UnaryOperationMap()->find(node_def.op()); if (op_pair == UnaryOperationMap()->end()) { - return tensorflow::errors::Unimplemented( - "Unary op: ", node_def.op(), " not supported at: ", node_def.name()); + return errors::Unimplemented("Unary op: ", node_def.op(), + " not supported at: ", node_def.name()); } - if (params->validation_only) return tensorflow::Status::OK(); + if (params->validation_only) return Status::OK(); // Start conversion. const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); @@ -3249,10 +3277,10 @@ tensorflow::Status ConvertUnary(OpConverterParams* params) { } params->outputs->push_back( TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertSquare(OpConverterParams* params) { +Status ConvertSquare(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"x", false}})); @@ -3275,10 +3303,10 @@ tensorflow::Status ConvertSquare(OpConverterParams* params) { nvinfer1::ITensor* output_tensor = layer->getOutput(0); params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertReduce(OpConverterParams* params) { +Status ConvertReduce(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR( @@ -3290,16 +3318,14 @@ tensorflow::Status ConvertReduce(OpConverterParams* params) { TRT_ShapedWeights index_list = inputs.at(1).weights(); TFAttrs attrs(node_def); - auto index_type = attrs.get("Tidx"); - // Only expect to handle INT32 as attributes for now - if (index_type != tensorflow::DataType::DT_INT32) { - return tensorflow::errors::Unimplemented("Tidx supports only DT_INT32"); + if (attrs.get("Tidx") != DataType::DT_INT32) { + return errors::Unimplemented("Tidx supports only DT_INT32"); } int axes = 0; if (index_list.count() == 0) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "TRT cannot support reduce on all (batch) dimensions, at", node_def.name()); } else { @@ -3309,7 +3335,7 @@ tensorflow::Status ConvertReduce(OpConverterParams* params) { int axis = index_list_data[i]; if (axis < 0) axis += tensor->getDimensions().nbDims + 1; if (axis == 0) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "TRT cannot reduce at batch dimension, at", node_def.name()); } axes |= (1 << (axis - 1)); @@ -3328,9 +3354,10 @@ tensorflow::Status ConvertReduce(OpConverterParams* params) { } else if (node_def.op() == "Mean") { reduce_operation = nvinfer1::ReduceOperation::kAVG; } else { - return tensorflow::errors::Unimplemented("Op not supported ", node_def.op(), - " , at ", node_def.name()); + return errors::Unimplemented("Op not supported ", node_def.op(), ", at ", + node_def.name()); } + if (params->validation_only) return Status::OK(); const auto keep_dims = attrs.get("keep_dims"); nvinfer1::ILayer* layer = params->converter->network()->addReduce( @@ -3339,10 +3366,10 @@ tensorflow::Status ConvertReduce(OpConverterParams* params) { TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); params->outputs->push_back(TRT_TensorOrWeights(layer->getOutput(0))); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertPad(OpConverterParams* params) { +Status ConvertPad(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR( @@ -3361,19 +3388,18 @@ tensorflow::Status ConvertPad(OpConverterParams* params) { TFAttrs attrs(node_def); // Padding type here is done through TF type // so I can leverage their EnumToDataType for my cast - auto padding_type = attrs.get("Tpaddings"); + auto padding_type = attrs.get("Tpaddings"); // TODO(jie): handle data type conversion for TRT? if (pads.shape_.d[0] != nb_dims || pads.shape_.d[1] != 2) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Pad only supports explicit padding on 4 dimensional tensor, at ", node_def.name()); } // Only expect to handle INT32 as attributes for now - if (padding_type != tensorflow::DataType::DT_INT32) { - return tensorflow::errors::Unimplemented( - "Tpaddings supports only DT_INT32"); + if (padding_type != DataType::DT_INT32) { + return errors::Unimplemented("Tpaddings supports only DT_INT32"); } auto pad_data = static_cast(const_cast(pads.GetValues())); @@ -3387,25 +3413,25 @@ tensorflow::Status ConvertPad(OpConverterParams* params) { // No padding at all, we should exit if (pad_index.empty()) { params->outputs->push_back(inputs.at(0)); - return tensorflow::Status::OK(); + return Status::OK(); } // Only supports padding on less than 2 axis GIE-2579 if (pad_index.size() > 2) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Padding layer does not support padding on > 2"); } // Padding on batch dimension is not supported if (pad_index[0] == 0) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Padding layer does not support padding on batch dimension"); } // Not doing the legit thing here. ignoring padding on dim 1 and 3; // TODO(jie): implement pad as uff parser if (pad_index.size() == 2 && pad_index[0] == 0 && pad_index[1] == 3) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "Padding layer does not support padding on dimension 1 and 3 yet"); } if (params->validation_only) return Status::OK(); @@ -3446,10 +3472,10 @@ tensorflow::Status ConvertPad(OpConverterParams* params) { params->outputs->push_back( TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertConcat(OpConverterParams* params) { +Status ConvertConcat(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; // TODO(tmorris): There is a bug with Concat and INT32 in TRT - it is supposed @@ -3460,7 +3486,7 @@ tensorflow::Status ConvertConcat(OpConverterParams* params) { int input_size = static_cast(inputs.size()) - 1; if (!inputs.at(0).is_tensor()) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Concat in TRT support only Tensor input, at ", node_def.name()); } @@ -3468,13 +3494,13 @@ tensorflow::Status ConvertConcat(OpConverterParams* params) { TRT_ShapedWeights axis = inputs.at(input_size).weights(); TFAttrs attrs(node_def); - auto index_type = attrs.get("Tidx"); + auto index_type = attrs.get("Tidx"); // TODO(jie): handle data type // Only expect to handle INT32 as index attributes for now - if (index_type != tensorflow::DataType::DT_INT32) - return tensorflow::errors::Unimplemented("Tidx supports only DT_INT32, at ", - node_def.name()); + if (index_type != DataType::DT_INT32) + return errors::Unimplemented("Tidx supports only DT_INT32, at ", + node_def.name()); int index = *(static_cast(const_cast(axis.GetValues()))); @@ -3483,11 +3509,11 @@ tensorflow::Status ConvertConcat(OpConverterParams* params) { auto dim = inputs.at(0).tensor()->getDimensions(); // dimension check if (index > dim.nbDims + 1) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Concatenate on axis out of dimension range, at ", node_def.name()); } if (index == 0) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Concatenate on batch dimension not supported, at ", node_def.name()); } if (index < 0) { @@ -3501,14 +3527,14 @@ tensorflow::Status ConvertConcat(OpConverterParams* params) { auto tensor_i = inputs.at(i).tensor(); auto dim_i = tensor_i->getDimensions(); if (dim_i.nbDims != dim.nbDims) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Concatenate receives inputs with inconsistent dimensions, at ", node_def.name()); } for (int j = 0; j < dim.nbDims; j++) { // check dimension consistency on non-concatenate axis if (j != index - 1 && dim_i.d[j] != dim.d[j]) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Concatenate receives inputs with inconsistent shape, at", node_def.name()); } @@ -3516,7 +3542,7 @@ tensorflow::Status ConvertConcat(OpConverterParams* params) { inputs_vec.push_back(tensor_i); } - if (params->validation_only) return tensorflow::Status::OK(); + if (params->validation_only) return Status::OK(); // nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); nvinfer1::IConcatenationLayer* layer = @@ -3527,10 +3553,10 @@ tensorflow::Status ConvertConcat(OpConverterParams* params) { layer->setAxis(index - 1); nvinfer1::ITensor* output_tensor = layer->getOutput(0); params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertFusedBatchNorm(OpConverterParams* params) { +Status ConvertFusedBatchNorm(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"x", false}, @@ -3544,7 +3570,7 @@ tensorflow::Status ConvertFusedBatchNorm(OpConverterParams* params) { float epsilon = attrs.get("epsilon"); auto data_format = attrs.get("data_format"); if (data_format != "NCHW") { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( node_def.op(), " only supports data_format=NCHW, at ", node_def.name()); } bool is_training = attrs.get("is_training"); @@ -3556,23 +3582,23 @@ tensorflow::Status ConvertFusedBatchNorm(OpConverterParams* params) { << "are using Keras, please call " << "keras.backend.set_learning_phase(0) before constructing " << "your model. At " << node_def.name(); - return tensorflow::errors::Unimplemented( - node_def.op(), " only supports is_training=false, at ", - node_def.name()); + return errors::Unimplemented(node_def.op(), + " only supports is_training=false, at ", + node_def.name()); } nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); // Check parameter types auto parameter_type = inputs.at(1).weights().type_; - if ((parameter_type != tensorflow::DataType::DT_FLOAT) && - (parameter_type != tensorflow::DataType::DT_HALF)) { - return tensorflow::errors::Unimplemented( + if ((parameter_type != DataType::DT_FLOAT) && + (parameter_type != DataType::DT_HALF)) { + return errors::Unimplemented( "only float32 or float16 weight data type is supported, for node " + - node_def.name() + " got " + tensorflow::DataTypeString(parameter_type)); + node_def.name() + " got " + DataTypeString(parameter_type)); } for (int i = 1; i < 5; i++) { if (inputs.at(i).weights().type_ != parameter_type) { - return tensorflow::errors::Unimplemented( + return errors::Unimplemented( "Inconsistent parameter type for batchnorm is not supported, at: " + node_def.name()); } @@ -3589,7 +3615,7 @@ tensorflow::Status ConvertFusedBatchNorm(OpConverterParams* params) { ptr_shape_weights = const_cast(&(inputs.at(i).weights())); } else if (inputs.at(i).weights().count() != 1) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Inconsistent batchnorm parameter count, at: " + node_def.name()); } } @@ -3623,16 +3649,16 @@ tensorflow::Status ConvertFusedBatchNorm(OpConverterParams* params) { float batchnorm_data[4]; for (int j = 0; j < 4; j++) { if (inputs.at(j + 1).weights().count() != 1) { - if (parameter_type == tensorflow::DT_FLOAT) { + if (parameter_type == DT_FLOAT) { batchnorm_data[j] = vals_array[j][i]; - } else if (parameter_type == tensorflow::DT_HALF) { + } else if (parameter_type == DT_HALF) { batchnorm_data[j] = Eigen::half_impl::half_to_float(cast_vals_array[j][i]); } } else { - if (parameter_type == tensorflow::DT_FLOAT) { + if (parameter_type == DT_FLOAT) { batchnorm_data[j] = vals_array[j][0]; - } else if (parameter_type == tensorflow::DT_HALF) { + } else if (parameter_type == DT_HALF) { batchnorm_data[j] = Eigen::half_impl::half_to_float(cast_vals_array[j][0]); } @@ -3644,10 +3670,10 @@ tensorflow::Status ConvertFusedBatchNorm(OpConverterParams* params) { float variance = batchnorm_data[3]; float combined_scale_val = scale / sqrtf(variance + epsilon); float combined_offset_val = offset - mean * combined_scale_val; - if (parameter_type == tensorflow::DT_FLOAT) { + if (parameter_type == DT_FLOAT) { combined_scale_vals[i] = combined_scale_val; combined_offset_vals[i] = combined_offset_val; - } else if (parameter_type == tensorflow::DT_HALF) { + } else if (parameter_type == DT_HALF) { cast_combined_scale_vals[i] = Eigen::half(combined_scale_val); cast_combined_offset_vals[i] = Eigen::half(combined_offset_val); } @@ -3663,10 +3689,10 @@ tensorflow::Status ConvertFusedBatchNorm(OpConverterParams* params) { TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); nvinfer1::ITensor* output_tensor = layer->getOutput(0); params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertGather(OpConverterParams* params) { +Status ConvertGather(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights( @@ -3676,8 +3702,8 @@ tensorflow::Status ConvertGather(OpConverterParams* params) { tensorflow::DataType::DT_INT32})); absl::Span axis = inputs.at(2).weights().GetSpan(); if (axis.size() != 1) { - return tensorflow::errors::InvalidArgument( - "Axis for GatherV2 must be a scalar, at ", node_def.name()); + return errors::InvalidArgument("Axis for GatherV2 must be a scalar, at ", + node_def.name()); } int trt_axis = 0; TF_RETURN_IF_ERROR(ConvertAxis(axis[0], inputs.at(0).GetTrtDims().nbDims, @@ -3692,14 +3718,13 @@ tensorflow::Status ConvertGather(OpConverterParams* params) { return Status::OK(); } -tensorflow::Status ConvertMatMulHelper(OpConverterParams* params, - TRT_TensorOrWeights tensor_input, - TRT_ShapedWeights weights_raw, - bool transpose_weight, - string node_name) { +Status ConvertMatMulHelper(OpConverterParams* params, + TRT_TensorOrWeights tensor_input, + TRT_ShapedWeights weights_raw, bool transpose_weight, + string node_name) { nvinfer1::ITensor* output_tensor; if (!tensor_input.is_tensor()) { - return tensorflow::errors::InvalidArgument("Input 0 expects tensor"); + return errors::InvalidArgument("Input 0 expects tensor"); } const nvinfer1::ITensor* tensor = tensor_input.tensor(); @@ -3719,7 +3744,7 @@ tensorflow::Status ConvertMatMulHelper(OpConverterParams* params, input_dim.d[input_dim.nbDims++] = 1; } TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( - tensor_input, input_dim, &tensor)); + tensor_input, input_dim, /*validation_only=*/false, &tensor)); nvinfer1::IFullyConnectedLayer* layer = params->converter->network()->addFullyConnected( @@ -3732,14 +3757,15 @@ tensorflow::Status ConvertMatMulHelper(OpConverterParams* params, auto output_dim = output_tensor->getDimensions(); output_dim.nbDims = 1; TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( - TRT_TensorOrWeights(output_tensor), output_dim, &temp_tensor)); + TRT_TensorOrWeights(output_tensor), output_dim, /*validation_only=*/false, + &temp_tensor)); output_tensor = const_cast(temp_tensor); params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); + return Status::OK(); } -// inputs are both two dimensional (tensorflow::ops::MatMul) -tensorflow::Status ConvertMatMul(OpConverterParams* params) { +// inputs are both two dimensional (ops::MatMul) +Status ConvertMatMul(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"a", false}, {"b", true}})); @@ -3761,7 +3787,7 @@ tensorflow::Status ConvertMatMul(OpConverterParams* params) { transpose_b, node_def.name()); } -tensorflow::Status ConvertBatchMatMul(OpConverterParams* params) { +Status ConvertBatchMatMul(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; // TODO(tmorris): Enable once false is updated to mean either tensor or weight @@ -3770,58 +3796,55 @@ tensorflow::Status ConvertBatchMatMul(OpConverterParams* params) { TF_RETURN_IF_ERROR(AllowDataTypes(*params, {tensorflow::DataType::DT_FLOAT, tensorflow::DataType::DT_HALF})); if (inputs.size() != 2) { - return tensorflow::errors::InvalidArgument( - node_def.op(), " got ", inputs.size(), " inputs but expected 2, at ", - node_def.name()); + return errors::InvalidArgument(node_def.op(), " got ", inputs.size(), + " inputs but expected 2, at ", + node_def.name()); + } + if (inputs[0].is_weights() && inputs[1].is_weights()) { + return errors::InvalidArgument( + "All inputs are weights, but Grappler is expected to fold them."); } TFAttrs attrs(node_def); - bool transpose_a = attrs.get("adj_x"); - bool transpose_b = attrs.get("adj_y"); - - auto dims = inputs.at(0).GetTrtDims(); + const bool transpose_a = attrs.get("adj_x"); + const bool transpose_b = attrs.get("adj_y"); + const auto dims = inputs.at(0).GetTrtDims(); if (dims.nbDims == 1) { // NC * CK is only supported through fully connected if (transpose_a == false && inputs.at(0).is_tensor() && inputs.at(1).is_weights()) { return ConvertMatMulHelper(params, inputs.at(0), inputs.at(1).weights(), transpose_b, node_def.name()); } else { - return tensorflow::errors::InvalidArgument( - "Invalid configuration for MatMul, at: " + node_def.name()); + return errors::InvalidArgument("Invalid configuration for MatMul, at: ", + node_def.name()); } } - const nvinfer1::ITensor* tensor_l; - const nvinfer1::ITensor* tensor_r; - auto dims_l = inputs.at(0).GetTrtDims(); - auto dims_r = inputs.at(1).GetTrtDims(); - if (inputs.at(0).is_weights()) { - if (inputs.at(0).GetTrtDims().d[0] != 1) { - return tensorflow::errors::InvalidArgument( - "Input 0 as weight assumes broadcast across batch for MatMul, at: " + - node_def.name()); - } else { - for (int i = 0; i < dims_l.nbDims - 1; i++) { - dims_l.d[i] = dims_l.d[i + 1]; - } - dims_l.nbDims--; - } - } - if (inputs.at(1).is_weights()) { - if (inputs.at(1).GetTrtDims().d[0] != 1) { - return tensorflow::errors::InvalidArgument( - "Input 1 as weight assumes broadcast across batch for MatMul, at: " + - node_def.name()); - } else { - for (int i = 0; i < dims_r.nbDims - 1; i++) { - dims_r.d[i] = dims_r.d[i + 1]; + auto get_tensor_with_proper_dims = [params]( + const TRT_TensorOrWeights& input, + const nvinfer1::ITensor** tensor) { + auto dims = input.GetTrtDims(); + if (input.is_weights()) { + // The other operand must be a tensor, this is ensured by earlier checks. + // Checks that the batch dimension is not changed by broadcasting. + if (dims.d[0] != 1) { + return errors::InvalidArgument( + "Input weight attempts to broadcast across batch dimension for " + "BatchMatMul, at ", + params->node_def.name()); } - dims_r.nbDims--; + // Remove the batch dimension from the weights. + TF_RETURN_IF_ERROR(RemoveBatchDimension(&dims)); } - } - TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( - inputs.at(0), dims_l, &tensor_l)); - TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( - inputs.at(1), dims_r, &tensor_r)); + // Create tensor and reshape if necessary. + TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( + input, dims, params->validation_only, tensor)); + return Status::OK(); + }; + const nvinfer1::ITensor* tensor_l; + const nvinfer1::ITensor* tensor_r; + TF_RETURN_IF_ERROR(get_tensor_with_proper_dims(inputs.at(0), &tensor_l)); + TF_RETURN_IF_ERROR(get_tensor_with_proper_dims(inputs.at(1), &tensor_r)); + if (params->validation_only) return Status::OK(); nvinfer1::IMatrixMultiplyLayer* layer = params->converter->network()->addMatrixMultiply( @@ -3830,10 +3853,10 @@ tensorflow::Status ConvertBatchMatMul(OpConverterParams* params) { TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); nvinfer1::ITensor* output_tensor = layer->getOutput(0); params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertSoftmax(OpConverterParams* params) { +Status ConvertSoftmax(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"logits", false}})); @@ -3843,8 +3866,8 @@ tensorflow::Status ConvertSoftmax(OpConverterParams* params) { int nbDims = tensor->getDimensions().nbDims; if (nbDims == 0) { - return tensorflow::errors::InvalidArgument( - "TensorRT Softmax cannot apply on batch dimension, at" + + return errors::InvalidArgument( + "TensorRT Softmax cannot apply on batch dimension, at", node_def.name()); } if (params->validation_only) return Status::OK(); @@ -3859,10 +3882,10 @@ tensorflow::Status ConvertSoftmax(OpConverterParams* params) { // Quantization range for SoftMax is always (0, 1) params->converter->ProvideQuantizationRange(output_tensor, 0.0f, 1.0f); params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertTopK(OpConverterParams* params) { +Status ConvertTopK(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; TF_RETURN_IF_ERROR( @@ -3897,12 +3920,11 @@ tensorflow::Status ConvertTopK(OpConverterParams* params) { nvinfer1::ITensor* output_indices_tensor = layer->getOutput(1); params->outputs->push_back(TRT_TensorOrWeights(output_value_tensor)); params->outputs->push_back(TRT_TensorOrWeights(output_indices_tensor)); - return tensorflow::Status::OK(); + return Status::OK(); } static void RegisterValidatableOpConverters( std::unordered_map* registration) { - // TODO(laigd): support all op types. (*registration)["BiasAdd"] = ConvertBiasAdd; (*registration)["ConcatV2"] = ConvertConcat; (*registration)["Const"] = ConvertConst; @@ -3924,6 +3946,18 @@ static void RegisterValidatableOpConverters( (*registration)["Transpose"] = ConvertTranspose; (*registration)["TopKV2"] = ConvertTopK; + // TODO(ben,jie): this is a temp hack. + (*registration)["Identity"] = ConvertIdentity; // Identity should be removed + (*registration)["Snapshot"] = ConvertIdentity; // Snapshot should be removed + + (*registration)["Sum"] = ConvertReduce; + (*registration)["Prod"] = ConvertReduce; + (*registration)["Max"] = ConvertReduce; + (*registration)["Min"] = ConvertReduce; + (*registration)["Mean"] = ConvertReduce; + (*registration)["Softmax"] = ConvertSoftmax; + (*registration)["BatchMatMul"] = ConvertBatchMatMul; + for (auto quantization_op_type : {"QuantizeAndDequantizeV2", "QuantizeAndDequantizeV3", "FakeQuantWithMinMaxVars", "FakeQuantWithMinMaxArgs"}) { @@ -3953,27 +3987,14 @@ void TrtNodeValidator::RegisterOpValidators() { void Converter::RegisterOpConverters() { RegisterValidatableOpConverters(&op_registry_); - // TODO(ben,jie): this is a temp hack. - op_registry_["Identity"] = ConvertIdentity; // Identity should be removed - op_registry_["Snapshot"] = ConvertIdentity; // Snapshot should be removed - - op_registry_["Sum"] = ConvertReduce; - op_registry_["Prod"] = ConvertReduce; - op_registry_["Max"] = ConvertReduce; - op_registry_["Min"] = ConvertReduce; - op_registry_["Mean"] = ConvertReduce; - op_registry_["Softmax"] = ConvertSoftmax; - op_registry_["BatchMatMul"] = ConvertBatchMatMul; - plugin_converter_ = ConvertPlugin; } -tensorflow::Status ConvertGraphDefToEngine( - const tensorflow::GraphDef& gdef, TrtPrecisionMode precision_mode, - int max_batch_size, size_t max_workspace_size_bytes, - const std::vector& input_shapes, - Logger* logger, nvinfer1::IGpuAllocator* allocator, - TRTInt8Calibrator* calibrator, +Status ConvertGraphDefToEngine( + const GraphDef& gdef, TrtPrecisionMode precision_mode, int max_batch_size, + size_t max_workspace_size_bytes, + const std::vector& input_shapes, Logger* logger, + nvinfer1::IGpuAllocator* allocator, TRTInt8Calibrator* calibrator, TrtUniquePtrType* engine, bool use_calibration, bool* convert_successfully) { engine->reset(); @@ -4004,8 +4025,7 @@ tensorflow::Status ConvertGraphDefToEngine( auto trt_network = TrtUniquePtrType(builder->createNetwork()); if (!trt_network) { - return tensorflow::errors::Internal( - "Failed to create TensorRT network object"); + return errors::Internal("Failed to create TensorRT network object"); } // Build the network @@ -4018,10 +4038,10 @@ tensorflow::Status ConvertGraphDefToEngine( VLOG(2) << "Converting op name=" << node_name << ", op=" << node_def.op(); if (IsEngineInput(node_name) && (node_def.op() == "Placeholder")) { int32 slot_number = -1; - if (!tensorflow::strings::safe_strto32( // non-absl ok + if (!strings::safe_strto32( // non-absl ok node_name.c_str() + strlen(kInputPHName), &slot_number)) { - return tensorflow::errors::InvalidArgument( - "Failed to parse slot number from ", node_name); + return errors::InvalidArgument("Failed to parse slot number from ", + node_name); } nvinfer1::DataType trt_dtype; nvinfer1::Dims trt_dims; @@ -4046,14 +4066,14 @@ tensorflow::Status ConvertGraphDefToEngine( converter.AddInputTensor(node_name, trt_dtype, trt_dims, batch_size)); } else if (IsEngineOutput(node_name) && (node_def.op() == "Identity")) { int32 slot_number = -1; - if (!tensorflow::strings::safe_strto32( // non-absl ok + if (!strings::safe_strto32( // non-absl ok node_name.c_str() + strlen(kOutputPHName), &slot_number)) { - return tensorflow::errors::InvalidArgument( - "Failed to parse slot number from ", node_name); + return errors::InvalidArgument("Failed to parse slot number from ", + node_name); } // Get output type that TensorFlow expects TFAttrs attrs(node_def); - tensorflow::DataType tf_dtype = attrs.get("T"); + DataType tf_dtype = attrs.get("T"); nvinfer1::DataType trt_dtype; TF_RETURN_IF_ERROR(ConvertDType(tf_dtype, &trt_dtype)); if (output_tensors.size() <= slot_number) { @@ -4077,18 +4097,17 @@ tensorflow::Status ConvertGraphDefToEngine( VLOG(1) << "Starting engine creation"; engine->reset(builder->buildCudaEngine(*converter.network())); if (engine->get() == nullptr) { - return tensorflow::errors::Internal("Failed to build TensorRT engine"); + return errors::Internal("Failed to build TensorRT engine"); } VLOG(1) << "Finished conversion"; - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status ConvertSegmentToGraphDef( - const tensorflow::Graph* graph, - const tensorflow::grappler::GraphProperties& graph_properties, +Status ConvertSegmentToGraphDef( + const Graph* graph, const grappler::GraphProperties& graph_properties, const std::vector& subgraph_nodes, // In topological order - std::vector* connections, - tensorflow::GraphDef* segment_def, string* common_scope) { + std::vector* connections, GraphDef* segment_def, + string* common_scope) { std::set marker_nodes; // Update connection shapes/data types and add corresponding input/output // nodes in the segment graphdef. @@ -4098,12 +4117,12 @@ tensorflow::Status ConvertSegmentToGraphDef( auto outside_node = graph->FindNodeId(connection.outside_id); if (!outside_node) { // This should never happen, unless the original graph is problematic. - return tensorflow::errors::NotFound( - "Cannot find node with id ", connection.outside_id, " in the graph."); + return errors::NotFound("Cannot find node with id ", + connection.outside_id, " in the graph."); } // Updates the shape and data types of input/output connections. - tensorflow::DataType dtype; - tensorflow::PartialTensorShape partial_shape; + DataType dtype; + PartialTensorShape partial_shape; if (connection.is_input_edge) { GetOutputProperties(graph_properties, graph->FindNodeId(connection.outside_id), @@ -4129,7 +4148,7 @@ tensorflow::Status ConvertSegmentToGraphDef( } marker_nodes.insert(node_name); auto seg_node = segment_def->add_node(); - tensorflow::NodeDefBuilder builder(node_name, "Placeholder"); + NodeDefBuilder builder(node_name, "Placeholder"); auto status = builder.Attr("shape", partial_shape) .Attr("dtype", dtype) .Finalize(seg_node); @@ -4148,7 +4167,7 @@ tensorflow::Status ConvertSegmentToGraphDef( } marker_nodes.insert(node_name); auto seg_node = segment_def->add_node(); - tensorflow::NodeDefBuilder builder(node_name, "Identity"); + NodeDefBuilder builder(node_name, "Identity"); auto status = builder .Input(connection.inside_node_name, connection.inside_port, dtype) @@ -4205,7 +4224,7 @@ tensorflow::Status ConvertSegmentToGraphDef( ++input_idx; continue; } else { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Found non control input outside the segment that is not an " "engine connection to ", snode->name(), ": ", input.first); @@ -4224,10 +4243,10 @@ tensorflow::Status ConvertSegmentToGraphDef( *common_scope = local_scope; VLOG(1) << "Converted TensorRT candidate segment @scope '" << local_scope << "' to a GraphDef"; - return tensorflow::Status::OK(); + return Status::OK(); } -bool OutputEdgeValidator::operator()(const tensorflow::Edge* out_edge) const { +bool OutputEdgeValidator::operator()(const Edge* out_edge) const { if (out_edge->IsControlEdge()) return true; if (out_edge->src()->type_string() == "Const") { VLOG(1) << "--> Need to remove output node " << out_edge->src()->name() diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h index 7b37173090519ff6fadd956942d7ea12a0644981..6333d9130a8035e0449446d11a37a63554222752 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h @@ -74,14 +74,14 @@ struct EngineConnection { const string outside_node_name; const int outside_id; const int outside_port; - tensorflow::PartialTensorShape outside_shape; // Only set for input edge. + PartialTensorShape outside_shape; // Only set for input edge. const string inside_node_name; const int inside_id; const int inside_port; - tensorflow::PartialTensorShape inside_shape; // Only set for output edge. + PartialTensorShape inside_shape; // Only set for output edge. - tensorflow::DataType connection_type; + DataType connection_type; const bool is_input_edge; // The port number of the TRT node connected with this edge. @@ -97,7 +97,7 @@ struct EngineInfo { string engine_name; string device; - tensorflow::GraphDef segment_graph_def; + GraphDef segment_graph_def; // Non-control input connections inside this vector are sorted in a way such // that, the segment nodes connecting to them are topological sorted. @@ -125,12 +125,11 @@ struct EngineInfo { // sorted in topological order. // // TODO(aaroey): add tests to validate these properties. -tensorflow::Status ConvertSegmentToGraphDef( - const tensorflow::Graph* graph, - const tensorflow::grappler::GraphProperties& graph_properties, +Status ConvertSegmentToGraphDef( + const Graph* graph, const grappler::GraphProperties& graph_properties, const std::vector& subgraph_nodes, - std::vector* connections, - tensorflow::GraphDef* segment_def, string* common_scope); + std::vector* connections, GraphDef* segment_def, + string* common_scope); // Converts given subgraph to a TRT engine saved in 'engine'. Returns ok iff // 'builder' successfully build the engine. If the result is not ok, 'engine' @@ -140,12 +139,11 @@ tensorflow::Status ConvertSegmentToGraphDef( // - convert_successfully: indicates whether the converson to TensorRT network // is successful. This is different than successfully building the engine: // building can still fail afterwards. -tensorflow::Status ConvertGraphDefToEngine( - const tensorflow::GraphDef& gdef, TrtPrecisionMode precision_mode, - int max_batch_size, size_t max_workspace_size_bytes, - const std::vector& input_shapes, - Logger* logger, nvinfer1::IGpuAllocator* allocator, - TRTInt8Calibrator* calibrator, +Status ConvertGraphDefToEngine( + const GraphDef& gdef, TrtPrecisionMode precision_mode, int max_batch_size, + size_t max_workspace_size_bytes, + const std::vector& input_shapes, Logger* logger, + nvinfer1::IGpuAllocator* allocator, TRTInt8Calibrator* calibrator, TrtUniquePtrType* engine, bool use_calibration, bool* convert_successfully); @@ -155,7 +153,7 @@ class OutputEdgeValidator { public: // Return true if the specified edge is eligible to be an output edge of the // TRT segment. - bool operator()(const tensorflow::Edge* out_edge) const; + bool operator()(const Edge* out_edge) const; }; string DebugString(const nvinfer1::DimensionType type); @@ -203,7 +201,7 @@ class TRT_ShapedWeights { // TODO(aaroey): make these private. nvinfer1::Dims shape_; // Note: shape.type[] is not used. - tensorflow::DataType type_; + DataType type_; private: // This constructor is only used by TrtWeightStore, which creates the @@ -229,8 +227,7 @@ class TRT_ShapedWeights { class TrtWeightStore { public: // Get a TRT_ShapedWeights with 'type' and 'dims'. - TRT_ShapedWeights GetTempWeights(tensorflow::DataType type, - const nvinfer1::Dims& dims); + TRT_ShapedWeights GetTempWeights(DataType type, const nvinfer1::Dims& dims); // Get a TRT_ShapedWeights with the same data type and dimensions as // 'weights'. @@ -341,8 +338,7 @@ class Converter; // Parameters for each op converter. struct OpConverterParams { - OpConverterParams(Converter* arg_converter, - const tensorflow::NodeDef& arg_node_def, + OpConverterParams(Converter* arg_converter, const NodeDef& arg_node_def, const std::vector& arg_inputs, std::vector* arg_outputs, bool arg_validation_only, TrtWeightStore* arg_weight_store) @@ -354,7 +350,7 @@ struct OpConverterParams { weight_store(arg_weight_store) {} Converter* converter; - const tensorflow::NodeDef& node_def; + const NodeDef& node_def; const std::vector& inputs; std::vector* outputs; const bool validation_only; @@ -379,9 +375,12 @@ class TrtNodeValidator { Status ValidateNode( const NodeDef& node_def, const std::vector>& input_node_and_ports, + const TrtPrecisionMode precision_mode, const grappler::GraphProperties& graph_properties); private: + static const std::set* quantize_ops; + void RegisterOpValidators(); // Convert a Const node to a TRT_TensorOrWeights. @@ -434,7 +433,7 @@ class Converter { // function/subgraph. // Convert the node to TRT network. - Status ConvertNode(const tensorflow::NodeDef& node_def); + Status ConvertNode(const NodeDef& node_def); // Add input tensor to the TRT network with given 'name', 'dtype', 'dims' and // 'batch_size'. @@ -487,8 +486,13 @@ class Converter { const nvinfer1::ITensor** output_tensor); // Converts 'input' into 'tensor' with shape specified by 'dims'. + // + // If validation_only is true, it doesn't do the conversion but only do some + // minimum validation for the eligibility of the conversion, and *tensor will + // be set to nullptr. Status PrepareTensorForShape(const TRT_TensorOrWeights& input, const nvinfer1::Dims& dims, + const bool validation_only, const nvinfer1::ITensor** tensor); // Return OK if the broadcast scheme is supported and compute the shapes after @@ -515,7 +519,7 @@ class Converter { Status GetTensorOrWeights(const string& name, TRT_TensorOrWeights* output); // Get the inputs of 'node_def' from trt_tensors_. - Status GetInputs(const tensorflow::NodeDef& node_def, + Status GetInputs(const NodeDef& node_def, std::vector* inputs) const; void RegisterOpConverters(); diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc index bda00e8c4f892ce54478dc86ef2be107936badfe..e89b31759f2e9882b6c0b52539f95677f0773a3c 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc @@ -123,7 +123,7 @@ template NodeDef MakeConstNodeDef(const string& name, const std::vector& vals, const TensorShape& shape) { Scope s = Scope::NewRootScope(); - Tensor t = ::tensorflow::test::AsTensor(vals, shape); + Tensor t = test::AsTensor(vals, shape); auto const_op = ops::Const(s.WithOpName(name), t); return const_op.node()->def(); } @@ -389,8 +389,8 @@ TEST(TRT_TensorOrWeights_Test, Basic) { class ValidatorTest : public ::testing::Test { public: - void AddOpValidator(const string& op_name, OpConverter op_validator) { - validator_.op_validators_[op_name] = op_validator; + std::unordered_map& op_validators() { + return validator_.op_validators_; } Status ConvertToTensorOrWeights( @@ -401,10 +401,18 @@ class ValidatorTest : public ::testing::Test { node_def, output_port, graph_properties, tensor_or_weights); } + const std::set* GetQuantizeOps() { return validator_.quantize_ops; } + protected: TrtNodeValidator validator_; }; +TEST_F(ValidatorTest, QuantizeOpsAreRegistered) { + for (const string& quantize_op : *GetQuantizeOps()) { + QCHECK(op_validators().count(quantize_op)); + } +} + TEST_F(ValidatorTest, ConvertToTensorOrWeights) { // Convert Const. { @@ -477,18 +485,30 @@ TEST_F(ValidatorTest, ValidateNode) { }; NodeDef node_def = MakeNodeDef("my_op", "MyOp", {}); - // Validator not registered, validation should pass. - TF_EXPECT_OK(validator_.ValidateNode(node_def, {}, graph_properties)); + // Validator not registered. + ExpectStatus(validator_.ValidateNode(node_def, {}, TrtPrecisionMode::FP32, + graph_properties), + error::UNIMPLEMENTED, "Op type MyOp is not supported."); // Register validator. - AddOpValidator("MyOp", op_converter); - TF_EXPECT_OK(validator_.ValidateNode(node_def, {}, graph_properties)); + op_validators()["MyOp"] = op_converter; + TF_EXPECT_OK(validator_.ValidateNode(node_def, {}, TrtPrecisionMode::FP32, + graph_properties)); EXPECT_EQ(false, start_conversion); // Let the converter return error. should_fail = true; - ExpectStatus(validator_.ValidateNode(node_def, {}, graph_properties), + ExpectStatus(validator_.ValidateNode(node_def, {}, TrtPrecisionMode::FP32, + graph_properties), error::INVALID_ARGUMENT); + + // Test quantization ops, they're only supported in INT8 mode. The success + // case is tested in OpConverterTest.ConvertQuantize. + node_def = MakeNodeDef("my_op", "FakeQuantWithMinMaxArgs", {}); + ExpectStatus(validator_.ValidateNode(node_def, {}, TrtPrecisionMode::FP32, + graph_properties), + error::UNIMPLEMENTED, + "Op type FakeQuantWithMinMaxArgs is not supported."); } class ConverterTest : public ::testing::Test { @@ -691,23 +711,34 @@ TEST_F(ConverterTest, PrepareTensorForShape_Tensor) { TRT_TensorOrWeights tw(input_tensor); const nvinfer1::ITensor* output_tensor = nullptr; - // Shape size doesn't match. - ExpectStatus(converter_->PrepareTensorForShape(tw, GetTestDims({2, 3, 6}), - &output_tensor), - error::INVALID_ARGUMENT, "Reshape shapes are not compatible"); - - // TODO(aaroey): we should check the case where uninferred dimensions are not - // an exact divisor of input dim ensions, e.g. for dims {-1, 7}. - - // Infer shape, ok. - TF_EXPECT_OK(converter_->PrepareTensorForShape(tw, GetTestDims({-1, 2}), - &output_tensor)); - ExpectTrtDimsEqualsArray({15, 2}, output_tensor->getDimensions()); + for (bool validation_only : {false, true}) { + // Shape size doesn't match. + ExpectStatus( + converter_->PrepareTensorForShape(tw, GetTestDims({2, 3, 6}), + validation_only, &output_tensor), + error::INVALID_ARGUMENT, "Reshape shapes are not compatible"); + + // TODO(aaroey): we should check the case where uninferred dimensions are + // not an exact divisor of input dim ensions, e.g. for dims {-1, 7}. + + // Infer shape, ok. + TF_EXPECT_OK(converter_->PrepareTensorForShape( + tw, GetTestDims({-1, 2}), validation_only, &output_tensor)); + if (validation_only) { + EXPECT_EQ(nullptr, output_tensor); + } else { + ExpectTrtDimsEqualsArray({15, 2}, output_tensor->getDimensions()); + } - // Regular shape. - TF_EXPECT_OK(converter_->PrepareTensorForShape(tw, GetTestDims({10, 3}), - &output_tensor)); - ExpectTrtDimsEqualsArray({10, 3}, output_tensor->getDimensions()); + // Regular shape. + TF_EXPECT_OK(converter_->PrepareTensorForShape( + tw, GetTestDims({10, 3}), validation_only, &output_tensor)); + if (validation_only) { + EXPECT_EQ(nullptr, output_tensor); + } else { + ExpectTrtDimsEqualsArray({10, 3}, output_tensor->getDimensions()); + } + } } TEST_F(ConverterTest, PrepareTensorForShape_Weights) { @@ -715,9 +746,15 @@ TEST_F(ConverterTest, PrepareTensorForShape_Weights) { weight_store_->GetTempWeights(DT_FLOAT, GetTestDims({2, 3, 5})); TRT_TensorOrWeights tw(weights); const nvinfer1::ITensor* output_tensor = nullptr; - TF_EXPECT_OK(converter_->PrepareTensorForShape(tw, GetTestDims({10, 3}), - &output_tensor)); - ExpectTrtDimsEqualsArray({10, 3}, output_tensor->getDimensions()); + for (bool validation_only : {false, true}) { + TF_EXPECT_OK(converter_->PrepareTensorForShape( + tw, GetTestDims({10, 3}), validation_only, &output_tensor)); + if (validation_only) { + EXPECT_EQ(nullptr, output_tensor); + } else { + ExpectTrtDimsEqualsArray({10, 3}, output_tensor->getDimensions()); + } + } } TEST_F(ConverterTest, MaybeUpdateBatchSize) { @@ -948,7 +985,7 @@ class ConvertGraphDefToEngineTest : public ::testing::Test { Status RunConvertGraphDefToEngine(Scope* s) { GraphDef gdef; TF_EXPECT_OK(s->ToGraphDef(&gdef)); - std::vector input_shapes; + std::vector input_shapes; int batch_size = -1; for (const NodeDef& node : gdef.node()) { absl::string_view node_name(node.name()); @@ -1049,7 +1086,7 @@ class OpConverterTest : public ::testing::Test { // Reset the validator and converter. validator_.reset(new TrtNodeValidator); - converter_.reset(new Converter(network_.get(), TrtPrecisionMode::FP32, + converter_.reset(new Converter(network_.get(), precision_mode_to_test_, /*use_calibration=*/false)); // Reset other related artifacts. @@ -1182,9 +1219,10 @@ class OpConverterTest : public ::testing::Test { grappler::GraphProperties graph_properties(item); TF_EXPECT_OK(graph_properties.InferStatically(true)); - ExpectStatus(validator_->ValidateNode(node_def, input_node_and_ports, - graph_properties), - expected_code, expected_msg_substr); + ExpectStatus( + validator_->ValidateNode(node_def, input_node_and_ports, + precision_mode_to_test_, graph_properties), + expected_code, expected_msg_substr); } void RunConversion(const NodeDef& node_def, @@ -1214,6 +1252,10 @@ class OpConverterTest : public ::testing::Test { std::unique_ptr converter_; std::unique_ptr validator_; + protected: + // TODO(laigd): parameterize the test and make the precision mode a parameter. + TrtPrecisionMode precision_mode_to_test_ = TrtPrecisionMode::FP32; + private: Logger logger_; TrtUniquePtrType builder_; @@ -1298,34 +1340,34 @@ void TestConvertConst(OpConverterTest* test) { reset_and_test(t, false, {}, {}); } { - Tensor t = ::tensorflow::test::AsScalar(12); + Tensor t = test::AsScalar(12); reset_and_test(t, false, {1}, {12}); reset_and_test(t, true, {1}, {12}); } { - Tensor t = ::tensorflow::test::AsTensor({1, 2}); + Tensor t = test::AsTensor({1, 2}); reset_and_test(t, false, {2}, {1, 2}); reset_and_test(t, true, {2}, {1, 2}); } { - Tensor t = ::tensorflow::test::AsTensor({1, 2, 3, 4, 5, 6}, - TensorShape({2, 3})); + Tensor t = + test::AsTensor({1, 2, 3, 4, 5, 6}, TensorShape({2, 3})); reset_and_test(t, false, {2, 3}, {1, 2, 3, 4, 5, 6}); reset_and_test(t, true, {2, 3}, {1, 2, 3, 4, 5, 6}); } { // Set all tensor elements to the same value. Such tensors are encoded // using a single element list in tensor proto. - Tensor t = ::tensorflow::test::AsTensor({1, 1, 1, 1, 1, 1}, - TensorShape({2, 3})); + Tensor t = + test::AsTensor({1, 1, 1, 1, 1, 1}, TensorShape({2, 3})); reset_and_test(t, false, {2, 3}, {1, 1, 1, 1, 1, 1}); reset_and_test(t, true, {2, 3}, {1, 1, 1, 1, 1, 1}); } { // Set trailing tensor elements to the same value. Such tensors are // encoded by truncating all equal elements except the first one. - Tensor t = ::tensorflow::test::AsTensor({2, 2, 1, 1, 1, 1}, - TensorShape({2, 3})); + Tensor t = + test::AsTensor({2, 2, 1, 1, 1, 1}, TensorShape({2, 3})); reset_and_test(t, false, {2, 3}, {2, 2, 1, 1, 1, 1}); reset_and_test(t, true, {2, 3}, {2, 2, 1, 1, 1, 1}); } @@ -2006,6 +2048,7 @@ TEST_F(OpConverterTest, ConvertBinary) { } TEST_F(OpConverterTest, ConvertQuantize) { + precision_mode_to_test_ = TrtPrecisionMode::INT8; const std::pair op_with_num_inputs[4] = { {"FakeQuantWithMinMaxArgs", 1}, {"FakeQuantWithMinMaxVars", 3}, @@ -2133,48 +2176,6 @@ TEST_F(OpConverterTest, ConvertQuantize) { } } -TEST_F(OpConverterTest, ConvertRelu6) { - { - // Input list is empty, should fail. - NodeDef node_def = MakeNodeDef("my_relu6", "Relu6", {}); - RunValidationAndConversion( - node_def, error::INVALID_ARGUMENT, - "Relu6 got 0 inputs but expected 1, at my_relu6"); - } - - // Get the NodeDef for Relu6. - Scope s = Scope::NewRootScope(); - auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); - auto relu6 = ops::Relu6(s.WithOpName("my_relu6"), input); - const NodeDef node_def = relu6.operation.node()->def(); - { - // Input is weights, should fail. - Reset(); - AddTestWeights("input", {1}, {1.0f}); - RunValidationAndConversion( - node_def, error::UNIMPLEMENTED, - "The input \"input\" for Relu6 must be a tensor, at my_relu6"); - } - { - // Clip tensor values and set quantization ranges, ok. - Reset(); - AddTestTensor("input", {1, 2, 3}); - RunValidationAndConversion(node_def); - TRT_TensorOrWeights output; - TF_EXPECT_OK(GetTensorOrWeights("my_relu6", &output)); - EXPECT_TRUE(output.is_tensor()); - auto ranges = quantization_ranges(); - EXPECT_EQ(ranges[output.tensor()], 6.0f); - - const DataVec input_data{ - {"input", test::AsTensor({-100, -1, 0, 3, 5, 9})}}; - DataVec output_data{{"my_relu6", ConstructTensor(6)}}; - BuildAndRun(input_data, &output_data); - EXPECT_THAT(GetSpanForData(output_data[0]), - ElementsAre(0, 0, 0, 3, 5, 6)); - } -} - template void TestConvertSquare(OpConverterTest* test) { test->Reset(); @@ -2255,13 +2256,23 @@ TEST_F(OpConverterTest, ConvertActivation) { "The input \"input\" for Relu must be a tensor, at my_act"); } + constexpr float kAlpha = 0.2f; + // Get nodedef for activation layer. auto get_act_nodedef = [](string op_name) -> NodeDef { Scope s = Scope::NewRootScope(); auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); - if (op_name == "Relu") { + if (op_name == "LeakyRelu") { + // LeakyRelu does not have a C++ API + NodeDef node_def = MakeNodeDef("my_act", "LeakyRelu", {"input"}); + (*node_def.mutable_attr())["alpha"].set_f(kAlpha); + return node_def; + } else if (op_name == "Relu") { auto act = ops::Relu(s.WithOpName("my_act"), input); return act.operation.node()->def(); + } else if (op_name == "Relu6") { + auto act = ops::Relu6(s.WithOpName("my_act"), input); + return act.operation.node()->def(); } else if (op_name == "Sigmoid") { auto act = ops::Sigmoid(s.WithOpName("my_act"), input); return act.operation.node()->def(); @@ -2274,8 +2285,12 @@ TEST_F(OpConverterTest, ConvertActivation) { }; // Get expected output for activation layer. auto get_act_output = [](string op_name, float input) -> float { - if (op_name == "Relu") { + if (op_name == "LeakyRelu") { + return (input > 0.0f) ? input : input * kAlpha; + } else if (op_name == "Relu") { return (input > 0.0f) ? input : 0.0f; + } else if (op_name == "Relu6") { + return std::min(std::max(input, 0.0f), 6.0f); } else if (op_name == "Sigmoid") { return 1.0f / (1.0f + std::exp(-input)); } else if (op_name == "Tanh") { @@ -2286,7 +2301,8 @@ TEST_F(OpConverterTest, ConvertActivation) { }; // Ok. - for (string op_name : {"Relu", "Sigmoid", "Tanh"}) { + for (const string& op_name : + {"LeakyRelu", "Relu", "Relu6", "Sigmoid", "Tanh"}) { Reset(); NodeDef node_def = get_act_nodedef(op_name); AddTestTensor("input", {1, 2, 3}); @@ -2295,6 +2311,11 @@ TEST_F(OpConverterTest, ConvertActivation) { TF_EXPECT_OK(GetTensorOrWeights("my_act", &output)); EXPECT_TRUE(output.is_tensor()); ExpectTrtDimsEqualsArray({1, 2, 3}, output.tensor()->getDimensions()); + if (op_name == "Relu6") { + // Relu6 should set quantization range automatically. + auto ranges = quantization_ranges(); + EXPECT_EQ(ranges[output.tensor()], 6.0f); + } const std::vector input = {-100, -2, -1, 0, 1, 100}; const DataVec input_data{{"input", test::AsTensor(input)}}; @@ -2549,8 +2570,8 @@ TEST_F(OpConverterTest, ConvertStridedSlice) { // Get nodedef for StridedSlice layer. auto get_strided_slice_nodedef = - [](int begin_mask = 0, int end_mask = 0, int ellipsis_mask = 0, - int new_axis_mask = 0, int shrink_axis_mask = 0) -> NodeDef { + [](int64 begin_mask = 0, int64 end_mask = 0, int64 ellipsis_mask = 0, + int64 new_axis_mask = 0, int64 shrink_axis_mask = 0) -> NodeDef { Scope s = Scope::NewRootScope(); auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); auto begin = ops::Placeholder(s.WithOpName("begin"), DT_INT32); diff --git a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc index 0eedfcacb4c11c8dc63fcfc13f044586b99b3c76..a4b2fad15593454823ac52060db8e6127b26a0a6 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc @@ -34,13 +34,13 @@ namespace convert { // TODO(sami): Remove VLOG messages once the code matures using absl::StrAppend; using absl::StrCat; -using tensorflow::str_util::Uppercase; +using str_util::Uppercase; -tensorflow::Status TRTOptimizationPass::Init( - const tensorflow::RewriterConfig_CustomGraphOptimizer* config) { +Status TRTOptimizationPass::Init( + const RewriterConfig_CustomGraphOptimizer* config) { VLOG(1) << "Called INIT for " << name_ << " with config = " << config; if (config == nullptr) { - return tensorflow::Status::OK(); + return Status::OK(); } const auto params = config->parameter_map(); if (params.count("minimum_segment_size")) { @@ -72,12 +72,11 @@ tensorflow::Status TRTOptimizationPass::Init( if (params.count("use_calibration")) { use_calibration_ = params.at("use_calibration").b(); } - return tensorflow::Status::OK(); + return Status::OK(); } -void TRTOptimizationPass::PrintDebugInfo( - tensorflow::grappler::Cluster* cluster, - const tensorflow::grappler::GrapplerItem& item) { +void TRTOptimizationPass::PrintDebugInfo(grappler::Cluster* cluster, + const grappler::GrapplerItem& item) { LOG(INFO) << "Cluster = " << cluster; string offset(" "); string offset2 = StrCat(offset, offset); @@ -95,7 +94,7 @@ void TRTOptimizationPass::PrintDebugInfo( } std::unordered_map peak_mem; auto status = cluster->GetPeakMemoryUsage(&peak_mem); - if (status == tensorflow::Status::OK()) { + if (status == Status::OK()) { LOG(INFO) << offset << "Peak Memory Usage :"; for (auto s : peak_mem) { LOG(INFO) << offset2 << s.first << " = " << s.second; @@ -177,9 +176,9 @@ void TRTOptimizationPass::PrintDebugInfo( } } -tensorflow::Status TRTOptimizationPass::Optimize( - tensorflow::grappler::Cluster* cluster, - const tensorflow::grappler::GrapplerItem& item, GraphDef* optimized_graph) { +Status TRTOptimizationPass::Optimize(grappler::Cluster* cluster, + const grappler::GrapplerItem& item, + GraphDef* optimized_graph) { VLOG(1) << "Called TRTOptimization Pass " << name_; // This is a hack to workaround optimizer issue. MetaOptimizer calls // optimization passes on function objects as well, we should not modify @@ -190,7 +189,7 @@ tensorflow::Status TRTOptimizationPass::Optimize( << " is probably called on funcdef! This optimizer must *NOT* " "be called on function objects."; *optimized_graph = item.graph; - return tensorflow::Status::OK(); + return Status::OK(); } if (VLOG_IS_ON(3)) { LOG(INFO) << CurrentStackTrace(); @@ -223,9 +222,9 @@ tensorflow::Status TRTOptimizationPass::Optimize( << " adjusting maximum batch size to match input batch size"; } } - tensorflow::grappler::GraphProperties static_graph_properties(item); + grappler::GraphProperties static_graph_properties(item); TF_RETURN_IF_ERROR(static_graph_properties.InferStatically(true)); - tensorflow::tensorrt::convert::ConversionParams cp; + ConversionParams cp; if (use_calibration_ && precision_mode_ != TrtPrecisionMode::INT8) { VLOG(1) << "Calibration with FP32 or FP16 is not implemented. " @@ -263,27 +262,23 @@ tensorflow::Status TRTOptimizationPass::Optimize( cp.cached_engine_batches = batches_; cp.max_cached_engines = max_cached_batches_; cp.use_calibration = use_calibration_; - auto status = tensorflow::tensorrt::convert::ConvertAfterShapes(cp); + auto status = ConvertAfterShapes(cp); VLOG(1) << "Returning from " << name_; return status; } -void TRTOptimizationPass::Feedback( - tensorflow::grappler::Cluster* cluster, - const tensorflow::grappler::GrapplerItem& item, - const GraphDef& optimized_graph, double result) {} - -} // namespace convert -} // namespace tensorrt -} // namespace tensorflow +void TRTOptimizationPass::Feedback(grappler::Cluster* cluster, + const grappler::GrapplerItem& item, + const GraphDef& optimized_graph, + double result) {} class VerboseCustomGraphOptimizerRegistrar - : public tensorflow::grappler::CustomGraphOptimizerRegistrar { + : public grappler::CustomGraphOptimizerRegistrar { public: VerboseCustomGraphOptimizerRegistrar( - const tensorflow::grappler::CustomGraphOptimizerRegistry::Creator& cr, - const tensorflow::string& name) - : tensorflow::grappler::CustomGraphOptimizerRegistrar(cr, name) { + const grappler::CustomGraphOptimizerRegistry::Creator& cr, + const string& name) + : grappler::CustomGraphOptimizerRegistrar(cr, name) { VLOG(1) << "Constructing a CustomOptimizationPass registration object for " << name; } @@ -293,10 +288,13 @@ static VerboseCustomGraphOptimizerRegistrar TRTOptimizationPass_Registrar( []() { VLOG(1) << "Instantiating CustomOptimizationPass object TensorRTOptimizer"; - return new tensorflow::tensorrt::convert::TRTOptimizationPass( - "TensorRTOptimizer"); + return new TRTOptimizationPass("TensorRTOptimizer"); }, ("TensorRTOptimizer")); +} // namespace convert +} // namespace tensorrt +} // namespace tensorflow + #endif #endif diff --git a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h index b2aed2a37afb6c01863f5617bad0bafe004eec24..a30516710055f7c12ee1a9211e56cbf174a44246 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h +++ b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h @@ -30,7 +30,7 @@ namespace tensorflow { namespace tensorrt { namespace convert { -class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer { +class TRTOptimizationPass : public grappler::CustomGraphOptimizer { public: TRTOptimizationPass(const string& name = "TRTOptimizationPass") : name_(name), @@ -46,19 +46,18 @@ class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer { string name() const override { return name_; }; - tensorflow::Status Init(const tensorflow::RewriterConfig_CustomGraphOptimizer* - config = nullptr) override; + Status Init( + const RewriterConfig_CustomGraphOptimizer* config = nullptr) override; - tensorflow::Status Optimize(tensorflow::grappler::Cluster* cluster, - const tensorflow::grappler::GrapplerItem& item, - GraphDef* optimized_graph) override; + Status Optimize(grappler::Cluster* cluster, + const grappler::GrapplerItem& item, + GraphDef* optimized_graph) override; - void Feedback(tensorflow::grappler::Cluster* cluster, - const tensorflow::grappler::GrapplerItem& item, + void Feedback(grappler::Cluster* cluster, const grappler::GrapplerItem& item, const GraphDef& optimized_graph, double result) override; - void PrintDebugInfo(tensorflow::grappler::Cluster* cluster, - const tensorflow::grappler::GrapplerItem& item); + void PrintDebugInfo(grappler::Cluster* cluster, + const grappler::GrapplerItem& item); private: const string name_; diff --git a/tensorflow/compiler/tf2tensorrt/kernels/get_serialized_resource_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/get_serialized_resource_op.cc index 81406b6e301ca350a3e52c97f5fcb575e88c3a90..e252f9111d61dce0b0821f72b3c56f2516fc20f3 100644 --- a/tensorflow/compiler/tf2tensorrt/kernels/get_serialized_resource_op.cc +++ b/tensorflow/compiler/tf2tensorrt/kernels/get_serialized_resource_op.cc @@ -46,7 +46,7 @@ class GetSerializedResourceOp : public OpKernel { SerializableResourceBase* resource = nullptr; OP_REQUIRES_OK(context, context->resource_manager()->Lookup( container, resource_name, &resource)); - ::tensorflow::core::ScopedUnref sc(resource); + core::ScopedUnref sc(resource); // Serialize the resource as output. string serialized_resource; diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc index f6d387c59cd04aa5c7ccad610290b7b1f1d2b11f..d949406344388a62b42f8ba3dec061f2578a02d1 100644 --- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc +++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc @@ -53,7 +53,7 @@ using ::nvinfer1::IRuntime; // A helper class to call done() when destructed for asynchronous execution. // Helps simultaneous execution of native and TRT engines. -class AsyncHelper : public tensorflow::core::RefCounted { +class AsyncHelper : public core::RefCounted { public: AsyncHelper(AsyncOpKernel::DoneCallback done) { done_ = done; } ~AsyncHelper() override { done_(); } @@ -139,37 +139,36 @@ class TRTEngineOp : public AsyncOpKernel { bool use_calibration_; }; -#define TYPECASE(dt, X, Y) \ - case dt: { \ - return (void*)X->flat::Type>().data(); \ +#define TYPECASE(dt, X, Y) \ + case dt: { \ + return (void*)X->flat::Type>().data(); \ } void* GetTensorAddress(const Tensor* tensor_ptr) { auto tensor_type = tensor_ptr->dtype(); switch (tensor_type) { - TYPECASE(tensorflow::DT_FLOAT, tensor_ptr, dest_ptr); - TYPECASE(tensorflow::DT_HALF, tensor_ptr, dest_ptr); - TYPECASE(tensorflow::DT_INT8, tensor_ptr, dest_ptr); + TYPECASE(DT_FLOAT, tensor_ptr, dest_ptr); + TYPECASE(DT_HALF, tensor_ptr, dest_ptr); + TYPECASE(DT_INT8, tensor_ptr, dest_ptr); default: { - LOG(ERROR) << "Unsupported Data type " - << tensorflow::DataTypeString(tensor_type); + LOG(ERROR) << "Unsupported Data type " << DataTypeString(tensor_type); return nullptr; } } } -tensorflow::Status TRTEngineOp::ConstructFunctionHandle(OpKernelContext* ctx) { +Status TRTEngineOp::ConstructFunctionHandle(OpKernelContext* ctx) { VLOG(1) << "Constructing function handle"; auto lib = ctx->function_library(); if (lib == nullptr) { - return tensorflow::errors::Internal("Context function library is null"); + return errors::Internal("Context function library is null"); } auto fdef = lib->GetFunctionLibraryDefinition()->Find(funcdef_name_); if (fdef == nullptr) { - return tensorflow::errors::Internal("Native FunctionDef ", funcdef_name_, - " can't be found in function library"); + return errors::Internal("Native FunctionDef ", funcdef_name_, + " can't be found in function library"); } - tensorflow::FunctionLibraryRuntime::InstantiateOptions inst_ops; + FunctionLibraryRuntime::InstantiateOptions inst_ops; inst_ops.overlay_lib = nullptr; inst_ops.state_handle = ""; inst_ops.target = ctx->device()->name(); @@ -194,11 +193,15 @@ TRTEngineOp::TRTEngineOp(OpKernelConstruction* context) if (!static_engine_) { if (!segment_graph_.ParseFromString(serialized_segment_)) { LOG(ERROR) << "Parsing segment graph failed!"; - context->SetStatus(tensorflow::errors::InvalidArgument( - "Failed to parse segment graphdef!")); + context->SetStatus( + errors::InvalidArgument("Failed to parse segment graphdef!")); return; } - serialized_segment_.resize(0); + VLOG(1) << "Size of serialized GraphDef: " + << serialized_segment_.capacity(); + string tmp; + // Swap with temporary empty string to deallocate the CPU memory. + serialized_segment_.swap(tmp); } VLOG(1) << "Constructing " << name(); string precision_string; @@ -220,7 +223,7 @@ TRTEngineOp::TRTEngineOp(OpKernelConstruction* context) calibrator_.reset(new TRTInt8Calibrator(calibration_data)); calibration_data.resize(0); } - native_func_ = tensorflow::kInvalidHandle; + native_func_ = kInvalidHandle; OP_REQUIRES_OK(context, context->GetAttr("max_cached_engines_count", &max_cached_engines_)); OP_REQUIRES_OK(context, context->GetAttr("cached_engine_batches", @@ -239,7 +242,7 @@ void TRTEngineOp::ExecuteNativeSegment(OpKernelContext* ctx, AsyncHelper* helper) { std::vector inputs; std::vector* outputs = new std::vector(); - if (native_func_ == tensorflow::kInvalidHandle) { + if (native_func_ == kInvalidHandle) { auto status = ConstructFunctionHandle(ctx); if (!status.ok()) { LOG(ERROR) << "Couldn't construct function handle " << funcdef_name_; @@ -248,7 +251,7 @@ void TRTEngineOp::ExecuteNativeSegment(OpKernelContext* ctx, } } auto lib = ctx->function_library(); - tensorflow::FunctionLibraryRuntime::Options opts; + FunctionLibraryRuntime::Options opts; opts.step_id = ctx->step_id(); opts.rendezvous = ctx->rendezvous(); opts.cancellation_manager = ctx->cancellation_manager(); @@ -260,8 +263,8 @@ void TRTEngineOp::ExecuteNativeSegment(OpKernelContext* ctx, helper->Ref(); // Increment count for calculating native graph VLOG(1) << "Executing native segment: " << name(); lib->Run(opts, native_func_, inputs, outputs, - [this, ctx, outputs, helper](const tensorflow::Status& s) { - tensorflow::core::ScopedUnref sc(helper); + [this, ctx, outputs, helper](const Status& s) { + core::ScopedUnref sc(helper); if (!s.ok()) { LOG(ERROR) << "Failed to execute native segment " << this->name() << ": " << s; @@ -282,18 +285,17 @@ void TRTEngineOp::ExecuteCalibration(OpKernelContext* ctx, AsyncHelper* helper) { VLOG(1) << "Executing TRT calibration: " << name(); helper->Ref(); - tensorflow::core::ScopedUnref sc(helper); + core::ScopedUnref sc(helper); auto res_mgr = ctx->resource_manager(); TRTCalibrationResource* calib_res = nullptr; - OP_REQUIRES_OK( - ctx, - res_mgr->LookupOrCreate( - "TF_TRT_Calibration", name(), - reinterpret_cast(&calib_res), - {[ctx, this](SerializableResourceBase** cr) -> tensorflow::Status { - return this->AllocateCalibrationResources(ctx, cr); - }})); - tensorflow::core::ScopedUnref calib_sc(calib_res); + OP_REQUIRES_OK(ctx, + res_mgr->LookupOrCreate( + "TF_TRT_Calibration", name(), + reinterpret_cast(&calib_res), + {[ctx, this](SerializableResourceBase** cr) -> Status { + return this->AllocateCalibrationResources(ctx, cr); + }})); + core::ScopedUnref calib_sc(calib_res); int num_inputs = ctx->num_inputs(); // Pass input data to calibrator std::unordered_map input_data; @@ -301,7 +303,7 @@ void TRTEngineOp::ExecuteCalibration(OpKernelContext* ctx, const Tensor& t = ctx->input(i); void* data_address = GetTensorAddress(&t); if (data_address == nullptr) { - ctx->SetStatus(tensorflow::errors::InvalidArgument( + ctx->SetStatus(errors::InvalidArgument( "Unsupported data type encountered in input ", i)); return; } @@ -357,13 +359,13 @@ bool TRTEngineOp::GetCompatibleCachedEngine( void TRTEngineOp::ComputeAsync(OpKernelContext* ctx, AsyncOpKernel::DoneCallback done) { auto helper = new AsyncHelper(done); - tensorflow::core::ScopedUnref sc(helper); + core::ScopedUnref sc(helper); if (calibration_mode_) { ExecuteCalibration(ctx, helper); return; } // Get shapes of inputs to engine. - std::vector input_shapes; + std::vector input_shapes; input_shapes.reserve(ctx->num_inputs()); for (int i = 0; i < ctx->num_inputs(); ++i) { input_shapes.push_back(ctx->input(i).shape()); @@ -492,7 +494,7 @@ bool TRTEngineOp::ExecuteTrtEngine(OpKernelContext* ctx, // nvinfer1::IExecutionContext::enqueue is not thread safe and we need a mutex // for it. - tensorflow::mutex_lock lock(engine_context->mu); + mutex_lock lock(engine_context->mu); // TODO(jie): trt enqueue does not return error auto ret = engine_context->execution_context->enqueue(num_batch, &buffers[0], *stream, nullptr); @@ -508,7 +510,7 @@ bool TRTEngineOp::ExecuteTrtEngine(OpKernelContext* ctx, EngineContext* TRTEngineOp::GetEngine( const std::vector& input_shapes, OpKernelContext* ctx) { static EngineContext empty_context; - tensorflow::mutex_lock lock(engine_mutex_); + mutex_lock lock(engine_mutex_); // TODO(tmorris): using first input to get batch size - is this reliable? const int batch_size = input_shapes[0].dim_size(0); @@ -516,7 +518,7 @@ EngineContext* TRTEngineOp::GetEngine( TRTEngineCacheResource* cache_res = nullptr; auto status = ctx->resource_manager()->LookupOrCreate( "TRTEngineCache", funcdef_name_, &cache_res, - {[this, ctx](TRTEngineCacheResource** cr) -> tensorflow::Status { + {[this, ctx](TRTEngineCacheResource** cr) -> Status { *cr = new TRTEngineCacheResource(ctx, this->max_cached_engines_); return Status::OK(); }}); @@ -524,7 +526,7 @@ EngineContext* TRTEngineOp::GetEngine( ctx->SetStatus(status); return &empty_context; } - tensorflow::core::ScopedUnref sc(cache_res); + core::ScopedUnref sc(cache_res); auto& cache = cache_res->cache_; auto allocator = cache_res->allocator_.get(); if (allocator == nullptr) { @@ -566,7 +568,11 @@ EngineContext* TRTEngineOp::GetEngine( TrtUniquePtrType( raw_static_engine->createExecutionContext()))); // Runtime is safe to delete after engine creation - serialized_segment_.clear(); + VLOG(1) << "Size of serialized TRT engine: " + << serialized_segment_.capacity(); + string tmp; + // Swap with temporary empty string to deallocate the CPU memory. + serialized_segment_.swap(tmp); if (max_batch_size < batch_size) { return &empty_context; } @@ -576,7 +582,7 @@ EngineContext* TRTEngineOp::GetEngine( // Handle the dynamic engine case. // See if there is a compatible engine cached. The batch size should be <= the // cached batch size. - std::vector engine_input_shapes; + std::vector engine_input_shapes; const bool matched_successfully = GetCompatibleCachedEngine(input_shapes, &engine_input_shapes); // If matched, use that engine. Otherwise, we will look in cache for that @@ -630,12 +636,12 @@ EngineContext* TRTEngineOp::GetEngine( return cache.at(engine_input_shapes).get(); } -tensorflow::Status TRTEngineOp::AllocateCalibrationResources( +Status TRTEngineOp::AllocateCalibrationResources( OpKernelContext* ctx, SerializableResourceBase** cr) { auto cres = new TRTCalibrationResource(); *cr = cres; // Get the allocator. - auto alloc = ctx->device()->GetAllocator(tensorflow::AllocatorAttributes()); + auto alloc = ctx->device()->GetAllocator(AllocatorAttributes()); if (!alloc) { LOG(WARNING) << "Can't get device allocator will not be able to " "allocate memory from TensorFlow memory pool"; @@ -646,12 +652,12 @@ tensorflow::Status TRTEngineOp::AllocateCalibrationResources( // Get the input shapes. const int batch_size = ctx->input(0).dim_size(0); const int num_inputs = ctx->num_inputs(); - std::vector shapes; + std::vector shapes; cres->device_tensors_.resize(num_inputs); VLOG(1) << " Constructing calibrator"; for (int i = 0; i < num_inputs; i++) { // allocate workspace on device for inputs - const tensorflow::Tensor& t = ctx->input(i); + const Tensor& t = ctx->input(i); shapes.emplace_back(t.shape()); Tensor* device_tensor; TF_RETURN_IF_ERROR(ctx->allocate_persistent( @@ -659,7 +665,7 @@ tensorflow::Status TRTEngineOp::AllocateCalibrationResources( CHECK_EQ(t.TotalBytes(), device_tensor->TotalBytes()); void* device_address = GetTensorAddress(device_tensor); if (device_address == nullptr) { - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Unsupported data type encountered in input ", i); } cres->device_buffers_.emplace( @@ -674,7 +680,7 @@ tensorflow::Status TRTEngineOp::AllocateCalibrationResources( ctx->device()->tensorflow_gpu_device_info()->gpu_id; if (platform_gpu_id < 0) { LOG(ERROR) << "Can't get gpu_device_info from context->device()"; - return tensorflow::errors::InvalidArgument( + return errors::InvalidArgument( "Context->device doesn't contain device info!"); } const int64 workspace_size_bytes = workspace_size_; @@ -709,7 +715,7 @@ tensorflow::Status TRTEngineOp::AllocateCalibrationResources( VLOG(1) << "Calibration loop terminated " << label; })); VLOG(1) << "initialized calibrator resource"; - return tensorflow::Status::OK(); + return Status::OK(); } REGISTER_KERNEL_BUILDER(Name("TRTEngineOp").Device(DEVICE_GPU), TRTEngineOp); diff --git a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.cc b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.cc index 871fb1210bd495dc3f5e8153bb6c3a361bf569f5..dd73d15029d6fe5515c823223ffe743e52dde6e9 100644 --- a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.cc +++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.cc @@ -33,7 +33,7 @@ PluginTensorRT* PluginFactoryTensorRT::createPlugin(const char* layer_name, return nullptr; } - tensorflow::mutex_lock lock(instance_m_); + mutex_lock lock(instance_m_); auto plugin_ptr = plugin_registry_[encoded_op_name].first(serial_data, serial_length); owned_plugins_.emplace_back(plugin_ptr); @@ -44,7 +44,7 @@ PluginTensorRT* PluginFactoryTensorRT::createPlugin(const char* layer_name, PluginTensorRT* PluginFactoryTensorRT::CreatePlugin(const string& op_name) { if (!IsPlugin(op_name)) return nullptr; - tensorflow::mutex_lock lock(instance_m_); + mutex_lock lock(instance_m_); auto plugin_ptr = plugin_registry_[op_name].second(); owned_plugins_.emplace_back(plugin_ptr); @@ -56,7 +56,7 @@ bool PluginFactoryTensorRT::RegisterPlugin( PluginConstructFunc construct_func) { if (IsPlugin(op_name)) return false; - tensorflow::mutex_lock lock(instance_m_); + mutex_lock lock(instance_m_); auto ret = plugin_registry_.emplace( op_name, std::make_pair(deserialize_func, construct_func)); @@ -64,7 +64,7 @@ bool PluginFactoryTensorRT::RegisterPlugin( } void PluginFactoryTensorRT::DestroyPlugins() { - tensorflow::mutex_lock lock(instance_m_); + mutex_lock lock(instance_m_); owned_plugins_.clear(); } diff --git a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h index 9aa99a40b80de92a4d9b9ad36e88e693b8aa42dc..cce4f52d9f1080fe0174b5fcb5dd0afdaf6e7769 100644 --- a/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h +++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h @@ -69,7 +69,7 @@ class PluginFactoryTensorRT : public nvinfer1::IPluginFactory { // TODO(jie): Owned plugin should be associated with different sessions; // should really hand ownership of plugins to resource management; std::vector> owned_plugins_; - tensorflow::mutex instance_m_; + mutex instance_m_; }; class TrtPluginRegistrar { @@ -89,9 +89,8 @@ class TrtPluginRegistrar { construct_func) \ REGISTER_TRT_PLUGIN_UNIQ(ctr, name, deserialize_func, construct_func) #define REGISTER_TRT_PLUGIN_UNIQ(ctr, name, deserialize_func, construct_func) \ - static ::tensorflow::tensorrt::TrtPluginRegistrar trt_plugin_registrar##ctr \ - TF_ATTRIBUTE_UNUSED = ::tensorflow::tensorrt::TrtPluginRegistrar( \ - name, deserialize_func, construct_func) + static TrtPluginRegistrar trt_plugin_registrar##ctr TF_ATTRIBUTE_UNUSED = \ + TrtPluginRegistrar(name, deserialize_func, construct_func) } // namespace tensorrt } // namespace tensorflow diff --git a/tensorflow/compiler/tf2tensorrt/segment/segment.cc b/tensorflow/compiler/tf2tensorrt/segment/segment.cc index 3794929b1df3fa999de6ab218dc2ddfb96e4ac81..9cab9d701293f2fd4704fea56ae5caf09c8f4d2f 100644 --- a/tensorflow/compiler/tf2tensorrt/segment/segment.cc +++ b/tensorflow/compiler/tf2tensorrt/segment/segment.cc @@ -39,7 +39,7 @@ namespace segment { using absl::StrAppend; using absl::StrCat; -// A simple graph representation to mirror tensorflow::Graph. This structure +// A simple graph representation to mirror Graph. This structure // helps saving memory since segmenter modifies the graph in place, preventing // the need to create a copy of the graph. It is composed of edges and nodes. // Nodes keep pointers to original TF nodes. @@ -75,7 +75,7 @@ class SimpleEdge { class SimpleNode { public: - SimpleNode(const tensorflow::Node* node, const int id); + SimpleNode(const Node* node, const int id); const std::vector& in_edges() const { return in_edges_; } const std::vector& out_edges() const { return out_edges_; } @@ -99,11 +99,11 @@ class SimpleNode { } const string& name() const { return node_->name(); } - const tensorflow::Node* tf_node() const { return node_; } + const Node* tf_node() const { return node_; } int id() const { return id_; } private: - const tensorflow::Node* node_; + const Node* node_; std::vector in_edges_; std::vector out_edges_; int id_; @@ -113,7 +113,7 @@ class SimpleNode { class SimpleGraph { public: - explicit SimpleGraph(const tensorflow::Graph* g); + explicit SimpleGraph(const Graph* g); ~SimpleGraph(); void AddControlEdge(SimpleNode* src, SimpleNode* dst); @@ -126,15 +126,11 @@ class SimpleGraph { return nodes_[node_id]; } int num_node_ids() const { return nodes_.size(); } - const SimpleNode* source_node() const { - return nodes_[tensorflow::Graph::kSourceId]; - } - const SimpleNode* sink_node() const { - return nodes_[tensorflow::Graph::kSinkId]; - } + const SimpleNode* source_node() const { return nodes_[Graph::kSourceId]; } + const SimpleNode* sink_node() const { return nodes_[Graph::kSinkId]; } private: - const tensorflow::Graph* g_; + const Graph* g_; std::vector nodes_; std::vector edges_; // free_edge_ids_ and free_node_ids_ contain freed indices. @@ -142,15 +138,14 @@ class SimpleGraph { std::set free_node_ids_; }; -SimpleNode::SimpleNode(const tensorflow::Node* node, const int id) - : node_(node), id_(id) { +SimpleNode::SimpleNode(const Node* node, const int id) : node_(node), id_(id) { if (node_) { in_edges_.reserve(node_->in_edges().size()); out_edges_.reserve(node_->out_edges().size()); } } -SimpleGraph::SimpleGraph(const tensorflow::Graph* g) : g_(g) { +SimpleGraph::SimpleGraph(const Graph* g) : g_(g) { int n_nodes = g_->num_node_ids(); nodes_.resize(n_nodes, nullptr); nodes_[g->kSourceId] = new SimpleNode(g->source_node(), g->kSourceId); @@ -194,8 +189,8 @@ void SimpleGraph::AddEdge(SimpleNode* src, int out_port, SimpleNode* dst, } else { edges_.push_back(nullptr); } - bool is_control = (out_port == tensorflow::Graph::kControlSlot); - is_control |= (in_port == tensorflow::Graph::kControlSlot); + bool is_control = (out_port == Graph::kControlSlot); + is_control |= (in_port == Graph::kControlSlot); auto edge = new SimpleEdge(i, src, out_port, dst, in_port, is_control); edges_[i] = edge; src->out_edges_.push_back(edge); @@ -203,8 +198,7 @@ void SimpleGraph::AddEdge(SimpleNode* src, int out_port, SimpleNode* dst, } void SimpleGraph::AddControlEdge(SimpleNode* src, SimpleNode* dst) { - AddEdge(src, tensorflow::Graph::kControlSlot, dst, - tensorflow::Graph::kControlSlot); + AddEdge(src, Graph::kControlSlot, dst, Graph::kControlSlot); } void SimpleGraph::RemoveEdge(const SimpleEdge* edge) { @@ -241,15 +235,14 @@ struct SimpleEdgePtrCompare { }; struct NodePtrCompare { - bool operator()(const tensorflow::Node* lhs, - const tensorflow::Node* rhs) const { + bool operator()(const Node* lhs, const Node* rhs) const { return lhs->name() < rhs->name(); } }; namespace { -// Copied from TF ReverseDFS, which only works for tensorflow::Graph. +// Copied from TF ReverseDFS, which only works for Graph. void StableDFS(const SimpleGraph& g, bool reverse, const std::vector& start, const std::function& enter, @@ -371,8 +364,7 @@ void ContractEdge(SimpleEdge* edge, SimpleGraph* graph, if (in_edge->src() != src) { SimpleEdge* e = const_cast(in_edge); if (e->src() == graph->source_node()) { - graph->AddEdge(e->src(), e->src_output(), src, - tensorflow::Graph::kControlSlot); + graph->AddEdge(e->src(), e->src_output(), src, Graph::kControlSlot); } else { graph->AddEdge(e->src(), e->src_output(), src, 0 /* input index */); } @@ -391,8 +383,7 @@ void ContractEdge(SimpleEdge* edge, SimpleGraph* graph, if (e->dst() == graph->sink_node()) { VLOG(1) << " edge to sink node " << src->name() << " -> " << e->dst()->name(); - graph->AddEdge(src, tensorflow::Graph::kControlSlot, e->dst(), - e->dst_input()); + graph->AddEdge(src, Graph::kControlSlot, e->dst(), e->dst_input()); } else { graph->AddEdge(src, 0 /* output index */, e->dst(), e->dst_input()); } @@ -410,12 +401,12 @@ void ContractEdge(SimpleEdge* edge, SimpleGraph* graph, } } -tensorflow::Status SegmentGraph( - const tensorflow::Graph* tf_graph, - const std::function& candidate_fn, - const std::function& input_candidate_fn, - const std::function& output_candidate_fn, - const SegmentOptions& options, SegmentNodesVector* segments) { +Status SegmentGraph(const Graph* tf_graph, + const std::function& candidate_fn, + const std::function& input_candidate_fn, + const std::function& output_candidate_fn, + const SegmentOptions& options, + SegmentNodesVector* segments) { // Steps: // 1. run the segmentation algorithm to find all the segments, which uses // candidate_fn to determine the candidates segment nodes; @@ -552,7 +543,7 @@ tensorflow::Status SegmentGraph( // A map from the segment identifier (currently the name of the root node of // the segment tree) to the segment nodes set. - std::map> sg_map; + std::map> sg_map; // A map from the segment identifier (currently the name of the root node of // the segment tree) to the device names that the nodes in the segment are @@ -578,7 +569,7 @@ tensorflow::Status SegmentGraph( device_maps[u.ParentValue()->name()].insert( tf_node->requested_device()); } else { - VLOG(1) << "Node " << tf_node->name() + VLOG(2) << "Node " << tf_node->name() << " has no device assigned requested device is: " << tf_node->requested_device(); } @@ -588,17 +579,16 @@ tensorflow::Status SegmentGraph( // --------------------------------- Step 2 --------------------------------- // Remove ineligible input/output nodes. for (auto& itr : sg_map) { - std::set& segment_nodes = - itr.second; + std::set& segment_nodes = itr.second; VLOG(1) << "Segment original size: " << segment_nodes.size(); while (true) { - std::deque in_nodes_que, out_nodes_que; + std::deque in_nodes_que, out_nodes_que; // Find an input node that is not eligible and add it to the queue. // Nodes that has no incoming edges should not be treated as "input", // as there are really no inputs to them. Similar for output nodes. for (auto node : segment_nodes) { bool added = false; - for (const tensorflow::Edge* edge : node->in_edges()) { + for (const Edge* edge : node->in_edges()) { if (!edge->IsControlEdge() && !edge->src()->IsSource() && !segment_nodes.count(edge->src())) { // 'node' is an input node. if (!input_candidate_fn(edge)) { @@ -609,7 +599,7 @@ tensorflow::Status SegmentGraph( } } if (added) continue; // Only adding the node once to either queue. - for (const tensorflow::Edge* edge : node->out_edges()) { + for (const Edge* edge : node->out_edges()) { if (!edge->dst()->IsSink() && !edge->IsControlEdge() && !segment_nodes.count(edge->dst())) { // 'node' is an output node. if (!output_candidate_fn(edge)) { @@ -637,13 +627,11 @@ tensorflow::Status SegmentGraph( // remove all their inputs, and for non-const output nodes remove all // their outputs. In this way, for common cases the number of removed // nodes should be minimum. - auto remove_nodes = [&segment_nodes]( - bool is_input_nodes, - std::deque* que) { + auto remove_nodes = [&segment_nodes](bool is_input_nodes, + std::deque* que) { // Run a BFS on the queue to find all the input/output nodes. - std::set visited; - std::set logged(que->begin(), - que->end()); + std::set visited; + std::set logged(que->begin(), que->end()); while (!que->empty()) { auto node = que->front(); que->pop_front(); @@ -722,7 +710,7 @@ tensorflow::Status SegmentGraph( VLOG(1) << "Devices " << s; } } - return tensorflow::Status::OK(); + return Status::OK(); } } // namespace segment diff --git a/tensorflow/compiler/tf2tensorrt/segment/segment.h b/tensorflow/compiler/tf2tensorrt/segment/segment.h index 9622ddd593990e93ba1b54e9dfd0052006e20ced..e31f1a989d9d9f203554811093e830ee8b139a6e 100644 --- a/tensorflow/compiler/tf2tensorrt/segment/segment.h +++ b/tensorflow/compiler/tf2tensorrt/segment/segment.h @@ -44,19 +44,19 @@ struct SegmentOptions { // Get the subgraphs of a graph that can be handled by TensorRT. // -// @param graph tensorflow::Graph of the network +// @param graph Graph of the network // @param candidate_fn A function that returns OK for a Node* if // that node can be handled by TensorRT. // @param segments Returns the TensorRT segments/subgraphs. Each entry // in the vector describes a subgraph by giving a set of the names of // all the NodeDefs in that subgraph. // @return the status. -tensorflow::Status SegmentGraph( - const tensorflow::Graph* tf_graph, - const std::function& candidate_fn, - const std::function& input_candidate_fn, - const std::function& output_candidate_fn, - const SegmentOptions& options, SegmentNodesVector* segments); +Status SegmentGraph(const Graph* tf_graph, + const std::function& candidate_fn, + const std::function& input_candidate_fn, + const std::function& output_candidate_fn, + const SegmentOptions& options, + SegmentNodesVector* segments); } // namespace segment } // namespace tensorrt diff --git a/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc b/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc index e11ad2719740d908f93ef580a6b308469365f402..84b690ecba6fcb9718a1008ee61383a84a381a46 100644 --- a/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc +++ b/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc @@ -33,13 +33,12 @@ namespace tensorflow { namespace tensorrt { namespace segment { namespace test { -namespace ops = ::tensorflow::ops; class SegmentTest : public ::testing::Test { protected: - std::function MakeCandidateFn( + std::function MakeCandidateFn( const std::set& node_names) { - return [node_names](const tensorflow::Node* node) -> Status { + return [node_names](const Node* node) -> Status { if (node_names.find(node->name()) != node_names.end()) { return Status::OK(); } @@ -47,22 +46,21 @@ class SegmentTest : public ::testing::Test { }; } - std::function MakeInputEdgeCandidateFn( + std::function MakeInputEdgeCandidateFn( const std::set& node_names) { - return [node_names](const tensorflow::Edge* in_edge) -> bool { + return [node_names](const Edge* in_edge) -> bool { return node_names.find(in_edge->dst()->name()) != node_names.end(); }; } - std::function MakeOutputEdgeCandidateFn( + std::function MakeOutputEdgeCandidateFn( const std::set& node_names) { - return [node_names](const tensorflow::Edge* out_edge) -> bool { + return [node_names](const Edge* out_edge) -> bool { return node_names.find(out_edge->src()->name()) != node_names.end(); }; } - void RunTest(const tensorflow::Graph* graph, - const std::set& candidates, + void RunTest(const Graph* graph, const std::set& candidates, const std::set& input_candidates, const std::set& output_candidates, const std::vector>& expected_segments) { @@ -106,7 +104,7 @@ std::set operator-(const std::set& lhs, const string& rhs) { TEST_F(SegmentTest, Empty) { Scope s = Scope::NewRootScope(); - tensorflow::Graph g(OpRegistry::Global()); + Graph g(OpRegistry::Global()); TF_EXPECT_OK(s.ToGraph(&g)); // Expect no segments/subgraphs. RunTest(&g, {}, {}, {}, {}); @@ -129,7 +127,7 @@ TEST_F(SegmentTest, Simple) { auto add2 = ops::Add(s.WithOpName("add2"), add0, add1); auto add3 = ops::Add(s.WithOpName("add3"), add0, add2); auto add4 = ops::Add(s.WithOpName("add4"), add2, add2); - tensorflow::Graph g(OpRegistry::Global()); + Graph g(OpRegistry::Global()); TF_EXPECT_OK(s.ToGraph(&g)); // All Add operations are candidates, and we expect all of them to be @@ -176,7 +174,7 @@ TEST_F(SegmentTest, AvoidCycle) { auto add2 = ops::Add(s.WithOpName("add2"), add0, add1); auto add3 = ops::Add(s.WithOpName("add3"), add0, add2); auto add4 = ops::Add(s.WithOpName("add4"), add2, add2); - tensorflow::Graph g(OpRegistry::Global()); + Graph g(OpRegistry::Global()); TF_EXPECT_OK(s.ToGraph(&g)); // add2 is not a TRT candidate so there should be no segments generated. @@ -207,7 +205,7 @@ TEST_F(SegmentTest, Multiple) { auto add3 = ops::Add(s.WithOpName("add3"), add0, add2); auto add4 = ops::Add(s.WithOpName("add4"), add2, add5); auto add6 = ops::Add(s.WithOpName("add6"), add5, add8); - tensorflow::Graph g(OpRegistry::Global()); + Graph g(OpRegistry::Global()); TF_EXPECT_OK(s.ToGraph(&g)); const std::set all_adds = {"add0", "add1", "add2", "add3", "add4", @@ -254,7 +252,7 @@ TEST_F(SegmentTest, BigIfElse) { auto add5 = ops::Add(s.WithOpName("add5"), add4, add4); auto add6 = ops::Add(s.WithOpName("add6"), add5, add5); auto add7 = ops::Add(s.WithOpName("add7"), add3, add6); - tensorflow::Graph g(OpRegistry::Global()); + Graph g(OpRegistry::Global()); TF_EXPECT_OK(s.ToGraph(&g)); // Make add2 not a TRT candidate, and we expect 2 segments. diff --git a/tensorflow/compiler/tf2tensorrt/utils/test_utils.h b/tensorflow/compiler/tf2tensorrt/utils/test_utils.h index d85875991b79014c4f173d3157ed02e6c96f045c..361eb01059fbc434cfc4ce499be495fe8230f74b 100644 --- a/tensorflow/compiler/tf2tensorrt/utils/test_utils.h +++ b/tensorflow/compiler/tf2tensorrt/utils/test_utils.h @@ -28,12 +28,11 @@ void ClearTestValues(const string& pattern); void AddTestValue(const string& label, const string& value); string GetTestValue(const string& label); -#define TRT_RETURN_IF_TEST_VALUE(label, value_to_return) \ - do { \ - if (::tensorflow::tensorrt::test::GetTestValue(label) == \ - value_to_return) { \ - return errors::Internal("Injected manually"); \ - } \ +#define TRT_RETURN_IF_TEST_VALUE(label, value_to_return) \ + do { \ + if (test::GetTestValue(label) == value_to_return) { \ + return errors::Internal("Injected manually"); \ + } \ } while (0) } // namespace test diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc index 1636cdc30c4df157ed124b160449af645f917252..a18f758a5512141ef180844dd4fabe960cbed4f2 100644 --- a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc @@ -72,7 +72,7 @@ void* TRTDeviceAllocator::allocate(uint64_t size, uint64_t alignment, uint32_t flags) { if (size == 0) return nullptr; // WAR for allocator alignment requirement. Certain cuda API calls require GPU - // memory with alignemtn to cudaDeviceProp::textureAlignment. + // memory with alignment to cudaDeviceProp::textureAlignment. // See issue #20856 alignment = 512; assert((alignment & (alignment - 1)) == 0); // zero or a power of 2. @@ -94,7 +94,7 @@ void* TRTDeviceAllocator::allocate(uint64_t size, uint64_t alignment, return mem; } -TRTDeviceAllocator::TRTDeviceAllocator(tensorflow::Allocator* allocator) +TRTDeviceAllocator::TRTDeviceAllocator(Allocator* allocator) : allocator_(allocator) { VLOG(1) << "Using " << allocator->Name() << " allocator from TensorFlow"; } diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h index 59ffb42bad348c78cde32035aff8c7081528b3a6..8ec06d7456c28505fe45859e42d83cc569d90dc5 100644 --- a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h @@ -59,7 +59,7 @@ class TRTCudaAllocator : public TRTBaseAllocator { class TRTDeviceAllocator : public TRTBaseAllocator { // Allocator implementation wrapping TF device allocators. public: - TRTDeviceAllocator(tensorflow::Allocator* allocator); + TRTDeviceAllocator(Allocator* allocator); // TODO(aaroey): base class doesn't have a virtual destructor, work with // Nvidia to fix it. @@ -70,7 +70,7 @@ class TRTDeviceAllocator : public TRTBaseAllocator { void free(void* memory) override; private: - tensorflow::Allocator* allocator_; + Allocator* allocator_; // supporting alignment from allocation request requires a map to free; std::unordered_map mem_map_; diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc index 5213fced1ea9220422245172f5b4a3f584a2a566..33a5c719ba9d750fc5ab173435512ef73ff3fce8 100644 --- a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc @@ -50,7 +50,7 @@ TRTInt8Calibrator::TRTInt8Calibrator(const string& calib_data) bool TRTInt8Calibrator::setBatch(const std::unordered_map& data, const cudaStream_t stream) { - tensorflow::mutex_lock lock(cond_mtx_); + mutex_lock lock(cond_mtx_); // Wait while the queue is full or calibration is running. while ((calib_running_ || batch_is_set_) && !done_) cond_.wait(lock); @@ -87,7 +87,7 @@ bool TRTInt8Calibrator::setBatch(const std::unordered_map& data, bool TRTInt8Calibrator::getBatch(void** bindings, const char** names, int num_bindings) { - tensorflow::mutex_lock lock(cond_mtx_); + mutex_lock lock(cond_mtx_); // Notify finish of last round of calibration. calib_running_ = false; cond_.notify_all(); @@ -111,7 +111,7 @@ bool TRTInt8Calibrator::getBatch(void** bindings, const char** names, } void TRTInt8Calibrator::waitAndSetDone() { - tensorflow::mutex_lock lock(cond_mtx_); + mutex_lock lock(cond_mtx_); // Wait while the queue is full or calibration is running, so we don't miss // the last batch. while ((calib_running_ || batch_is_set_) && !done_) cond_.wait(lock); @@ -128,7 +128,7 @@ const void* TRTInt8Calibrator::readCalibrationCache(std::size_t& length) { } void TRTInt8Calibrator::setDone() { - tensorflow::mutex_lock lock(cond_mtx_); + mutex_lock lock(cond_mtx_); done_ = true; cond_.notify_all(); } diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h index aa70b07f8d79848c362275815004db32cca128be..d34e244f6c7fe201915cb4b52808d3e0e3c57fa0 100644 --- a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h @@ -78,10 +78,10 @@ struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator { const int batch_size_; // mutex for condition_variable - tensorflow::mutex cond_mtx_; + mutex cond_mtx_; // condition variable to implement producer-consumer queue for calibration - tensorflow::condition_variable cond_; + condition_variable cond_; // Is calibration finished? bool done_; diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h index 09c47b36b0ad8074e749342e7d08f139da7ea1f4..8ece326446d9f3cb20d5ea02406e71e6e346446e 100644 --- a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h @@ -100,26 +100,24 @@ class LRUCache { } // Creates n free positions in cache - tensorflow::Status DiscardOld(size_t n = 0) { + Status DiscardOld(size_t n = 0) { if (n > capacity_) { - return tensorflow::errors::Internal( - "Insufficient capacity in cache (capacity = ", capacity_, - ", requested ", n, ")"); + return errors::Internal("Insufficient capacity in cache (capacity = ", + capacity_, ", requested ", n, ")"); } while (objects_.size() > (capacity_ - n)) { key_type discard_key = keys_.back(); keys_.pop_back(); objects_.erase(discard_key); } - return tensorflow::Status::OK(); + return Status::OK(); } }; // Define a hash function for vector because it is used as the key // for the engine cache. struct VectorTensorShapeHasher { - std::size_t operator()( - const std::vector& key) const { + std::size_t operator()(const std::vector& key) const { return std::hash()(TensorShapeUtils::ShapeListString(key)); } }; @@ -141,12 +139,12 @@ struct EngineContext { GUARDED_BY(mu); }; -class TRTEngineCacheResource : public tensorflow::ResourceBase { +class TRTEngineCacheResource : public ResourceBase { public: TRTEngineCacheResource(OpKernelContext* ctx, size_t capacity) : cache_(capacity) { auto device = ctx->device(); - auto alloc = device->GetAllocator(tensorflow::AllocatorAttributes()); + auto alloc = device->GetAllocator(AllocatorAttributes()); if (!alloc) { LOG(ERROR) << "Can't find device allocator for gpu device " << device->name(); diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_resources.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_resources.cc index 2e553079b19a3e5d0739cc6ac79a84f3b6a1fc4e..534e59f06b7d8f6768d1fc58e6a96cfe692fa14f 100644 --- a/tensorflow/compiler/tf2tensorrt/utils/trt_resources.cc +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_resources.cc @@ -49,7 +49,7 @@ Status TRTCalibrationResource::SerializeToString(string* serialized) { thr_->join(); *serialized = calibrator_->getCalibrationTableAsString(); if (serialized->empty()) { - return tensorflow::errors::Unknown("Calibration table is empty."); + return errors::Unknown("Calibration table is empty."); } return Status::OK(); } diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_resources.h b/tensorflow/compiler/tf2tensorrt/utils/trt_resources.h index 5e8d4b3b738df09b0c2ea82dcc06e9b23a708385..abfed2c1816732a6e7d7ef396d1923edf0d90f32 100644 --- a/tensorflow/compiler/tf2tensorrt/utils/trt_resources.h +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_resources.h @@ -37,7 +37,7 @@ limitations under the License. namespace tensorflow { namespace tensorrt { -class SerializableResourceBase : public tensorflow::ResourceBase { +class SerializableResourceBase : public ResourceBase { public: virtual Status SerializeToString(string* serialized) = 0; }; @@ -60,7 +60,7 @@ class TRTCalibrationResource : public SerializableResourceBase { TrtUniquePtrType builder_; TrtUniquePtrType engine_; std::unique_ptr allocator_; - tensorflow::tensorrt::Logger logger_; + Logger logger_; // TODO(sami): Use threadpool threads! std::unique_ptr thr_; }; diff --git a/tensorflow/compiler/tf2xla/const_analysis.cc b/tensorflow/compiler/tf2xla/const_analysis.cc index a57095f91e43f6b31b58e5a5f36331241451b545..6aff436da4f613a399c006b922b8aba3ce65a2e5 100644 --- a/tensorflow/compiler/tf2xla/const_analysis.cc +++ b/tensorflow/compiler/tf2xla/const_analysis.cc @@ -20,15 +20,26 @@ limitations under the License. #include "absl/algorithm/container.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/core/common_runtime/function.h" +#include "tensorflow/core/framework/attr_value.pb.h" +#include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/graph/algorithm.h" +#include "tensorflow/core/lib/core/errors.h" namespace tensorflow { + +Status GetCompileTimeConstInputs(const Node* node, + std::vector* const_input_idxs, + FunctionLibraryRuntime* flib_runtime); + // Backwards dataflow analysis that finds arguments to a graph that must be // compile-time constants. Status BackwardsConstAnalysis(const Graph& g, std::vector* compile_time_const_arg_indices, std::vector* compile_time_const_nodes, + FunctionLibraryRuntime* flib_runtime, std::function edge_filter) { std::vector compile_time_const_nodes_impl; if (compile_time_const_nodes) { @@ -61,7 +72,18 @@ Status BackwardsConstAnalysis(const Graph& g, } for (const Edge* pred : node->in_edges()) { if (!pred->IsControlEdge() && edge_filter(*pred)) { - (*compile_time_const_nodes)[pred->src()->id()] = true; + // If the src node of the `pred` is an IdentityN do not mark it as a + // compile-time const. Only mark the corresponding input to the + // IdentityN node as a const. + // Note: XLA IdentityN op simply forwards its inputs so this is safe. + while (edge_filter(*pred) && + pred->src()->type_string() == "IdentityN") { + status = pred->src()->input_edge(pred->src_output(), &pred); + if (!status.ok()) return; + } + if (edge_filter(*pred)) { + (*compile_time_const_nodes)[pred->src()->id()] = true; + } } } return; @@ -69,17 +91,29 @@ Status BackwardsConstAnalysis(const Graph& g, // Mark any compile-time constant operator arguments as const. std::vector const_input_idxs; - status = XlaOpRegistry::CompileTimeConstantInputs( - node->def(), node->op_def(), &const_input_idxs); + status = GetCompileTimeConstInputs(node, &const_input_idxs, flib_runtime); if (!status.ok()) { return; } for (Edge const* edge : node->in_edges()) { - if (absl::c_binary_search(const_input_idxs, edge->dst_input()) && + if (!edge->IsControlEdge() && + absl::c_binary_search(const_input_idxs, edge->dst_input()) && edge_filter(*edge)) { - (*compile_time_const_nodes)[edge->src()->id()] = true; + // Do not mark IdentityN nodes as compile-time const. + // If the src node of the `pred` is an IdentityN do not mark it as a + // compile-time const. Only mark the corresponding input to the + // IdentityN node as a const. + // Note: XLA IdentityN op simply forwards its inputs so this is safe. + while (edge_filter(*edge) && + edge->src()->type_string() == "IdentityN") { + status = edge->src()->input_edge(edge->src_output(), &edge); + if (!status.ok()) return; + } + if (edge_filter(*edge)) { + (*compile_time_const_nodes)[edge->src()->id()] = true; + } } } }; @@ -91,4 +125,61 @@ Status BackwardsConstAnalysis(const Graph& g, return status; } +Status GetCompileTimeConstInputs(const Node* node, + std::vector* const_input_idxs, + FunctionLibraryRuntime* flib_runtime) { + if (node->type_string() != "While") { + return XlaOpRegistry::CompileTimeConstantInputs(node->def(), node->op_def(), + const_input_idxs); + } + // For While nodes, recurse into the body and cond graphs. + // TODO(b/124403063): Implement similar functionality for cond nodes and other + // functional ops. + NameAttrList cond_function; + TF_RETURN_IF_ERROR(GetNodeAttr(node->def(), "cond", &cond_function)); + NameAttrList body_function; + TF_RETURN_IF_ERROR(GetNodeAttr(node->def(), "body", &body_function)); + FunctionLibraryRuntime::Handle cond_handle; + FunctionLibraryRuntime::Handle body_handle; + TF_RETURN_IF_ERROR(flib_runtime->Instantiate( + cond_function.name(), AttrSlice(&cond_function.attr()), &cond_handle)); + TF_RETURN_IF_ERROR(flib_runtime->Instantiate( + body_function.name(), AttrSlice(&body_function.attr()), &body_handle)); + const FunctionBody* fcond = flib_runtime->GetFunctionBody(cond_handle); + const FunctionBody* fbody = flib_runtime->GetFunctionBody(body_handle); + TF_RET_CHECK(fcond); + TF_RET_CHECK(fbody); + int num_inputs = fbody->fdef.signature().input_arg_size(); + + // Stores which of the loop inputs are expected to be compile time constants. + std::vector compile_time_const_arg_indices(num_inputs); + TF_RETURN_IF_ERROR(BackwardsConstAnalysis( + *(fcond->graph), &compile_time_const_arg_indices, + /*compile_time_const_nodes=*/nullptr, flib_runtime)); + TF_RETURN_IF_ERROR(BackwardsConstAnalysis( + *(fbody->graph), &compile_time_const_arg_indices, + /*compile_time_const_nodes=*/nullptr, flib_runtime)); + for (int i = 0; i < num_inputs; i++) { + if (compile_time_const_arg_indices[i]) { + // Check that this input is actually a loop invariant. + // NOTE(srbs): Ideally this should raise an error if the loop body + // requires the input at this index to be a compile time const but it is + // not a loop invariant. However, that causes problems because const + // analysis is performed for the entire graph (in the + // MarkForCompilationPass for example) and not just for the ops + // that will actually be run using XLA kernels. So we silently return here + // and let the error be raised during the actual compilation of the + // XLA graph. + Node* arg_i = fbody->arg_nodes[i]; + Node* ret_i = fbody->ret_nodes[i]; + const Node* ret_i_input_0; + TF_RETURN_IF_ERROR(ret_i->input_node(0, &ret_i_input_0)); + if (ret_i_input_0->id() == arg_i->id()) { + const_input_idxs->push_back(i); + } + } + } + return Status::OK(); +} + } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/const_analysis.h b/tensorflow/compiler/tf2xla/const_analysis.h index 49b3c6d413c6b637fa825bf182be7cc36e49b6c8..1663cbff41c3e10ba586c60eca475b760dee4896 100644 --- a/tensorflow/compiler/tf2xla/const_analysis.h +++ b/tensorflow/compiler/tf2xla/const_analysis.h @@ -34,11 +34,13 @@ namespace tensorflow { // `compile_time_const_nodes`, if `compile_time_const_nodes` is not null. // // Only propagate const-ness along edges for which `edge_filter` returns true. -Status BackwardsConstAnalysis(const Graph& g, - std::vector* compile_time_const_arg_indices, - std::vector* compile_time_const_nodes, - std::function edge_filter = - [](const Edge& e) { return true; }); +Status BackwardsConstAnalysis( + const Graph& g, std::vector* compile_time_const_arg_indices, + std::vector* compile_time_const_nodes, + FunctionLibraryRuntime* flib_runtime, + std::function edge_filter = [](const Edge& e) { + return true; + }); } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/const_analysis_test.cc b/tensorflow/compiler/tf2xla/const_analysis_test.cc index 40c6d0e01701d9104a200d9ea27706a0a7c12146..ed5f004550f0cb57e1545436c90bb6a9e8c19652 100644 --- a/tensorflow/compiler/tf2xla/const_analysis_test.cc +++ b/tensorflow/compiler/tf2xla/const_analysis_test.cc @@ -44,8 +44,8 @@ TEST(ConstAnalysisTest, Basics) { std::vector const_args(4, false); std::vector const_nodes(root.graph()->num_node_ids(), false); - TF_ASSERT_OK( - BackwardsConstAnalysis(*root.graph(), &const_args, &const_nodes)); + TF_ASSERT_OK(BackwardsConstAnalysis(*root.graph(), &const_args, &const_nodes, + /*flib_runtime=*/nullptr)); // Arg 0 doesn't need to be constant since the graph only uses its shape. // Arg 1 must be constant because it flows to the shape argument of a Reshape. @@ -82,7 +82,8 @@ TEST(ConstAnalysisTest, TopologicalOrder) { std::vector const_args(3, false); TF_ASSERT_OK(BackwardsConstAnalysis(graph, &const_args, - /*compile_time_const_nodes=*/nullptr)); + /*compile_time_const_nodes=*/nullptr, + /*flib_runtime=*/nullptr)); EXPECT_EQ(const_args, std::vector({true, true, false})); } @@ -103,7 +104,8 @@ TEST(ConstAnalysisTest, DontFollowControlDependencies) { std::vector const_args(2, false); TF_ASSERT_OK(BackwardsConstAnalysis(graph, &const_args, - /*compile_time_const_nodes=*/nullptr)); + /*compile_time_const_nodes=*/nullptr, + /*flib_runtime=*/nullptr)); EXPECT_EQ(const_args, std::vector({false, true})); } @@ -128,7 +130,8 @@ TEST(ConstAnalysisTest, RespectExplicitAttr_0) { std::vector const_args(2, false); TF_ASSERT_OK(BackwardsConstAnalysis(graph, &const_args, - /*compile_time_const_nodes=*/nullptr)); + /*compile_time_const_nodes=*/nullptr, + /*flib_runtime=*/nullptr)); EXPECT_EQ(const_args, std::vector({false, false})); } @@ -152,7 +155,8 @@ TEST(ConstAnalysisTest, RespectExplicitAttr_1) { std::vector const_args(1, false); TF_ASSERT_OK(BackwardsConstAnalysis(graph, &const_args, - /*compile_time_const_nodes=*/nullptr)); + /*compile_time_const_nodes=*/nullptr, + /*flib_runtime=*/nullptr)); EXPECT_EQ(const_args, std::vector({true})); } diff --git a/tensorflow/compiler/tf2xla/graph_compiler.cc b/tensorflow/compiler/tf2xla/graph_compiler.cc index 5e4699bbb6218089d2e76a36c7351bf7fbd23264..b3cb23003ec92ff2142fa4bcb71d943b236447ac 100644 --- a/tensorflow/compiler/tf2xla/graph_compiler.cc +++ b/tensorflow/compiler/tf2xla/graph_compiler.cc @@ -56,9 +56,9 @@ Status PrepareArguments(XlaOpKernelContext* ctx, Graph* graph, auto client = ctx->compiler()->client(); std::vector arg_must_be_compile_time_constant(expressions.size()); - TF_RETURN_IF_ERROR( - BackwardsConstAnalysis(*graph, &arg_must_be_compile_time_constant, - /*compile_time_const_nodes=*/nullptr)); + TF_RETURN_IF_ERROR(BackwardsConstAnalysis( + *graph, &arg_must_be_compile_time_constant, + /*compile_time_const_nodes=*/nullptr, ctx->function_library())); args->resize(expressions.size()); for (int i = 0; i < args->size(); ++i) { @@ -284,6 +284,7 @@ void GraphCompiler::PartiallySetupParams(OpKernelContext::Params* params) { params->inputs = &tensor_inputs_; params->step_container = step_container_; params->resource_manager = device_->resource_manager(); + params->function_library = flib_; } } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD index 343568b2392595a2347bde41f0a2e2559fb1de19..7ad4f912fd01b9ae3d293be05a2d088c085c9b8b 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -115,6 +115,7 @@ tf_kernel_library( ], tags = ["optonly"], deps = [ + ":case_op", ":conv_op_helpers", ":if_op", ":while_op", @@ -136,7 +137,6 @@ tf_kernel_library( "//tensorflow/compiler/xla/client:xla_builder", "//tensorflow/compiler/xla/client:xla_computation", "//tensorflow/compiler/xla/client/lib:arithmetic", - "//tensorflow/compiler/xla/client/lib:cholesky", "//tensorflow/compiler/xla/client/lib:constants", "//tensorflow/compiler/xla/client/lib:loops", "//tensorflow/compiler/xla/client/lib:math", @@ -151,6 +151,7 @@ tf_kernel_library( "//tensorflow/core:control_flow_ops_op_lib", "//tensorflow/core:data_flow_ops_op_lib", "//tensorflow/core:framework", + "//tensorflow/core:framework_bounds_check", "//tensorflow/core:functional_ops_op_lib", "//tensorflow/core:image_ops_op_lib", "//tensorflow/core:lib", @@ -168,14 +169,11 @@ tf_kernel_library( "//tensorflow/core:state_ops_op_lib", "//tensorflow/core:stateless_random_ops_op_lib", "//tensorflow/core:training_ops_op_lib", - "//tensorflow/core/kernels:bounds_check", - "//tensorflow/core/kernels:concat_lib", "//tensorflow/core/kernels:constant_op", "//tensorflow/core/kernels:control_flow_ops", "//tensorflow/core/kernels:conv_ops", "//tensorflow/core/kernels:cwise_op", "//tensorflow/core/kernels:list_kernels", - "//tensorflow/core/kernels:no_op", "//tensorflow/core/kernels:ops_util", "//tensorflow/core/kernels:pooling_ops", "//tensorflow/core/kernels:random_op", @@ -229,7 +227,7 @@ cc_library( "//tensorflow/compiler/xla/client/lib:arithmetic", "//tensorflow/compiler/xla/client/lib:constants", "//tensorflow/core:framework", - "//tensorflow/core/kernels:bounds_check", + "//tensorflow/core:framework_bounds_check", "//tensorflow/core/kernels:conv_ops", "//tensorflow/core/kernels:ops_util", "@com_google_absl//absl/types:span", @@ -251,6 +249,7 @@ tf_kernel_library( "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", + "@com_google_absl//absl/strings", ], ) @@ -271,6 +270,23 @@ tf_kernel_library( ], ) +tf_kernel_library( + name = "case_op", + srcs = ["case_op.cc"], + hdrs = ["case_op.h"], + deps = [ + "//tensorflow/compiler/tf2xla:common", + "//tensorflow/compiler/tf2xla:side_effect_util", + "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/tf2xla/ops:xla_ops", + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + ], +) + # Kernels that have a dummy (no-op) implementation. tf_kernel_library( name = "xla_dummy_ops", @@ -304,9 +320,8 @@ tf_kernel_library( "//tensorflow/compiler/xla/client:xla_builder", "//tensorflow/compiler/xla/client/lib:arithmetic", "//tensorflow/core:framework", + "//tensorflow/core:framework_bounds_check", "//tensorflow/core:lib", - "//tensorflow/core/kernels:argmax_op", - "//tensorflow/core/kernels:bounds_check", ], ) diff --git a/tensorflow/compiler/tf2xla/kernels/case_op.cc b/tensorflow/compiler/tf2xla/kernels/case_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..24623768f3897179575fe4cec6190a9a877a5202 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/case_op.cc @@ -0,0 +1,297 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/tf2xla/kernels/case_op.h" + +#include "tensorflow/compiler/tf2xla/shape_util.h" +#include "tensorflow/compiler/tf2xla/side_effect_util.h" +#include "tensorflow/compiler/tf2xla/xla_context.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/core/lib/core/errors.h" + +namespace tensorflow { + +XlaCaseOp::XlaCaseOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("branches", &branches_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("Tin", &input_types_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("Tout", &output_types_)); + if (!ctx->GetAttr(kXlaTokenInputNodesAttrName, &token_input_nodes_).ok()) { + has_token_input_output_ = false; + } else { + has_token_input_output_ = !token_input_nodes_.empty(); + } +} + +// TODO(b/35949885): There is duplication here with the handling of the +// while_op. Refactor the common code out/rework. +void XlaCaseOp::Compile(XlaOpKernelContext* ctx) { + xla::XlaBuilder* b = ctx->builder(); + int num_branches = branches_.size(); + OP_REQUIRES(ctx, num_branches >= 1, + errors::InvalidArgument("Must provide at least one case branch")); + OP_REQUIRES(ctx, input_type(0) == DT_INT32, + errors::InvalidArgument( + "branch_index argument must be a int32 for XLA compilation")); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(ctx->InputShape(0)), + errors::InvalidArgument( + "branch_index argument must be scalar for XLA compilation")); + + VLOG(1) << "Building Case: " << input_types_.size() << " inputs"; + + std::vector arguments(input_types_.size()); + int num_resource_args = 0; + for (int i = 0; i < input_types_.size(); ++i) { + XlaCompiler::Argument& arg = arguments[i]; + DataType type = ctx->input_type(i + 1); + + if (type == DT_RESOURCE) { + XlaResource* resource; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(i + 1, &resource)); + + arg.initialized = resource->initialized(); + arg.kind = XlaCompiler::Argument::kResource; + arg.resource_kind = resource->kind(); + + arg.type = resource->type(); + arg.shape = resource->shape(); + OP_REQUIRES(ctx, arg.initialized, + errors::Unimplemented("Uninitialized arguments: ", arg.name)); + arg.max_array_size = resource->max_array_size(); + for (const auto& gradient : resource->tensor_array_gradients()) { + arg.tensor_array_gradients.insert(gradient.first); + } + arg.name = resource->name(); + VLOG(2) << "Resource " << resource->name() + << " type: " << DataTypeString(arg.type) + << " shape: " << arg.HumanString() + << " initialized: " << arg.initialized; + + num_resource_args++; + } else { + arg.kind = XlaCompiler::Argument::kParameter; + arg.type = input_types_[i]; + arg.shape = ctx->InputShape(i + 1); + VLOG(2) << "Arg type: " << DataTypeString(arg.type) + << " shape: " << arg.HumanString(); + } + } + + // Compile each branch of the conditional. + XlaCompiler::CompileOptions options; + options.use_tuple_arg = true; + options.resolve_compile_time_constants = false; + options.return_updated_values_for_all_resources = true; + options.is_entry_computation = false; + options.add_token_input_output = has_token_input_output_; + XlaCompiler* compiler = ctx->compiler(); + + std::vector branch_results(num_branches); + std::vector branch_results_p(num_branches); + for (int j = 0; j < num_branches; ++j) { + OP_REQUIRES_OK(ctx, + compiler->CompileFunction(options, branches_[j], arguments, + &branch_results[j])); + branch_results_p[j] = &branch_results[j]; + } + + bool has_tensor_array_gradients = false; + for (XlaCompiler::CompilationResult* result : branch_results_p) { + for (const XlaCompiler::ResourceUpdate& update : result->resource_updates) { + XlaResource* resource; + OP_REQUIRES_OK(ctx, + ctx->GetResourceInput(update.input_index + 1, &resource)); + XlaCompiler::Argument& arg = arguments[update.input_index]; + + // Add any TensorArray gradients touched by the then/else computation to + // the enclosing graph. + for (const string& grad_source : update.tensor_array_gradients_accessed) { + VLOG(5) << "TensorArray " << resource->name() << " accessed gradient " + << grad_source; + XlaResource* gradient; + OP_REQUIRES_OK(ctx, resource->GetOrCreateTensorArrayGradient( + grad_source, b, &gradient)); + } + // Add all of the TensorArray gradients to the argument. For simplicity, + // we always pass all known gradients. + for (const auto& gradient : resource->tensor_array_gradients()) { + arg.tensor_array_gradients.insert(gradient.first); + } + if (!resource->tensor_array_gradients().empty()) { + has_tensor_array_gradients = true; + } + } + } + + // Recompile the functions to update the argument shapes for tensor arrays. + if (has_tensor_array_gradients) { + for (int j = 0; j < num_branches; ++j) { + branch_results[j] = {}; + OP_REQUIRES_OK(ctx, + compiler->CompileFunction(options, branches_[j], arguments, + &branch_results[j])); + } + } + + xla::Shape branch0_input_shape; + std::vector result_computations(num_branches); + for (int j = 0; j < num_branches; ++j) { + // Check that all branches have identical input shapes. + OP_REQUIRES(ctx, branch_results[j].xla_input_shapes.size() == 1, + errors::FailedPrecondition("Expected one input shape")); + xla::Shape branch_input_shape = branch_results[j].xla_input_shapes[0]; + if (j == 0) { + branch0_input_shape = branch_input_shape; + } + OP_REQUIRES(ctx, branch_input_shape.IsTuple(), + errors::FailedPrecondition("Expected tuple shape")); + OP_REQUIRES(ctx, branch_results[j].xla_input_shapes.size() == 1, + errors::FailedPrecondition("Expected one input shape")); + OP_REQUIRES( + ctx, + xla::ShapeUtil::Compatible(branch0_input_shape, branch_input_shape), + errors::InvalidArgument( + "Input shapes of 0 and ", j, " branches do not match: ", + xla::ShapeUtil::HumanString(branch0_input_shape), " vs. ", + xla::ShapeUtil::HumanString(branch_input_shape))); + + // Check that all branches have identical output shapes. + OP_REQUIRES( + ctx, + xla::ShapeUtil::Compatible(branch_results[0].xla_output_shape, + branch_results[j].xla_output_shape), + errors::InvalidArgument( + "Output shapes of 0 and ", j, " branches do not match: ", + xla::ShapeUtil::HumanString(branch_results[0].xla_output_shape), + " vs. ", + xla::ShapeUtil::HumanString(branch_results[j].xla_output_shape))); + + if (j == 0) { + VLOG(2) << "Input shape: " + << xla::ShapeUtil::HumanString(branch0_input_shape); + VLOG(2) << "Output shape: " + << xla::ShapeUtil::HumanString( + branch_results[0].xla_output_shape); + } + + // We set return_updated_values_for_all_resources=true and we pass the same + // arguments to both computations, so the resource update count must match. + OP_REQUIRES(ctx, + branch_results[0].resource_updates.size() == + branch_results[j].resource_updates.size(), + errors::FailedPrecondition( + "Different number of resources in 0 and ", j, " branch")); + for (int i = 0; i < branch_results[0].resource_updates.size(); ++i) { + const auto& lhs = branch_results[0].resource_updates[i]; + const auto& rhs = branch_results[j].resource_updates[i]; + bool equal = lhs.input_index == rhs.input_index && + lhs.shape == rhs.shape && + lhs.tensor_array_gradients_accessed == + rhs.tensor_array_gradients_accessed; + OP_REQUIRES(ctx, equal, + errors::FailedPrecondition("Mismatch in resource of 0 and ", + j, " branch for resource ", i)); + } + result_computations[j] = branch_results[j].computation.get(); + } + + // Prepare the input arg Tuple. + int num_inputs = branch_results[0].input_mapping.size(); + std::vector inputs(num_inputs); + for (int i = 0; i < num_inputs; ++i) { + int input_num = branch_results[0].input_mapping[i] + 1; + if (has_token_input_output_ && i == num_inputs - 1) { + // Set token input for this "case" op. + std::vector token_inputs; + for (const string& node_name : token_input_nodes_) { + auto token_or = compiler->GetNodeToken(node_name); + OP_REQUIRES_OK(ctx, token_or.status()); + token_inputs.push_back(token_or.ValueOrDie()); + } + inputs[i] = xla::AfterAll(b, token_inputs); + } else if (ctx->input_type(input_num) == DT_RESOURCE) { + XlaResource* resource; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(input_num, &resource)); + OP_REQUIRES_OK(ctx, resource->Pack(&inputs[i], b)); + } else { + inputs[i] = ctx->Input(i + 1); + } + } + auto input_tuple = xla::Tuple(b, inputs); + + xla::XlaOp outputs = + xla::Conditional(ctx->Input(0), absl::MakeSpan(result_computations), + std::vector(num_branches, input_tuple)); + // Sets non-variable outputs. + for (int i = 0; i < output_types_.size(); ++i) { + xla::XlaOp output_handle = xla::GetTupleElement(outputs, i); + if (VLOG_IS_ON(2)) { + LOG(INFO) << "Setting output " << i; + auto shape_or = b->GetShape(output_handle); + if (shape_or.ok()) { + LOG(INFO) << "Shape for output " << i << ": " + << xla::ShapeUtil::HumanString(shape_or.ValueOrDie()); + } else { + LOG(INFO) << "Shape unknown for output " << i; + } + } + ctx->SetOutput(i, output_handle); + } + if (has_token_input_output_) { + // Set token output for this "Case" op. Token output is the last output of + // XLA computation, which comes after all "normal" TF outputs and resource + // updates. For "Case" node, num of resource updates equals to number of + // resource args because we set `return_updated_values_for_all_resources` + // to true in XlaCompiler option. + xla::XlaOp token_output = + xla::GetTupleElement(outputs, output_types_.size() + num_resource_args); + auto shape_or = b->GetShape(token_output); + OP_REQUIRES_OK(ctx, shape_or.status()); + OP_REQUIRES(ctx, shape_or.ValueOrDie().IsToken(), + errors::FailedPrecondition( + "Token output is not token type: ", + xla::ShapeUtil::HumanString(shape_or.ValueOrDie()))); + OP_REQUIRES_OK(ctx, compiler->SetNodeToken(name(), token_output)); + } + + // Updates the values of any resource variables modified by the conditional + // bodies. + for (const XlaCompiler::CompilationResult& result : branch_results) { + for (int i = 0; i < result.resource_updates.size(); ++i) { + const XlaCompiler::ResourceUpdate& update = result.resource_updates[i]; + XlaResource* resource; + OP_REQUIRES_OK(ctx, + ctx->GetResourceInput(update.input_index + 1, &resource)); + if (update.modified) { + int pos = static_cast(result.outputs.size()) + i; + OP_REQUIRES_OK(ctx, + resource->SetFromPack( + arguments[update.input_index].tensor_array_gradients, + xla::GetTupleElement(outputs, pos), b)); + } + VLOG(2) << "Case variable: pos: " << update.input_index + << " name: " << resource->name() + << " modified: " << update.modified + << " type: " << DataTypeString(update.type) + << " shape: " << update.shape.DebugString(); + } + } + VLOG(1) << "Done building Case"; +} + +REGISTER_XLA_OP(Name("Case").AllowResourceTypes(), XlaCaseOp); + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/case_op.h b/tensorflow/compiler/tf2xla/kernels/case_op.h new file mode 100644 index 0000000000000000000000000000000000000000..ea14b18149cb5bc9162d42b384eb3a5e943ad8be --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/case_op.h @@ -0,0 +1,62 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_CASE_OP_H_ +#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_CASE_OP_H_ + +#include +#include + +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/core/framework/attr_value.pb.h" +#include "tensorflow/core/framework/types.h" + +namespace tensorflow { + +// This TensorFlow op provides a functional switch/case primitive. +// +// The outputs of the branches must agree on the number, types, and +// shapes of the Tensors carried around the two bodies. +// +// Computations in branch bodies may read from and write to resource variables. +// Resource variables may be passed as arguments to the branch function's +// bodies. The XlaCompiler converts resource variable arguments +// into parameters to the XLA computation and moves them to the end of the +// parameter list, and by using the `return_updated_values_for_all_variables` +// we ensure that all variables that appear in the input also appear at the +// end of the branch bodies output. This ensures the branch bodies output +// signatures match. +// +// It is the user's responsibility to ensure that each non-variable _Arg matches +// the corresponding _Retval. +class XlaCaseOp : public XlaOpKernel { + public: + explicit XlaCaseOp(OpKernelConstruction* ctx); + + void Compile(XlaOpKernelContext* ctx) override; + + private: + TF_DISALLOW_COPY_AND_ASSIGN(XlaCaseOp); + + std::vector branches_; + DataTypeVector input_types_; + DataTypeVector output_types_; + bool has_token_input_output_; + std::vector token_input_nodes_; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_TF2XLA_KERNELS_CASE_OP_H_ diff --git a/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc b/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc index 0ed3044efa5b1060d2b0ad2d5563b0e02ebf66ec..e6b30a38e0379fc09af07af686f4c5f3a737ecda 100644 --- a/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc @@ -15,7 +15,8 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" -#include "tensorflow/compiler/xla/client/lib/cholesky.h" +#include "tensorflow/compiler/xla/client/lib/matrix.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" namespace tensorflow { namespace { @@ -24,7 +25,9 @@ class CholeskyOp : public XlaOpKernel { public: explicit CholeskyOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} void Compile(XlaOpKernelContext* ctx) override { - ctx->SetOutput(0, xla::Cholesky(ctx->Input(0))); + ctx->SetOutput(0, + xla::Triangle(xla::Cholesky(ctx->Input(0), /*lower=*/true), + /*lower=*/true)); } }; diff --git a/tensorflow/compiler/tf2xla/kernels/concat_op.cc b/tensorflow/compiler/tf2xla/kernels/concat_op.cc index 91e4d9cea7cbf6075e30250587044174c4b8e7f4..09c97de13eb2ed951ca705cda89b7f293808cdf0 100644 --- a/tensorflow/compiler/tf2xla/kernels/concat_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/concat_op.cc @@ -31,7 +31,6 @@ limitations under the License. #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/concat_lib.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" diff --git a/tensorflow/compiler/tf2xla/kernels/elu_op.cc b/tensorflow/compiler/tf2xla/kernels/elu_op.cc index 5fdb1d972c55efb876972d3f472b53a1f7cde1c2..87bb9d49c0c97181bac33da01ec7e0b10cf5d6fc 100644 --- a/tensorflow/compiler/tf2xla/kernels/elu_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/elu_op.cc @@ -22,7 +22,6 @@ limitations under the License. #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/core/framework/kernel_def_builder.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/no_op.h" namespace tensorflow { namespace { diff --git a/tensorflow/compiler/tf2xla/kernels/index_ops_kernel_argmax_float_1d.cc b/tensorflow/compiler/tf2xla/kernels/index_ops_kernel_argmax_float_1d.cc index 47cf8c6675bc120653c2a5ab6d4b07376dc382ee..39d96e748b3a2a852c03c0dd53ec175f0c66a43a 100644 --- a/tensorflow/compiler/tf2xla/kernels/index_ops_kernel_argmax_float_1d.cc +++ b/tensorflow/compiler/tf2xla/kernels/index_ops_kernel_argmax_float_1d.cc @@ -25,9 +25,6 @@ limitations under the License. namespace tensorflow { EIGEN_STRONG_INLINE void argmax_float_1d_xla_impl(void* out, void** data) { - // data is managed by the JIT code so msan can't tell it's initialized. - TF_ANNOTATE_MEMORY_IS_INITIALIZED(data, 2 * sizeof(void*)); - float* input = static_cast(data[0]); int64 input_size = *static_cast(data[1]); diff --git a/tensorflow/compiler/tf2xla/kernels/l2loss_op.cc b/tensorflow/compiler/tf2xla/kernels/l2loss_op.cc index 93f029731c34e84000a3dc00df8af05654cccf2d..7f25d34c3ef82e5360fd2d7c1cd12dd8c6f40507 100644 --- a/tensorflow/compiler/tf2xla/kernels/l2loss_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/l2loss_op.cc @@ -19,7 +19,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/no_op.h" namespace tensorflow { namespace { diff --git a/tensorflow/compiler/tf2xla/kernels/no_op.cc b/tensorflow/compiler/tf2xla/kernels/no_op.cc index 65ab9da8d7ca0509a4a69c43727a0e6c0435908a..da50b75251beb2f97400cc7d2ffb5f4d05a3fb6e 100644 --- a/tensorflow/compiler/tf2xla/kernels/no_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/no_op.cc @@ -13,12 +13,22 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/core/kernels/no_op.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/core/framework/kernel_def_builder.h" namespace tensorflow { +namespace { + +class NoOp : public OpKernel { + public: + explicit NoOp(OpKernelConstruction* context) : OpKernel(context) {} + void Compute(OpKernelContext* context) override {} + bool IsExpensive() override { return false; } +}; + +} // namespace + // XLA_* devices also register a "real" NoOp operator so we suppress the // dummy operator using CompilationOnly(). REGISTER_XLA_OP(Name("NoOp").CompilationOnly(), NoOp); diff --git a/tensorflow/compiler/tf2xla/kernels/pack_op.cc b/tensorflow/compiler/tf2xla/kernels/pack_op.cc index 426a0941df57f19072d1cb9f3fa3d0079db465c5..6ca100a2f2bf90e1d61829aa45a44cbc97090ed1 100644 --- a/tensorflow/compiler/tf2xla/kernels/pack_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/pack_op.cc @@ -30,7 +30,6 @@ limitations under the License. #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/concat_lib.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" diff --git a/tensorflow/compiler/tf2xla/kernels/scan_ops.cc b/tensorflow/compiler/tf2xla/kernels/scan_ops.cc index daefdfc58a4957d9e685d25aa90da6218f2041ad..8431724f438f67c07740212e1e31926777fef3ae 100644 --- a/tensorflow/compiler/tf2xla/kernels/scan_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/scan_ops.cc @@ -30,7 +30,6 @@ limitations under the License. #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/concat_lib.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" diff --git a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc index b6c96b1f582710e1cc39e6e1e0e800ef8170743d..a93d137e96519837ae289f08ff4d32960970aad9 100644 --- a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc @@ -31,7 +31,6 @@ limitations under the License. #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/concat_lib.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc index 77a3e5c001e1c715f23ae5148f94dae2faa81acf..b98b98ce50af2cb811297989899b06d33296bf13 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc @@ -34,7 +34,6 @@ limitations under the License. #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/concat_lib.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc index 8958a48bc79dce91c41ab7d0a5fc0fbb401112ba..4feb17d2c86e63bc9f2a6c0ee57764d35505971a 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc @@ -35,7 +35,6 @@ limitations under the License. #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/concat_lib.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" diff --git a/tensorflow/compiler/tf2xla/kernels/topk_op.cc b/tensorflow/compiler/tf2xla/kernels/topk_op.cc index ee3bdf3394e37c757f31724e73e95417becaa534..22cfd16008899c1ad3c73453bec34a0b0d2e8c78 100644 --- a/tensorflow/compiler/tf2xla/kernels/topk_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/topk_op.cc @@ -20,7 +20,6 @@ limitations under the License. #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/core/framework/kernel_def_builder.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/no_op.h" namespace tensorflow { namespace { diff --git a/tensorflow/compiler/tf2xla/kernels/training_ops.cc b/tensorflow/compiler/tf2xla/kernels/training_ops.cc index 26d4214099d1d07c1b2e275d783654d9cd948e28..247db8d5d172b04e414b1ff0e53f12b533f36944 100644 --- a/tensorflow/compiler/tf2xla/kernels/training_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/training_ops.cc @@ -22,7 +22,6 @@ limitations under the License. #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/core/framework/kernel_def_builder.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/no_op.h" namespace tensorflow { namespace { @@ -856,15 +855,12 @@ class ResourceApplyAdadelta : public XlaOpKernel { xla::XlaOp grad = ctx->Input(6); xla::XlaBuilder* b = ctx->builder(); - xla::XlaOp neg_half = XlaHelpers::FloatLiteral(b, dtype_, -0.5); - xla::XlaOp half = XlaHelpers::FloatLiteral(b, dtype_, 0.5); xla::XlaOp one = XlaHelpers::FloatLiteral(b, dtype_, 1.0); - xla::XlaOp two = XlaHelpers::FloatLiteral(b, dtype_, 2.0); - accum = rho * accum + (one - rho) * xla::Pow(grad, two); - xla::XlaOp update = xla::Pow(accum_update + epsilon, half) * - xla::Pow(accum + epsilon, neg_half) * grad; - accum_update = rho * accum_update + (one - rho) * xla::Pow(update, two); + accum = rho * accum + (one - rho) * xla::Square(grad); + xla::XlaOp update = + xla::Sqrt(accum_update + epsilon) * xla::Rsqrt(accum + epsilon) * grad; + accum_update = rho * accum_update + (one - rho) * xla::Square(update); var = var - update * lr; OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, dtype_, var)); OP_REQUIRES_OK(ctx, ctx->AssignVariable(1, dtype_, accum)); diff --git a/tensorflow/compiler/tf2xla/kernels/unpack_op.cc b/tensorflow/compiler/tf2xla/kernels/unpack_op.cc index 2fc5619de737b8977e4249e4d2297a0303c339ce..2d95f2f30a86f3a9c95e528858c53ab48d7a02e8 100644 --- a/tensorflow/compiler/tf2xla/kernels/unpack_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/unpack_op.cc @@ -30,7 +30,6 @@ limitations under the License. #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/concat_lib.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" diff --git a/tensorflow/compiler/tf2xla/kernels/while_op.cc b/tensorflow/compiler/tf2xla/kernels/while_op.cc index f49da9683b3622bdda708cc305306baafa1639df..58637302d4a33caf358dea8efd8cae10948eb4c5 100644 --- a/tensorflow/compiler/tf2xla/kernels/while_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/while_op.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/kernels/while_op.h" +#include "absl/strings/str_split.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/side_effect_util.h" #include "tensorflow/compiler/tf2xla/type_util.h" @@ -25,11 +26,14 @@ limitations under the License. #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/client/xla_computation.h" #include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/op_kernel.h" namespace tensorflow { +const char kPropagateCompileTimeConsts[] = "_xla_propagate_compile_time_consts"; + namespace { // Builds XlaCompiler argument descriptions `args` from `ctx`. @@ -89,6 +93,154 @@ Status MakeXlaCompilerArgumentsFromInputs( return Status::OK(); } +// Populates loop invariant indices to true in `loop_invariants`. +void GetLoopInvariants(XlaOpKernelContext* ctx, + const NameAttrList& body_name_attr, + std::vector* const loop_invariants) { + const FunctionBody* body; + OP_REQUIRES_OK(ctx, ctx->compiler()->FindFunctionBody(body_name_attr, &body)); + for (int i = 0; i < body->ret_nodes.size(); i++) { + const Node* arg = body->arg_nodes[i]; + const Node* ret = body->ret_nodes[i]; + const Node* ret_input_0; + OP_REQUIRES_OK(ctx, ret->input_node(0, &ret_input_0)); + (*loop_invariants)[i] = ret_input_0->id() == arg->id(); + } +} + +// Converts entries in `args` which are loop invariants and have compile +// time constant inputs to constants so that they can be propagated in the loop +// body. +Status ConvertLoopInvariantsToConst( + XlaOpKernelContext* ctx, const NameAttrList& body_name_attr, + std::vector* args, + std::vector* compile_time_const_arg_indices, + int* num_compile_time_const_args, xla::Client* client) { + std::vector loop_invariants(ctx->num_inputs()); + GetLoopInvariants(ctx, body_name_attr, &loop_invariants); + for (int i = 0; i < ctx->num_inputs(); i++) { + XlaCompiler::Argument& arg = (*args)[i]; + const XlaExpression& expression = ctx->InputExpression(i); + // If this is a loop invariant and the input tensor is a compile time + // constant build a kConstant type argument. + if (arg.kind != XlaCompiler::Argument::kResource && loop_invariants[i]) { + // NOTE: We can not simple check that this is Kind::kConstant because + // this could be the output of a MetadataOnly op e.g. Size. + xla::StatusOr> maybe_constant = + expression.ResolveConstant(client); + if (maybe_constant.ok() && maybe_constant.ValueOrDie().has_value()) { + arg.kind = XlaCompiler::Argument::kConstant; + arg.type = expression.dtype(); + arg.constant_value = std::move(maybe_constant.ValueOrDie().value()); + arg.shape = expression.GetShape().ValueOrDie(); + compile_time_const_arg_indices->at(i) = true; + (*num_compile_time_const_args)++; + } + } + } + return Status::OK(); +} + +Status VerifyBodyInputAndOutputShapeMatch( + XlaOpKernelContext* ctx, + const std::vector& compile_time_const_arg_indices, + const XlaCompiler::CompilationResult& body, bool has_token_input_output) { + xla::Shape body_input_shape = body.xla_input_shapes[0]; + xla::Shape body_output_shape; + body_output_shape.set_element_type(xla::TUPLE); + for (int i = 0; i < ctx->num_outputs(); i++) { + if (!compile_time_const_arg_indices[i]) { + *(body_output_shape.add_tuple_shapes()) = + body.xla_output_shape.tuple_shapes(i); + } + } + // If `body` has a token output, append its shape to `body_output_shape`. + if (has_token_input_output) { + *(body_output_shape.add_tuple_shapes()) = + body.xla_output_shape.tuple_shapes(ctx->num_inputs()); + } + if (!xla::ShapeUtil::Compatible(body_input_shape, body_output_shape)) { + return errors::InvalidArgument( + "Input and output shapes of loop body do not match: ", + xla::ShapeUtil::HumanString(body_input_shape), " vs. ", + xla::ShapeUtil::HumanString(body_output_shape)); + } + return Status::OK(); +} + +xla::StatusOr BuildWrappedCond( + XlaOpKernelContext* ctx, const XlaCompiler::CompilationResult& cond) { + xla::Shape cond_input_shape = cond.xla_input_shapes[0]; + std::unique_ptr cb = + ctx->builder()->CreateSubBuilder("cond_wrapper"); + auto inputs = xla::Parameter(cb.get(), 0, cond_input_shape, "inputs"); + auto outputs = xla::Call(cb.get(), *cond.computation, {inputs}); + xla::GetTupleElement(outputs, 0); + return cb->Build(); +} + +xla::StatusOr BuildWrappedBody( + XlaOpKernelContext* ctx, const XlaCompiler::CompilationResult& body, + const std::vector& compile_time_const_arg_indices, + int num_compile_time_const_args, bool has_token_input_output) { + if (num_compile_time_const_args <= 0) { + return xla::XlaComputation(body.computation->proto()); + } + xla::XlaComputation body_wrapper; + std::unique_ptr cb = + ctx->builder()->CreateSubBuilder("body_wrapper"); + xla::Shape body_input_shape = body.xla_input_shapes[0]; + auto inputs = xla::Parameter(cb.get(), 0, body_input_shape, "inputs"); + // Call the original body function which has mismatched inputs and outputs + // and strip the compile time consts from the list of outputs. While requires + // the inputs and outputs of its body function to match. + auto outputs = xla::Call(cb.get(), *body.computation, {inputs}); + std::vector non_compile_time_const_outputs; + for (int i = 0; i < compile_time_const_arg_indices.size(); i++) { + if (!compile_time_const_arg_indices[i]) { + non_compile_time_const_outputs.push_back( + xla::GetTupleElement(outputs, i)); + } + } + // If `body` has a token output, append it to + // `non_compile_time_const_outputs`. + if (has_token_input_output) { + non_compile_time_const_outputs.push_back( + xla::GetTupleElement(outputs, ctx->num_outputs())); + } + xla::Tuple(cb.get(), non_compile_time_const_outputs); + return cb->Build(); +} + +xla::XlaOp BuildWhile(XlaOpKernelContext* ctx, + const xla::XlaComputation& wrapped_cond, + const xla::XlaComputation& wrapped_body, + const xla::XlaOp& initial_values, + const std::vector& input_mapping, + const std::vector& compile_time_const_arg_indices, + int num_compile_time_const_args, + bool has_token_input_output) { + xla::XlaOp while_result = + xla::While(wrapped_cond, wrapped_body, initial_values); + std::vector padded_while_outputs(ctx->num_outputs()); + int while_result_index = 0; + for (int i = 0; i < ctx->num_inputs(); i++) { + if (!compile_time_const_arg_indices[i]) { + padded_while_outputs[input_mapping[while_result_index]] = + xla::GetTupleElement(while_result, while_result_index); + while_result_index++; + } else { + padded_while_outputs[i] = ctx->Input(i); + } + } + // If `body` has a token output, append it to `padded_while_outputs`. + if (has_token_input_output) { + padded_while_outputs.push_back(xla::GetTupleElement( + while_result, ctx->num_inputs() - num_compile_time_const_args)); + } + return xla::Tuple(ctx->builder(), padded_while_outputs); +} + } // anonymous namespace XlaWhileOp::XlaWhileOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { @@ -102,6 +254,10 @@ XlaWhileOp::XlaWhileOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { } else { has_token_input_output_ = !token_input_nodes_.empty(); } + if (ctx->HasAttr(kPropagateCompileTimeConsts)) { + OP_REQUIRES_OK(ctx, ctx->GetAttr(kPropagateCompileTimeConsts, + &propagate_compile_time_consts_)); + } } void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { @@ -117,6 +273,25 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { xla::XlaBuilder* builder = ctx->builder(); XlaCompiler* compiler = ctx->compiler(); + // Indices of loop vars which satisfy the following conditions: + // 1. They are loop invariants. + // 2. The op inputs at these indices are compile time constants. + // + // These compile time consts do not appear as _Args in the cond/body functions + // and are replaced by kConstant nodes instead. As as result, the compiled + // body function does not have matching input and output shape. We fix this + // by rewriting the body computation (see body_wrapper below) to output + // just the non compile-time-const values and later pad up the while output + // with the const args. + std::vector compile_time_const_arg_indices(ctx->num_inputs()); + int num_compile_time_const_args = 0; + if (propagate_compile_time_consts_) { + OP_REQUIRES_OK(ctx, ConvertLoopInvariantsToConst( + ctx, body_name_attr_, &arguments, + &compile_time_const_arg_indices, + &num_compile_time_const_args, compiler->client())); + } + VLOG(1) << "Compiling body"; // All resource that are inputs to the loop's body must also be @@ -232,12 +407,13 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { "Input shapes of loop body and condition do not match: ", xla::ShapeUtil::HumanString(body_input_shape), " vs. ", xla::ShapeUtil::HumanString(cond_input_shape))); - OP_REQUIRES( - ctx, xla::ShapeUtil::Compatible(body_input_shape, body.xla_output_shape), - errors::InvalidArgument( - "Input and output shapes of loop body do not match: ", - xla::ShapeUtil::HumanString(body_input_shape), " vs. ", - xla::ShapeUtil::HumanString(body.xla_output_shape))); + + // Check that the shape of the body outputs excluding the compile time const + // args (which are pruned from the body outputs in body_wapper) matches the + // shape of the inputs. + OP_REQUIRES_OK(ctx, VerifyBodyInputAndOutputShapeMatch( + ctx, compile_time_const_arg_indices, body, + has_token_input_output_)); xla::Shape expected_cond_output_shape_without_side_effect = xla::ShapeUtil::MakeTupleShape( @@ -275,7 +451,7 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { OP_REQUIRES_OK(ctx, ctx->GetResourceInput(input_num, &resource)); OP_REQUIRES_OK(ctx, resource->Pack(&inputs[i], builder)); } else { - inputs[i] = ctx->Input(i); + inputs[i] = ctx->Input(input_num); } } @@ -284,26 +460,28 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { VLOG(1) << "Building while loop"; // Wraps the condition in a computation that unpacks the output tuple. - xla::XlaComputation cond_wrapper; - { - std::unique_ptr cb = - builder->CreateSubBuilder("cond_wrapper"); - auto inputs = xla::Parameter(cb.get(), 0, cond_input_shape, "inputs"); - auto outputs = xla::Call(cb.get(), *cond.computation, {inputs}); - xla::GetTupleElement(outputs, 0); - xla::StatusOr result = cb->Build(); - OP_REQUIRES_OK(ctx, result.status()); - cond_wrapper = std::move(result.ValueOrDie()); - } - - xla::XlaOp while_result = xla::While(cond_wrapper, *body.computation, init); + xla::StatusOr cond_result = BuildWrappedCond(ctx, cond); + OP_REQUIRES_OK(ctx, cond_result.status()); + xla::XlaComputation wrapped_cond = std::move(cond_result.ValueOrDie()); + + // Remove compile time const args from the list of body outputs. + xla::StatusOr body_result = + BuildWrappedBody(ctx, body, compile_time_const_arg_indices, + num_compile_time_const_args, has_token_input_output_); + OP_REQUIRES_OK(ctx, body_result.status()); + xla::XlaComputation wrapped_body = std::move(body_result.ValueOrDie()); + + // Builds the While op and pads its output with the compile time const args. + xla::XlaOp while_result = + BuildWhile(ctx, wrapped_cond, wrapped_body, init, body.input_mapping, + compile_time_const_arg_indices, num_compile_time_const_args, + has_token_input_output_); // Sets non-variable outputs and determine when resource variables start. int resource_index = 0; for (int i = 0; i < ctx->num_outputs(); ++i) { if (ctx->input_type(i) != DT_RESOURCE) { - ctx->SetOutput(body.input_mapping[i], - xla::GetTupleElement(while_result, i)); + ctx->SetOutput(i, xla::GetTupleElement(while_result, i)); ++resource_index; } else { break; diff --git a/tensorflow/compiler/tf2xla/kernels/while_op.h b/tensorflow/compiler/tf2xla/kernels/while_op.h index aeeff40e68f8b778628b9e85bd9b4ddcb73883a5..16ec8d0e520b5a282318f8e5225bcec65818e3e8 100644 --- a/tensorflow/compiler/tf2xla/kernels/while_op.h +++ b/tensorflow/compiler/tf2xla/kernels/while_op.h @@ -21,6 +21,8 @@ limitations under the License. namespace tensorflow { +extern const char kPropagateCompileTimeConsts[]; + // This TensorFlow op provides a functional iteration primitive. // // The inputs and outputs of the loop body must agree on the number, types, and @@ -58,6 +60,10 @@ class XlaWhileOp : public XlaOpKernel { NameAttrList body_name_attr_; bool has_token_input_output_; std::vector token_input_nodes_; + // Whether to propagate compile time consts into the loop body. + // This is not supported by default now since it may cause HBM memory + // overheads. + bool propagate_compile_time_consts_ = false; TF_DISALLOW_COPY_AND_ASSIGN(XlaWhileOp); }; diff --git a/tensorflow/compiler/tf2xla/tf2xla.cc b/tensorflow/compiler/tf2xla/tf2xla.cc index 28a4566c9d284fb8410a2d618f368c4dd2c1d893..6db027e2acd61f1f52a2da33ee31aad8eb070ab9 100644 --- a/tensorflow/compiler/tf2xla/tf2xla.cc +++ b/tensorflow/compiler/tf2xla/tf2xla.cc @@ -265,8 +265,11 @@ Status ConvertGraphToXla(std::unique_ptr graph, std::vector xla_args; TF_RETURN_IF_ERROR(CreateXlaArgs(*graph, &xla_args)); + std::vector xla_aliases; // Populate arguments with resource variables from the config. The variables // get turned into inputs and outputs. + int64 input_num = xla_args.size(); + int64 output_num = config.fetch_size(); for (const tf2xla::Variable& variable : config.variable()) { XlaCompiler::Argument arg; arg.type = variable.type(); @@ -276,6 +279,13 @@ Status ConvertGraphToXla(std::unique_ptr graph, arg.resource_kind = XlaResource::kVariable; arg.initialized = true; xla_args.push_back(std::move(arg)); + + // We want to alias the input and output of the variable, so the updates are + // carried out in-place. + xla_aliases.push_back({/*output_index=*/{output_num}, + /*param_number=*/input_num, /*param_index=*/{}}); + ++input_num; + ++output_num; } // Compile the graph into an XLA computation. @@ -290,7 +300,7 @@ Status ConvertGraphToXla(std::unique_ptr graph, XlaCompiler::CompilationResult result; TF_RETURN_IF_ERROR(compiler.CompileGraph(XlaCompiler::CompileOptions(), "tfcompile", std::move(graph), - xla_args, &result)); + xla_args, xla_aliases, &result)); *computation = std::move(*result.computation); int num_const_results = 0; diff --git a/tensorflow/compiler/tf2xla/xla_compilation_device.cc b/tensorflow/compiler/tf2xla/xla_compilation_device.cc index 5bd0277c051711f2677b90a2679662899521e94a..f98d07d196ea8551f1a5b53fa2e88e7bc43639de 100644 --- a/tensorflow/compiler/tf2xla/xla_compilation_device.cc +++ b/tensorflow/compiler/tf2xla/xla_compilation_device.cc @@ -42,7 +42,7 @@ class XlaCompilationAllocator : public Allocator { void* AllocateRaw(size_t alignment, size_t num_bytes) override { // Regardless of the size requested, always allocates an XlaExpression. - // Respects the aligment request because there is alignment checking even + // Respects the alignment request because there is alignment checking even // for Tensors whose data is never accessed. void* p = port::AlignedMalloc(sizeof(XlaExpression), alignment); XlaExpression* expression = reinterpret_cast(p); diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc index 3221ec5b727de1f792cd61b792ee917588d56cf9..f4a0f456ed5420d60d8284fe9f235a1d3171600b 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler.cc @@ -603,7 +603,7 @@ Status XlaCompiler::CompileFunction( VLOG(1) << "===================================================="; TF_RETURN_IF_ERROR( - CompileGraph(options, function_id, std::move(graph), args, result)); + CompileGraph(options, function_id, std::move(graph), args, {}, result)); VLOG(1) << "===================================================="; cache_[{function_id, arg_vector}] = *result; @@ -936,7 +936,8 @@ Status XlaCompiler::CompileSingleOp( } FixupSourceAndSinkEdges(graph.get()); - return CompileGraph(options, node_def.name(), std::move(graph), args, result); + return CompileGraph(options, node_def.name(), std::move(graph), args, {}, + result); } namespace { @@ -1019,11 +1020,11 @@ void ConvertConstantsToExpressions(xla::XlaBuilder* builder, } // namespace -Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, - string const& name, - std::unique_ptr graph, - absl::Span args, - CompilationResult* result) { +Status XlaCompiler::CompileGraph( + const XlaCompiler::CompileOptions& options, string const& name, + std::unique_ptr graph, absl::Span args, + absl::Span user_aliases, + CompilationResult* result) { VLOG(1) << "Executing graph symbolically to populate XlaBuilder."; TF_RETURN_IF_ERROR(PropagateConstIntoFunctionalNodes( @@ -1071,6 +1072,12 @@ Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, options.is_entry_computation)); context->set_args(std::move(arg_expressions)); + // Propagate any aliases given to us by the user. + for (const xla::XlaBuilder::InputOutputAlias& alias : user_aliases) { + builder.SetUpAlias(alias.output_index, alias.param_number, + alias.param_index); + } + PushNodeTokenMapping(); // Use std::set instead of std::unordered_set to ensure determinism. std::set output_node_token_inputs; diff --git a/tensorflow/compiler/tf2xla/xla_compiler.h b/tensorflow/compiler/tf2xla/xla_compiler.h index ad3144b41bdf3fc8b75ab5230e8e128df2962884..0b0908e9d6913f2664e4d976611b1218be44ff2b 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.h +++ b/tensorflow/compiler/tf2xla/xla_compiler.h @@ -339,10 +339,11 @@ class XlaCompiler { // Compiles a tensorflow::Graph into an xla::XlaComputation. // Similar to CompileFunction, but takes a Graph as input rather than a // function. - Status CompileGraph(const CompileOptions& options, string const& name, - std::unique_ptr graph, - absl::Span args, - CompilationResult* result); + Status CompileGraph( + const CompileOptions& options, string const& name, + std::unique_ptr graph, absl::Span args, + absl::Span user_aliases, + CompilationResult* result); // Compiles a single Op, given by `node_def`, into an // xla::XlaComputation. Similar to CompileFunction but takes a single Op as @@ -416,11 +417,11 @@ class XlaCompiler { Status SetNodeToken(const string& node_name, const xla::XlaOp& op); xla::StatusOr GetNodeToken(const string& node_name); - private: // Sets the function body `fbody` to the one registered as `function`. Status FindFunctionBody(const NameAttrList& function, const FunctionBody** fbody); + private: // Returns the optimized graph object in this function body. std::unique_ptr GetGraph(const FunctionBody* fbody); diff --git a/tensorflow/compiler/tf2xla/xla_compiler_test.cc b/tensorflow/compiler/tf2xla/xla_compiler_test.cc index b31137867d738944eaaa73e142ad8538ec6b854a..1818d4290324aa398f8f90ff11725dc48948b621 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler_test.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler_test.cc @@ -175,9 +175,9 @@ TEST_F(XlaCompilerTest, EmptyReturnValues) { std::unique_ptr graph(new Graph(OpRegistry::Global())); XlaCompiler::CompilationResult result; - TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", - std::move(graph), - /*args=*/{}, &result)); + TF_ASSERT_OK(compiler.CompileGraph( + XlaCompiler::CompileOptions(), "add", std::move(graph), + /*args=*/{}, /*user_aliases=*/{}, &result)); TF_ASSERT_OK(client_->Execute(*result.computation, {}).status()); } @@ -207,7 +207,8 @@ TEST_F(XlaCompilerTest, Simple) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", - std::move(graph), args, &result)); + std::move(graph), args, + /*user_aliases=*/{}, &result)); // Tests that the generated computation works. xla::Literal param0_literal = xla::LiteralUtil::CreateR1({7, 42}); @@ -258,7 +259,7 @@ TEST_F(XlaCompilerTest, OutOfOrderGraph) { compile_options.always_return_tuple = false; XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(compile_options, "add", std::move(graph), - args, &result)); + args, /*user_aliases=*/{}, &result)); // Tests that the generated computation works. xla::Literal param0_literal = xla::LiteralUtil::CreateR1({7, 42}); @@ -319,7 +320,8 @@ TEST_F(XlaCompilerTest, HonorShapeRepresentationFnForRetVal) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", - std::move(graph), args, &result)); + std::move(graph), args, + /*user_aliases=*/{}, &result)); xla::Shape transposed = xla::ShapeUtil::MakeShapeWithLayout(xla::S32, {2, 3}, {0, 1}); // Check that the return shapes are correctly tranposed. @@ -360,7 +362,8 @@ TEST_F(XlaCompilerTest, TransposeVariables) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "transpose", - std::move(graph), args, &result)); + std::move(graph), args, + /*user_aliases=*/{}, &result)); xla::Shape transposed = xla::ShapeUtil::MakeShapeWithLayout(xla::S32, {2, 3}, {1, 0}); // Check that the return shapes are correctly tranposed. @@ -410,7 +413,7 @@ TEST_F(XlaCompilerTest, MixedOrderArguments) { compile_options.always_return_tuple = false; XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(compile_options, "add", std::move(graph), - args, &result)); + args, /*user_aliases=*/{}, &result)); EXPECT_THAT(result.input_mapping, ::testing::ElementsAre(0, 1)); } @@ -440,9 +443,9 @@ TEST_F(XlaCompilerTest, HasSaneErrorOnNonCompileTimeConstantInputToReshape) { XlaCompiler compiler(DefaultOptions()); XlaCompiler::CompilationResult result; - Status status = - compiler.CompileGraph(XlaCompiler::CompileOptions(), "reshape", - std::move(graph), args, &result); + Status status = compiler.CompileGraph(XlaCompiler::CompileOptions(), + "reshape", std::move(graph), args, + /*user_aliases=*/{}, &result); EXPECT_FALSE(status.ok()); EXPECT_TRUE( absl::StrContains(status.error_message(), "depends on a parameter")) @@ -486,7 +489,8 @@ TEST_F(XlaCompilerTest, ConstantOutputs) { compile_options.resolve_compile_time_constants = true; XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(compile_options, "constants", - std::move(graph_copy), args, &result)); + std::move(graph_copy), args, + /*user_aliases=*/{}, &result)); ASSERT_EQ(2, result.outputs.size()); EXPECT_TRUE(result.outputs[0].is_constant); @@ -519,7 +523,8 @@ TEST_F(XlaCompilerTest, ConstantOutputs) { compile_options.resolve_compile_time_constants = false; XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(compile_options, "constants", - std::move(graph_copy), args, &result)); + std::move(graph_copy), args, + /*user_aliases=*/{}, &result)); ASSERT_EQ(2, result.outputs.size()); EXPECT_FALSE(result.outputs[0].is_constant); @@ -605,7 +610,8 @@ TEST_F(XlaCompilerTest, ConstantOutputsOfFunctionalNode) { compile_options.resolve_compile_time_constants = true; XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(compile_options, "constants", - std::move(graph), args, &result)); + std::move(graph), args, + /*user_aliases=*/{}, &result)); ASSERT_EQ(2, result.outputs.size()); EXPECT_TRUE(result.outputs[0].is_constant); @@ -647,7 +653,8 @@ TEST_F(XlaCompilerTest, ResourceManager) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "dummy", - std::move(graph), args, &result)); + std::move(graph), args, + /*user_aliases=*/{}, &result)); EXPECT_EQ(1, resource->Get()); @@ -683,7 +690,8 @@ TEST_F(XlaCompilerTest, DeterministicCompilation) { XlaCompiler compiler(options); TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "dummy", - std::move(graph), args, &results[i])); + std::move(graph), args, + /*user_aliases=*/{}, &results[i])); } for (int64 i = 1; i < test_count; ++i) { @@ -749,7 +757,8 @@ TEST_F(XlaCompilerTest, CanPassTensorArraysToAndFromComputation) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", - std::move(graph), args, &result)); + std::move(graph), args, + /*user_aliases=*/{}, &result)); ASSERT_EQ(1, result.resource_updates.size()); const XlaCompiler::ResourceUpdate& update = result.resource_updates[0]; @@ -808,7 +817,8 @@ TEST_F(XlaCompilerTest, UnwrittenTensorArrayGradientsAreNotComputationOutputs) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", - std::move(graph), args, &result)); + std::move(graph), args, + /*user_aliases=*/{}, &result)); EXPECT_EQ(0, result.resource_updates.size()); } @@ -840,7 +850,8 @@ TEST_F(XlaCompilerTest, NewTensorArrayGradientsAreComputationOutputs) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", - std::move(graph), args, &result)); + std::move(graph), args, + /*user_aliases=*/{}, &result)); EXPECT_EQ(1, result.resource_updates.size()); } @@ -915,7 +926,8 @@ TEST_F(XlaCompilerTest, FunctionCallWithConstants) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "fill", - std::move(graph), args, &result)); + std::move(graph), args, + /*user_aliases=*/{}, &result)); } // Tests CompileFunction with a local function lookup failing, fails with @@ -998,7 +1010,8 @@ TEST_F(XlaCompilerTest, Variables) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", - std::move(graph), args, &result)); + std::move(graph), args, + /*user_aliases=*/{}, &result)); RunAndCheckVariablesComputation(client_, result); } @@ -1033,7 +1046,7 @@ TEST_F(XlaCompilerTest, ResultLayoutSingle) { auto compile_options = XlaCompiler::CompileOptions(); compile_options.always_return_tuple = false; TF_ASSERT_OK(compiler.CompileGraph(compile_options, "id", std::move(graph), - args, &result)); + args, /*user_aliases=*/{}, &result)); EXPECT_TRUE(xla::ShapeUtil::Equal( result.xla_output_shape, xla::ShapeUtil::MakeShapeWithLayout(xla::S32, {2, 3}, {0, 1}))); @@ -1069,7 +1082,8 @@ TEST_F(XlaCompilerTest, ResultLayoutMultiple) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "id", - std::move(graph), args, &result)); + std::move(graph), args, + /*user_aliases=*/{}, &result)); xla::Shape result_shape = xla::ShapeUtil::MakeShapeWithLayout(xla::S32, {2, 3}, {0, 1}); @@ -1099,7 +1113,8 @@ TEST_F(XlaCompilerTest, ReturnResourceHandleOnly) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", - std::move(graph), args, &result)); + std::move(graph), args, + /*user_aliases=*/{}, &result)); // Tests that the generated computation works. xla::Literal param1_literal = xla::LiteralUtil::CreateR1({-3, 101}); @@ -1149,7 +1164,8 @@ TEST_F(XlaCompilerTest, ReturnResourceHandle) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", - std::move(graph), args, &result)); + std::move(graph), args, + /*user_aliases=*/{}, &result)); RunAndCheckVariablesComputation(client_, result); } @@ -1200,7 +1216,7 @@ TEST_F(XlaCompilerTest, VariableRepresentationShapeFunction) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(compile_options, "add", std::move(graph), - args, &result)); + args, /*user_aliases=*/{}, &result)); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr program_shape, client_->GetComputationShape(*result.computation)); @@ -1270,7 +1286,7 @@ TEST_F(XlaCompilerTest, ArgRetvalShapeRepresentationFunction) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(compile_options, "add", std::move(graph), - args, &result)); + args, /*user_aliases=*/{}, &result)); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr program_shape, client_->GetComputationShape(*result.computation)); @@ -1351,7 +1367,8 @@ TEST_F(XlaCompilerTest, FunctionWithInvalidOp) { std::vector args; XlaCompiler::CompilationResult result; status = compiler.CompileGraph(XlaCompiler::CompileOptions(), "fill", - std::move(graph), args, &result); + std::move(graph), args, /*user_aliases=*/{}, + &result); ASSERT_FALSE(status.ok()); EXPECT_TRUE(absl::StrContains(status.error_message(), "InvalidOp")) << status.error_message(); @@ -1376,7 +1393,8 @@ TEST_F(XlaCompilerTest, NodeWithInvalidDataType) { XlaCompiler::CompilationResult result; XlaCompiler compiler(DefaultOptions()); status = compiler.CompileGraph(XlaCompiler::CompileOptions(), "invalid_type", - std::move(graph), args, &result); + std::move(graph), args, /*user_aliases=*/{}, + &result); ASSERT_FALSE(status.ok()); EXPECT_TRUE(absl::StrContains(status.error_message(), "is not in the list of allowed values")) @@ -1402,7 +1420,8 @@ TEST_F(XlaCompilerTest, SingleOpWithoutInputs) { CopyGraph(*graph, graph_copy.get()); XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "NoOp", - std::move(graph_copy), args, &result)); + std::move(graph_copy), args, + /*user_aliases=*/{}, &result)); } } @@ -1451,7 +1470,7 @@ TEST_F(XlaCompilerTest, TokenInputAndOutput) { CopyGraph(*graph, graph_copy.get()); XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(options, "NoOp", std::move(graph_copy), - args, &result)); + args, /*user_aliases=*/{}, &result)); EXPECT_EQ(result.xla_input_shapes.size(), 1); EXPECT_TRUE(result.xla_output_shape.IsTuple()); EXPECT_EQ(xla::ShapeUtil::TupleElementCount(result.xla_output_shape), 1); @@ -1469,7 +1488,7 @@ TEST_F(XlaCompilerTest, TokenInputAndOutput) { CopyGraph(*graph, graph_copy.get()); XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(options, "NoOp", std::move(graph_copy), - args, &result)); + args, /*user_aliases=*/{}, &result)); EXPECT_EQ(result.xla_input_shapes.size(), 2); EXPECT_TRUE(result.xla_input_shapes[1].IsToken()); EXPECT_TRUE(result.xla_output_shape.IsTuple()); diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.h b/tensorflow/compiler/tf2xla/xla_op_registry.h index c5e078a02d1ca6fdd8405ae6556a5205e387421e..80d022b592c911bff740577b46c765dbab1c2833 100644 --- a/tensorflow/compiler/tf2xla/xla_op_registry.h +++ b/tensorflow/compiler/tf2xla/xla_op_registry.h @@ -51,10 +51,10 @@ constexpr std::array kNumericTypes = { {DT_UINT8, DT_UINT32, DT_UINT64, DT_INT8, DT_INT32, DT_INT64, DT_HALF, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64, DT_COMPLEX128, DT_BFLOAT16}}; -constexpr std::array kCpuAllTypes = { +constexpr std::array kCpuAllTypes = { {DT_UINT8, DT_QUINT8, DT_UINT32, DT_UINT64, DT_INT8, DT_QINT8, DT_INT32, DT_QINT32, DT_INT64, DT_HALF, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64, - DT_COMPLEX128, DT_BOOL}}; + DT_COMPLEX128, DT_BOOL, DT_BFLOAT16}}; constexpr std::array kGpuAllTypes = { {DT_UINT8, DT_QUINT8, DT_UINT32, DT_UINT64, DT_INT8, DT_QINT8, DT_INT32, diff --git a/tensorflow/compiler/xla/client/lib/BUILD b/tensorflow/compiler/xla/client/lib/BUILD index c5dea5f18030f2d226c86e3408ea85b2b5989728..d24a0d652a4a7265fc352a45a10eb8489887f740 100644 --- a/tensorflow/compiler/xla/client/lib/BUILD +++ b/tensorflow/compiler/xla/client/lib/BUILD @@ -49,47 +49,6 @@ xla_test( ], ) -cc_library( - name = "cholesky", - srcs = ["cholesky.cc"], - hdrs = ["cholesky.h"], - deps = [ - ":math", - "//tensorflow/compiler/xla:literal", - "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:status_macros", - "//tensorflow/compiler/xla:statusor", - "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:xla_builder", - "//tensorflow/compiler/xla/client/lib:constants", - "//tensorflow/compiler/xla/client/lib:loops", - "//tensorflow/compiler/xla/client/lib:matrix", - "//tensorflow/compiler/xla/client/lib:slicing", - "//tensorflow/core:lib", - ], -) - -xla_test( - name = "cholesky_test", - srcs = ["cholesky_test.cc"], - tags = ["optonly"], - deps = [ - ":arithmetic", - ":cholesky", - ":matrix", - "//tensorflow/compiler/xla:array2d", - "//tensorflow/compiler/xla:literal", - "//tensorflow/compiler/xla:statusor", - "//tensorflow/compiler/xla:test", - "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla/client:xla_builder", - "//tensorflow/compiler/xla/tests:client_library_test_base", - "//tensorflow/compiler/xla/tests:literal_test_util", - "//tensorflow/compiler/xla/tests:xla_internal_test_main", - "//tensorflow/core:test", - ], -) - cc_library( name = "comparators", srcs = ["comparators.cc"], @@ -195,6 +154,7 @@ xla_test( name = "math_test", srcs = ["math_test.cc"], deps = [ + ":constants", ":math", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", @@ -207,22 +167,6 @@ xla_test( ], ) -xla_test( - name = "math_exhaustive_test", - srcs = ["math_exhaustive_test.cc"], - shard_count = 16, - deps = [ - ":math", - "//tensorflow/compiler/xla:literal_util", - "//tensorflow/compiler/xla:test", - "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:xla_builder", - "//tensorflow/compiler/xla/tests:client_library_test_base", - "//tensorflow/compiler/xla/tests:xla_internal_test_main", - ], -) - cc_library( name = "matrix", srcs = ["matrix.cc"], diff --git a/tensorflow/compiler/xla/client/lib/cholesky.h b/tensorflow/compiler/xla/client/lib/cholesky.h deleted file mode 100644 index 0bae26837c0f14dd0cfab82cf426becc787ec11c..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/client/lib/cholesky.h +++ /dev/null @@ -1,39 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_CHOLESKY_H_ -#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_CHOLESKY_H_ - -#include "tensorflow/compiler/xla/client/xla_builder.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" - -namespace xla { - -// Computes the Cholesky decompositions of a batch of symmetric positive -// definite matrices. -// `a` must be a (batched) square matrix; i.e., it must have rank >= 2 with the -// two minor dimensions equal. -// The algorithm implements a blocked Cholesky decomposition; `block_size` is -// the block size to use. -// TODO(phawkins): check for negative values on the diagonal and return an -// error, instead of silently yielding NaNs. -// TODO(znado): handle the complex Hermitian case -xla::XlaOp Cholesky( - xla::XlaOp a, int64 block_size = 256, - xla::PrecisionConfig::Precision precision = xla::PrecisionConfig::HIGHEST); - -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_CHOLESKY_H_ diff --git a/tensorflow/compiler/xla/client/lib/math.cc b/tensorflow/compiler/xla/client/lib/math.cc index 19d98d100191fcba590d64c643d76b2bc5d5e5c5..20d3c0fc549d9cbb14c8d8e271ff386a06b5ecab 100644 --- a/tensorflow/compiler/xla/client/lib/math.cc +++ b/tensorflow/compiler/xla/client/lib/math.cc @@ -27,6 +27,29 @@ limitations under the License. namespace xla { +// Returns operation(operand), except if `operand` is one of the types in +// upcast_types, in which case first converts it to F32, and then converts the +// result down to the original type. +static XlaOp DoWithUpcastToF32(XlaOp operand, + absl::Span upcast_types, + const std::function& operation) { + auto& b = *operand.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(auto shape, b.GetShape(operand)); + PrimitiveType elem_ty = shape.element_type(); + bool needs_upcast = absl::c_linear_search(upcast_types, elem_ty); + + if (needs_upcast) { + operand = ConvertElementType(operand, F32); + } + XlaOp result = operation(operand); + if (needs_upcast) { + result = ConvertElementType(result, elem_ty); + } + return result; + }); +} + // TODO(jlebar): Use this function in more places in this file to restrict the // domain of other functions. static Status EnsureOperandIsRealFp(absl::string_view op_name, XlaOp operand) { @@ -107,10 +130,6 @@ XlaOp IsNegZero(XlaOp operand) { }); } -XlaOp Sqrt(XlaOp operand) { return Pow(operand, ScalarLike(operand, 0.5)); } - -XlaOp Rsqrt(XlaOp operand) { return Pow(operand, ScalarLike(operand, -0.5)); } - XlaOp Square(XlaOp operand) { return operand * operand; } XlaOp Reciprocal(XlaOp operand) { return ScalarLike(operand, 1.0) / operand; } @@ -187,11 +206,17 @@ XlaOp Erfc(XlaOp x) { auto& b = *x.builder(); return b.ReportErrorOrReturn([&]() -> StatusOr { TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("Erfc", x)); + // erfc(x) = // erfc_impl(x) if x > 1 // 1 - erf_impl(x) otherwise - return Select(Gt(Abs(x), ScalarLike(x, 1)), ErfcImpl(x), - ScalarLike(x, 1) - ErfImpl(x)); + // + // Erf(c)Impl don't have enough precision when run with bf16 intermediates + // (not surprising!), so upcast to f32 in this case. + return DoWithUpcastToF32(x, {BF16}, [](XlaOp x) { + return Select(Gt(Abs(x), ScalarLike(x, 1)), ErfcImpl(x), + ScalarLike(x, 1) - ErfImpl(x)); + }); }); } @@ -202,8 +227,13 @@ XlaOp Erf(XlaOp x) { // erf(x) = // erf_impl(x) if x < 1 // 1 - erfc_impl(x) otherwise - return Select(Lt(Abs(x), ScalarLike(x, 1)), ErfImpl(x), - ScalarLike(x, 1) - ErfcImpl(x)); + // + // Erf(c)Impl don't have enough precision when run with bf16 intermediates + // (not surprising!), so upcast to f32 in this case. + return DoWithUpcastToF32(x, {BF16}, [](XlaOp x) { + return Select(Lt(Abs(x), ScalarLike(x, 1)), ErfImpl(x), + ScalarLike(x, 1) - ErfcImpl(x)); + }); }); } @@ -243,7 +273,18 @@ XlaOp ErfInv(XlaOp x) { for (int i = 1; i < kDegree; ++i) { p = coefficient(i) + p * w; } - return p * x; + + // Result modulo edge cases. + XlaOp result = p * x; + + // Handle edge cases, namely erfinv(+/-1) = +/-inf. (The above computation is + // indeterminate, and can give nan or -/+inf.) + auto& b = *x.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape shape, b.GetShape(x)); + return Select(Eq(Abs(x), ScalarLike(x, 1)), + x * MaxValue(&b, shape.element_type()), result); + }); } namespace { @@ -270,10 +311,7 @@ static constexpr std::array kLanczosCoefficients = { // t(z) = z + kLanczosGamma + 1/2 // A(z) = kBaseLanczosCoeff + sigma(k = 1, n, kLanczosCoefficients[i] / (z + k)) XlaOp Lgamma(XlaOp input) { - auto& b = *input.builder(); - return b.ReportErrorOrReturn([&]() -> StatusOr { - TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("Lgamma", input)); - + auto do_it = [](XlaOp input) { XlaOp one_half = ScalarLike(input, 0.5); XlaOp one = ScalarLike(input, 1); @@ -309,7 +347,16 @@ XlaOp Lgamma(XlaOp input) { XlaOp log_t = log_lanczos_gamma_plus_one_half + Log1p(z / lanczos_gamma_plus_one_half); - XlaOp log_y = log_sqrt_two_pi + (z + one_half) * log_t - t + Log(x); + // Compute the final result (modulo reflection). t(z) may be large, and we + // need to be careful not to overflow to infinity in the first term of + // + // (z + 1/2) * log(t(z)) - t(z). + // + // Therefore we compute this as + // + // (z + 1/2 - t(z) / log(t(z))) * log(t(z)). + // + XlaOp log_y = log_sqrt_two_pi + (z + one_half - t / log_t) * log_t + Log(x); // Compute the reflected value, used when x < 0.5: // @@ -324,18 +371,27 @@ XlaOp Lgamma(XlaOp input) { // important. // // Because abs(sin(pi * x)) has period 1, we can equivalently use - // abs(sin(pi * frac(x))) = sin(pi * frac(x)), where frac(x) is the - // fractional part of x. This is more numerically accurate: It doesn't - // overflow to inf like pi * x can, and if x is an integer, it evaluates to - // 0 exactly, which is significant because we then take the log of this - // value, and log(0) is inf. + // abs(sin(pi * frac(x))), where frac(x) is the fractional part of x. This + // is more numerically accurate: It doesn't overflow to inf like pi * x can, + // and if x is an integer, it evaluates to 0 exactly, which is significant + // because we then take the log of this value, and log(0) is inf. // // We don't have a frac(x) primitive in XLA and computing it is tricky, but // because abs(sin(pi * x)) = abs(sin(pi * abs(x))), it's good enough for // our purposes to use abs(frac(x)) = abs(x) - floor(abs(x)). // + // Furthermore, pi * abs(frac(x)) loses precision when abs(frac(x)) is close + // to 1. To remedy this, we can use the fact that sin(pi * x) in the domain + // [0, 1] is symmetric across the line Y=0.5. + // XlaOp abs_input = Abs(input); - XlaOp reflection_denom = Log(Sin(pi * (abs_input - Floor(abs_input)))); + XlaOp abs_frac_input = abs_input - Floor(abs_input); + // Convert values of abs_frac_input > 0.5 to (1 - frac_input) to improve + // precision of pi * abs_frac_input for values of abs_frac_input close to 1. + XlaOp reduced_frac_input = + Select(Gt(abs_frac_input, ScalarLike(abs_frac_input, 0.5)), + ScalarLike(abs_frac_input, 1) - abs_frac_input, abs_frac_input); + XlaOp reflection_denom = Log(Sin(pi * reduced_frac_input)); // Avoid computing -inf - inf, which is nan. If reflection_denom is +/-inf, // then it "wins" and the result is +/-inf. @@ -346,9 +402,16 @@ XlaOp Lgamma(XlaOp input) { // lgamma(+/-inf) = +inf. XlaOp inf_bcast = FullLike(input, std::numeric_limits::infinity()); - return Select(Or(IsFinite(input), // is finite, or - Not(Or(Lt(input, one), Ge(input, one)))), // is nan - result, inf_bcast); + return Select(IsInf(input), inf_bcast, result); + }; + + auto& b = *input.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("Lgamma", input)); + // F16 and BF16 don't provide sufficient precision for intermediate results + // here (although it's better than you might expect!), so do the + // computations in F32. + return DoWithUpcastToF32(input, {BF16, F16}, do_it); }); } @@ -360,10 +423,7 @@ XlaOp Lgamma(XlaOp input) { // A(z) = kBaseLanczosCoeff + sigma(k = 1, n, kLanczosCoefficients[i] / (z + k)) // A'(z) = sigma(k = 1, n, kLanczosCoefficients[i] / (z + k) / (z + k)) XlaOp Digamma(XlaOp input) { - auto& b = *input.builder(); - return b.ReportErrorOrReturn([&]() -> StatusOr { - TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("Digamma", input)); - + auto do_it = [](XlaOp input) { XlaOp zero = ScalarLike(input, 0); XlaOp one_half = ScalarLike(input, 0.5); XlaOp one = ScalarLike(input, 1); @@ -401,8 +461,28 @@ XlaOp Digamma(XlaOp input) { Log1p(z / lanczos_gamma_plus_one_half); XlaOp y = log_t + num / denom - lanczos_gamma / t; - XlaOp reflection = y - pi * Cos(pi * input) / Sin(pi * input); - return Select(need_to_reflect, reflection, y); + + // We need to be careful how we compute cot(pi * input) below: For + // near-integral values of `input`, pi * input can lose precision. + // + // Input is already known to be less than 0.5 (otherwise we don't have to + // reflect). We shift values smaller than -0.5 into the range [-.5, .5] to + // increase precision of pi * input and the resulting cotangent. + XlaOp reduced_input = input + Abs(Floor(input + ScalarLike(input, 0.5))); + XlaOp reflection = + y - pi * Cos(pi * reduced_input) / Sin(pi * reduced_input); + XlaOp real_result = Select(need_to_reflect, reflection, y); + + // Digamma has poles at negative integers and zero; return nan for those. + return Select(And(Le(input, zero), Eq(input, Floor(input))), + FullLike(input, std::numeric_limits::quiet_NaN()), + real_result); + }; + + auto& b = *input.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("Digamma", input)); + return DoWithUpcastToF32(input, {BF16, F16}, do_it); }); } diff --git a/tensorflow/compiler/xla/client/lib/math.h b/tensorflow/compiler/xla/client/lib/math.h index b036fa299d92988439dfecbb5415865071d5577d..71a3acedcec0a8e65561d4139baeaf532ec8bf46 100644 --- a/tensorflow/compiler/xla/client/lib/math.h +++ b/tensorflow/compiler/xla/client/lib/math.h @@ -37,12 +37,6 @@ XlaOp IsNegZero(XlaOp operand); // std::nextafter(from, to) would. XlaOp NextAfter(XlaOp from, XlaOp to); -// Computes the square root of 'operand'. -XlaOp Sqrt(XlaOp operand); - -// Computes the reciprocal of the square root of 'operand'. -XlaOp Rsqrt(XlaOp operand); - // Computes the square of 'operand'. XlaOp Square(XlaOp operand); diff --git a/tensorflow/compiler/xla/client/lib/math_exhaustive_test.cc b/tensorflow/compiler/xla/client/lib/math_exhaustive_test.cc deleted file mode 100644 index f4ed0db56f7026d2c397c5beb1cc7ea3e4b06fee..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/client/lib/math_exhaustive_test.cc +++ /dev/null @@ -1,187 +0,0 @@ -/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/xla/client/lib/math.h" -#include "tensorflow/compiler/xla/client/xla_builder.h" -#include "tensorflow/compiler/xla/literal_util.h" -#include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/client_library_test_base.h" -#include "tensorflow/compiler/xla/tests/test_macros.h" -#include "tensorflow/compiler/xla/types.h" - -namespace xla { -namespace { - -using Eigen::half; - -struct Testcase { - Testcase(string name, XlaOp (*op)(XlaOp), float (*host_op)(float)) - : name(name), op(op), host_op(host_op) {} - - Testcase& set_tolerance(float abs_err, float rel_err) { - error.abs = abs_err; - error.rel = rel_err; - return *this; - } - - Testcase& set_relaxed_nans() { - error.relaxed_nans = true; - return *this; - } - - Testcase& set_fewer_infs_ok() { - error.fewer_infs_ok = true; - return *this; - } - - Testcase& set_skip_pos_inf() { - skip_pos_inf = true; - return *this; - } - - Testcase& set_skip_neg_inf() { - skip_neg_inf = true; - return *this; - } - - Testcase& set_skip_infs() { - skip_pos_inf = true; - skip_neg_inf = true; - return *this; - } - - Testcase& set_skip_neg_zero() { - skip_neg_zero = true; - return *this; - } - - string name; - XlaOp (*op)(XlaOp); - float (*host_op)(float); - - ErrorSpec error{0.01, 0.01}; - - // If true, don't test +/-infinity or negative 0. - bool skip_pos_inf = false; - bool skip_neg_inf = false; - bool skip_neg_zero = false; -}; - -void PrintTo(const Testcase& tc, std::ostream* os) { *os << tc.name; } - -class MathExhaustiveTest : public ClientLibraryTestBase, - public ::testing::WithParamInterface { - public: - MathExhaustiveTest() { - // Disable fast-math, otherwise we get the wrong results for e.g. - // sqrt(-inf). - SetFastMathDisabled(true); - } -}; - -// Checks a function's behavior on all fp16 values. -// -// TODO(jlebar): asin and lgamma tests fail on interpreter. -XLA_TEST_P(MathExhaustiveTest, DISABLED_ON_INTERPRETER(F16)) { - const Testcase& tc = GetParam(); - XlaBuilder b(TestName()); - - std::vector input; - for (uint32 i = 0; i < 1 << 16; ++i) { - half h; - h.x = i; - - // If we're not using infinity as an input, use 0 as a placeholder rather - // than simply skipping this element. We do this because when the test - // framework reports an incorrect answer, it tells us which index failed. - // So long as our inputs are a simple list of all possible float16s, we can - // convert an index to a half with e.g. the following Python: - // - // np.frombuffer(array('H', [12345]), dtype=np.float16)[0] - // - // but as soon as our list of inputs has any gaps, this doesn't work. - if (std::isinf(static_cast(h)) && - ((tc.skip_pos_inf && h > half{0}) || - (tc.skip_neg_inf && h < half{0}))) { - h = half{0}; - } - - if (h == half{0} && tc.skip_neg_zero && - std::signbit(static_cast(h))) { - h = half{0}; - } - - input.push_back(h); - } - - std::vector expected_result; - for (const auto& h : input) { - expected_result.push_back( - static_cast(tc.host_op(static_cast(h)))); - } - - XlaOp param = AddParam(LiteralUtil::CreateR1(input), &b); - tc.op(param); - ComputeAndCompareR1(&b, expected_result, {}, tc.error); -} - -// TODO(b/123355973): The following tests from math.cc are missing. -// -// - Many failures. -// -// Testcase{"acosh", Acosh, std::acosh}.set_relaxed_nans(), -// Testcase{"asinh", Asinh, std::asinh}, -// Testcase{"sinh", Sinh, std::sinh}, -// Testcase{"cosh", Cosh, std::cosh}.set_fewer_infs_ok(), -// Testcase{"round_to_even", RoundToEven, -// [](float x) { return std::nearbyint(x / 2) * 2; }}, -// -// - No equivalent std function to compare with. -// -// Testcase{"erfinv", ErfInv, std::erfinv}, -// Testcase{"digamma", Digamma, std::digamma}, -// -// - Needs a special test (function takes two args, and simply computing in f32 -// and downcasting to f16 doesn't give the correct answer). -// -// Testcase{"nextafter", NextAfter, std::nextafter}, -// -// TODO(b/123355973): Test math functions not from math.cc (e.g. log). -// TODO(b/123355973): Test bf16 and f32. -// TODO(b/123355973): Get rid of skip_infs / skip_neg_zero below if possible. -// TODO(b/123355973): Reduce lgamma error if possible; it is very high. -INSTANTIATE_TEST_CASE_P( - MathExhaustiveTest_Instantiation, MathExhaustiveTest, - ::testing::ValuesIn(std::vector{ - Testcase{"sqrt", Sqrt, std::sqrt}.set_skip_neg_inf(), - Testcase{"rsqrt", Rsqrt, [](float x) { return 1 / std::sqrt(x); }} - .set_tolerance(0.05, 0.05) - .set_skip_infs() - .set_skip_neg_zero(), - Testcase{"square", Square, [](float x) { return x * x; }}, - Testcase{"reciprocal", Reciprocal, [](float x) { return 1 / x; }}, - Testcase{"erf", Erf, std::erf}.set_tolerance(0.001, 0.0001), - Testcase{"erfc", Erfc, std::erfc}.set_tolerance(0.001, 0.0001), - Testcase{"lgamma", Lgamma, std::lgamma} - .set_tolerance(0.1, 0.15) - .set_fewer_infs_ok(), - Testcase{"asin", Asin, std::asin}.set_skip_infs(), - Testcase{"acos", Acos, std::acos}.set_skip_infs(), - Testcase{"atan", Atan, std::atan}, - Testcase{"tan", Tan, std::tan}.set_tolerance(0.05, 0.05), - })); - -} // namespace -} // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/math_test.cc b/tensorflow/compiler/xla/client/lib/math_test.cc index bdfb0575f573716b54cf9116d155d8a3a55056e8..50613ce50255b8e211f6e64afbe0add290dfc647 100644 --- a/tensorflow/compiler/xla/client/lib/math_test.cc +++ b/tensorflow/compiler/xla/client/lib/math_test.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/client/lib/math.h" +#include "tensorflow/compiler/xla/client/lib/constants.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/primitive_util.h" @@ -104,6 +105,49 @@ class MathTypedTest : public MathTest { {true, false, false, false, false, false, false}), {}, error_spec_); } + + // sqrt(x) == pow(x, 0.5) except that + // + // pow(-inf, 0.5) == inf, while + // sqrt(-inf) == nan. + // + // Check that none of our backends are incorrectly assuming that sqrt(x) == + // pow(x, 0.5) without checking this edge case. + // + // For good measure, we also check pow with an exponent other than 0.5. + void TestSqrtPowInequivalence() { + SetFastMathDisabled(true); + + // Tests disable constant folding by default, but this test needs it + // enabled, otherwise we don't tickle the bug we're trying to catch. + // Specifically, without constant folding, the constants we pass to Pow + // below are hidden behind a reshape that's never folded away! + mutable_debug_options()->clear_xla_disable_hlo_passes(); + + const T inf(std::numeric_limits::infinity()); + const T nan(std::numeric_limits::quiet_NaN()); + + XlaBuilder b(TestName()); + auto x = AddParam(LiteralUtil::CreateR1({-inf}), &b); + ConcatInDim( + &b, {Sqrt(x), Pow(x, ScalarLike(x, 0.5)), Pow(x, ScalarLike(x, 0.3))}, + 0); + std::vector expected = {nan, inf, inf}; + ComputeAndCompareR1(&b, expected, {}, error_spec_); + } + + void TestErfEdgeCases() { + SetFastMathDisabled(true); + + XlaBuilder b(TestName()); + auto x = AddParam(LiteralUtil::CreateR1({T{-1}, T{1}, T{0}}), &b); + ErfInv(x); + + const T inf(std::numeric_limits::infinity()); + std::vector expected = {-inf, inf, T{0}}; + + ComputeAndCompareR1(&b, expected, {}, error_spec_); + } }; // TODO(b/123355973): Add bfloat16 to TestTypes once it's working. @@ -119,6 +163,10 @@ XLA_TYPED_TEST(MathTypedTest, LogEdgeCases) { this->TestLogEdgeCases(); } XLA_TYPED_TEST(MathTypedTest, Log1pEdgeCases) { this->TestLog1pEdgeCases(); } XLA_TYPED_TEST(MathTypedTest, IsInfOrNan) { this->TestIsInfOrNan(); } XLA_TYPED_TEST(MathTypedTest, IsNegZero) { this->TestIsNegZero(); } +XLA_TYPED_TEST(MathTypedTest, SqrtPowInequivalence) { + this->TestSqrtPowInequivalence(); +} +XLA_TYPED_TEST(MathTypedTest, ErfInvEdgeCases) { this->TestErfEdgeCases(); } // Check that certain ops only support real, floating-point inputs. // @@ -239,6 +287,7 @@ XLA_TEST_F(MathTest, Lgamma) { ComputeAndCompareR1(&builder, expected, {}, error_spec_); } +#if !defined(XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16) XLA_TEST_F(MathTest, LgammaF16) { SetFastMathDisabled(true); @@ -259,6 +308,7 @@ XLA_TEST_F(MathTest, LgammaF16) { }; ComputeAndCompareR1(&b, expected, {}, ErrorSpec{0.1}); } +#endif XLA_TEST_F(MathTest, Digamma) { XlaBuilder builder(TestName()); diff --git a/tensorflow/compiler/xla/client/lib/slicing.cc b/tensorflow/compiler/xla/client/lib/slicing.cc index d7b33c5af25606c4e7e443027b913f7ca13a013c..0878cbeaf9ae1d85051ea3b5844f5837286c7dc2 100644 --- a/tensorflow/compiler/xla/client/lib/slicing.cc +++ b/tensorflow/compiler/xla/client/lib/slicing.cc @@ -161,4 +161,24 @@ XlaOp TorchGather(XlaOp input, XlaOp index, int64 dim) { }); } +XlaOp TorchIndexSelect(XlaOp input, XlaOp index, int64 dim) { + XlaBuilder* builder = input.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape input_shape, builder->GetShape(input)); + TF_ASSIGN_OR_RETURN(Shape index_shape, builder->GetShape(index)); + std::vector slice_sizes = input_shape.dimensions(); + slice_sizes[dim] = 1; + GatherDimensionNumbers gather_dnums; + for (int64 i = 0; i < input_shape.rank(); ++i) { + if (i != dim) { + gather_dnums.add_offset_dims(i); + } + } + gather_dnums.set_index_vector_dim(index_shape.rank()); + gather_dnums.add_collapsed_slice_dims(dim); + gather_dnums.add_start_index_map(dim); + return Gather(input, index, gather_dnums, slice_sizes); + }); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/slicing.h b/tensorflow/compiler/xla/client/lib/slicing.h index 69f98a6f43fa167adf6f77b28645a3460b292633..bb6191df7c442f23a63f0d0b80c9b534c31e30fc 100644 --- a/tensorflow/compiler/xla/client/lib/slicing.h +++ b/tensorflow/compiler/xla/client/lib/slicing.h @@ -57,6 +57,14 @@ XlaOp DynamicUpdateSliceInMinorDims(XlaOp x, XlaOp update, // `index`. XlaOp TorchGather(XlaOp input, XlaOp index, int64 dim); +// Returns a new tensor which indexes the input tensor along dimension dim using +// the entries in index. +// +// The returned tensor has the same number of dimensions as the original tensor +// (input). The dimth dimension has the same size as the length of index; other +// dimensions have the same size as in the original tensor. +XlaOp TorchIndexSelect(XlaOp input, XlaOp index, int64 dim); + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_SLICING_H_ diff --git a/tensorflow/compiler/xla/client/lib/slicing_test.cc b/tensorflow/compiler/xla/client/lib/slicing_test.cc index db6ebb9df18372260a64a3e9fd17b0c30b35667d..408a82ca3c6eeeae7edac8511769ec9c0d5a5f44 100644 --- a/tensorflow/compiler/xla/client/lib/slicing_test.cc +++ b/tensorflow/compiler/xla/client/lib/slicing_test.cc @@ -115,5 +115,43 @@ XLA_TEST_F(SlicingTest, TorchGather) { ComputeAndCompareR2(&builder, {{1, 1}, {4, 3}}, {input_data.get(), index_data.get()}); } + +XLA_TEST_F(SlicingTest, TorchIndexSelectOn0) { + xla::XlaBuilder builder(TestName()); + + xla::XlaOp input, index; + auto input_data = + CreateR2Parameter({{0.1427, 0.0231, -0.5414, -1.0009}, + {-0.4664, 0.2647, -0.1228, -1.1068}, + {-1.1734, -0.6571, 0.7230, -0.6004}}, + 0, "input", &builder, &input); + auto index_data = + CreateR1Parameter({0, 2}, 1, "index", &builder, &index); + TorchIndexSelect(input, index, 0); + + ComputeAndCompareR2( + &builder, + {{0.1427, 0.0231, -0.5414, -1.0009}, {-1.1734, -0.6571, 0.7230, -0.6004}}, + {input_data.get(), index_data.get()}); +} + +XLA_TEST_F(SlicingTest, TorchIndexSelectOn1) { + xla::XlaBuilder builder(TestName()); + + xla::XlaOp input, index; + auto input_data = + CreateR2Parameter({{0.1427, 0.0231, -0.5414, -1.0009}, + {-0.4664, 0.2647, -0.1228, -1.1068}, + {-1.1734, -0.6571, 0.7230, -0.6004}}, + 0, "input", &builder, &input); + auto index_data = + CreateR1Parameter({0, 2}, 1, "index", &builder, &index); + TorchIndexSelect(input, index, 1); + + ComputeAndCompareR2( + &builder, {{0.1427, -0.5414}, {-0.4664, -0.1228}, {-1.1734, 0.7230}}, + {input_data.get(), index_data.get()}); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/client/xla_builder.cc b/tensorflow/compiler/xla/client/xla_builder.cc index b371b5af37b3b1bf911133a485554f87c8e09183..9d856c2996f421d51238fd2f1e18921ff3f1f5d0 100644 --- a/tensorflow/compiler/xla/client/xla_builder.cc +++ b/tensorflow/compiler/xla/client/xla_builder.cc @@ -299,12 +299,17 @@ XlaComputation XlaBuilder::BuildAndNoteError() { return build_status.ConsumeValueOrDie(); } -StatusOr XlaBuilder::Build(bool remove_dynamic_dimensions) { +Status XlaBuilder::GetCurrentStatus() const { if (!first_error_.ok()) { string backtrace; first_error_backtrace_.Dump(tensorflow::DebugWriteToString, &backtrace); return AppendStatus(first_error_, backtrace); } + return Status::OK(); +} + +StatusOr XlaBuilder::Build(bool remove_dynamic_dimensions) { + TF_RETURN_IF_ERROR(GetCurrentStatus()); return Build(instructions_.back().id(), remove_dynamic_dimensions); } @@ -318,11 +323,7 @@ StatusOr XlaBuilder::Build(XlaOp root, StatusOr XlaBuilder::Build(int64 root_id, bool remove_dynamic_dimensions) { - if (!first_error_.ok()) { - string backtrace; - first_error_backtrace_.Dump(tensorflow::DebugWriteToString, &backtrace); - return AppendStatus(first_error_, backtrace); - } + TF_RETURN_IF_ERROR(GetCurrentStatus()); // TODO(b/121223198): XLA backend cannot handle dynamic dimensions yet, remove // all dynamic dimensions before building xla program until we have support in @@ -1879,32 +1880,46 @@ XlaOp XlaBuilder::Conditional(const XlaOp& predicate, const XlaOp& true_operand, const XlaComputation& true_computation, const XlaOp& false_operand, const XlaComputation& false_computation) { + // The index of true_computation must be 0 and that of false computation + // must be 1. + return Conditional(predicate, {&true_computation, &false_computation}, + {true_operand, false_operand}); +} + +XlaOp XlaBuilder::Conditional( + const XlaOp& branch_index, + absl::Span branch_computations, + absl::Span branch_operands) { return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; - TF_ASSIGN_OR_RETURN(const Shape& predicate_shape, GetShape(predicate)); - TF_ASSIGN_OR_RETURN(const Shape& true_operand_shape, - GetShape(true_operand)); - TF_ASSIGN_OR_RETURN(const ProgramShape& true_computation_shape, - true_computation.GetProgramShape()); - TF_ASSIGN_OR_RETURN(const Shape& false_operand_shape, - GetShape(false_operand)); - TF_ASSIGN_OR_RETURN(const ProgramShape& false_computation_shape, - false_computation.GetProgramShape()); - TF_ASSIGN_OR_RETURN( - Shape shape, - ShapeInference::InferConditionalShape( - predicate_shape, true_operand_shape, false_operand_shape, - true_computation_shape, false_computation_shape)); + TF_ASSIGN_OR_RETURN(const Shape& branch_index_shape, + GetShape(branch_index)); + std::vector branch_operand_shapes(branch_operands.size()); + std::vector branch_computation_shapes( + branch_computations.size()); + for (int j = 0; j < branch_operands.size(); ++j) { + TF_ASSIGN_OR_RETURN(branch_operand_shapes[j], + GetShape(branch_operands[j])); + TF_ASSIGN_OR_RETURN(branch_computation_shapes[j], + branch_computations[j]->GetProgramShape()); + } + TF_ASSIGN_OR_RETURN(const Shape shape, + ShapeInference::InferConditionalShape( + branch_index_shape, branch_computation_shapes, + branch_operand_shapes)); *instr.mutable_shape() = shape.ToProto(); - // The index of true_computation must be 0 and that of false computation - // must be 1. - AddCalledComputation(true_computation, &instr); - AddCalledComputation(false_computation, &instr); + for (const XlaComputation* branch_computation : branch_computations) { + AddCalledComputation(*branch_computation, &instr); + } + std::vector operands(1, branch_index); + for (const XlaOp branch_operand : branch_operands) { + operands.emplace_back(branch_operand); + } return AddInstruction(std::move(instr), HloOpcode::kConditional, - {predicate, true_operand, false_operand}); + absl::MakeSpan(operands)); }); } @@ -3021,6 +3036,21 @@ XlaOp TriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower, }); } +XlaOp Cholesky(XlaOp a, bool lower) { + XlaBuilder* builder = a.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + TF_ASSIGN_OR_RETURN(const Shape& a_shape, builder->GetShape(a)); + xla::CholeskyOptions& options = *instr.mutable_cholesky_options(); + options.set_lower(lower); + TF_ASSIGN_OR_RETURN(Shape shape, + ShapeInference::InferCholeskyShape(a_shape)); + *instr.mutable_shape() = shape.ToProto(); + + return builder->AddInstruction(std::move(instr), HloOpcode::kCholesky, {a}); + }); +} + XlaOp Infeed(XlaBuilder* builder, const Shape& shape, const string& config) { return builder->Infeed(shape, config); } @@ -3286,6 +3316,12 @@ XlaOp Real(const XlaOp& operand) { XlaOp Imag(const XlaOp& operand) { return operand.builder()->UnaryOp(HloOpcode::kImag, operand); } +XlaOp Sqrt(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kSqrt, operand); +} +XlaOp Rsqrt(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kRsqrt, operand); +} XlaOp Pow(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { @@ -3359,6 +3395,13 @@ XlaOp Conditional(const XlaOp& predicate, const XlaOp& true_operand, false_computation); } +XlaOp Conditional(const XlaOp& branch_index, + absl::Span branch_computations, + absl::Span branch_operands) { + return branch_index.builder()->Conditional(branch_index, branch_computations, + branch_operands); +} + XlaOp ReducePrecision(const XlaOp& operand, const int exponent_bits, const int mantissa_bits) { return operand.builder()->ReducePrecision(operand, exponent_bits, diff --git a/tensorflow/compiler/xla/client/xla_builder.h b/tensorflow/compiler/xla/client/xla_builder.h index fd2e9816e8a0b755b0a1060e8ed4e30c317bd208..56e85e394c515b87de12f915202db0c7a9ccd867 100644 --- a/tensorflow/compiler/xla/client/xla_builder.h +++ b/tensorflow/compiler/xla/client/xla_builder.h @@ -238,6 +238,10 @@ class XlaBuilder { // See also set_die_immediately_on_error(). Status first_error() const { return first_error_; } + // Returns the current status of the builder, complete with the stack trace + // information. + Status GetCurrentStatus() const; + // Returns the shape of the given op. StatusOr GetShape(const XlaOp& op) const; @@ -299,14 +303,17 @@ class XlaBuilder { input_output_aliases_.push_back({output_index, param_number, param_index}); } - private: // Describes an input/output alias as inserted by the SetUpAlias() API. struct InputOutputAlias { + // Specifies the index of the aliased buffer in the result tuple. ShapeIndex output_index; + // Specifies the parameter containing the buffer to be aliased. int64 param_number; + // Specifies the index of the aliased buffer in the parameter ShapeIndex param_index; }; + private: // Build helper which takes the id of the root operation.. StatusOr Build(int64 root_id, bool remove_dynamic_dimensions); @@ -525,6 +532,10 @@ class XlaBuilder { const XlaOp& false_operand, const XlaComputation& false_computation); + XlaOp Conditional(const XlaOp& branch_index, + absl::Span branch_computations, + absl::Span branch_operands); + XlaOp ReducePrecision(const XlaOp& operand, const int exponent_bits, const int mantissa_bits); @@ -797,6 +808,7 @@ class XlaBuilder { friend XlaOp TriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower, bool unit_diagonal, TriangularSolveOptions::Transpose transpose_a); + friend XlaOp Cholesky(XlaOp a, bool lower); friend XlaOp Infeed(XlaBuilder* builder, const Shape& shape, const string& config); friend void Outfeed(const XlaOp& operand, const Shape& shape_with_layout, @@ -906,6 +918,8 @@ class XlaBuilder { friend XlaOp Tanh(const XlaOp& operand); friend XlaOp Real(const XlaOp& operand); friend XlaOp Imag(const XlaOp& operand); + friend XlaOp Sqrt(const XlaOp& operand); + friend XlaOp Rsqrt(const XlaOp& operand); friend XlaOp Pow(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions); friend XlaOp IsFinite(const XlaOp& operand); @@ -939,6 +953,10 @@ class XlaBuilder { const XlaComputation& true_computation, const XlaOp& false_operand, const XlaComputation& false_computation); + friend XlaOp Conditional( + const XlaOp& branch_index, + absl::Span branch_computations, + absl::Span branch_operands); friend XlaOp ReducePrecision(const XlaOp& operand, const int exponent_bits, const int mantissa_bits); friend XlaOp Gather(const XlaOp& input, const XlaOp& start_indices, @@ -1336,8 +1354,7 @@ XlaOp Fft(const XlaOp& operand, FftType fft_type, // * `left_side` is a boolean, indicating whether to solve a system of the form // op(a) * x = b (true) or x * op(a) = b (false). // * `lower` is a boolean, indicating whether the argument `a` is -// lower-triangular -// (true) or upper-triangular (false). +// lower-triangular (true) or upper-triangular (false). // * If `unit_diagonal` is true, the diagonal elements of `a` are assumed to be // 1 and not accessed. // * `transpose_a` indicates which function `op` we use to transform the tensor @@ -1346,6 +1363,20 @@ XlaOp TriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower, bool unit_diagonal, TriangularSolveOptions::Transpose transpose_a); +// Computes the Cholesky decompositions of a batch of symmetric (Hermitian) +// positive definite matrices. +// `a` must be a (batched) square matrix; i.e., it must have rank >= 2 with the +// two minor dimensions equal. +// If `lower` is true, the data from the lower triangle is used; if false, the +// upper triangle is used. The input data in the other triangle of the input +// does not affect the output. Returns the output in the same lower/uppper +// triangle. The data returned in the other output triangle is arbitrary and +// implementation-defined. +// +// The value returned if `a` is not Hermitian positive definite is +// implementation-defined. +XlaOp Cholesky(XlaOp a, bool lower); + // Enqueues an infeed instruction onto the computation, which writes data of // the given shape to the infeed buffer of the device. XlaOp Infeed(XlaBuilder* builder, const Shape& shape, @@ -1634,6 +1665,12 @@ XlaOp Real(const XlaOp& operand); // Enqueues an imaginary-part instruction onto the computation. XlaOp Imag(const XlaOp& operand); +// Enqueues a sqrt computation onto the computation. +XlaOp Sqrt(const XlaOp& operand); + +// Enqueues a rsqrt computation onto the computation. +XlaOp Rsqrt(const XlaOp& operand); + // Enqueues a lhs^rhs computation onto the computation. XlaOp Pow(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); @@ -1715,9 +1752,11 @@ XlaOp Sort(const XlaOp& keys, absl::Span values = {}, // * All operands must be tensors with the same dimensions. The element types of // the tensors may be different. // * The result is a tuple that consists of the operands in sorted order (along -// the provided dimension, as above). When comparing, 'comparator' is called -// with 2 * n scalar parameters, where parameter 2 * i and 2 * i + 1 -// correspond to the value of operand i at two index positions. +// the provided dimension, as above). The same permutation as implied by the +// comparison computation is applied to all operand tensors. When comparing +// two index positions, 'comparator' is called with 2 * n scalar parameters, +// where parameter 2 * i and 2 * i + 1 correspond to the value of operand i at +// two index positions. // Default comparator computations can be found in lib/comparators.h XlaOp Sort(absl::Span operands, const XlaComputation& comparator, int64 dimension = -1, bool is_stable = false); @@ -1748,6 +1787,15 @@ XlaOp Conditional(const XlaOp& predicate, const XlaOp& true_operand, const XlaOp& false_operand, const XlaComputation& false_computation); +// Enqueues either a predicated (if/else) or indexed (switch/case/default) +// conditional node onto the computation. N >= 1 branch_computations and +// branch_operands are matched by index. branch_index selects the branch that +// will be executed. Out of range branch_index uses the N-1'th +// branch_computation as default. +XlaOp Conditional(const XlaOp& branch_index, + absl::Span branch_computations, + absl::Span branch_operands); + // Enqueues a ReducePrecision node onto the computation. XlaOp ReducePrecision(const XlaOp& operand, const int exponent_bits, const int mantissa_bits); diff --git a/tensorflow/compiler/xla/g3doc/_project.yaml b/tensorflow/compiler/xla/g3doc/_project.yaml index 33d8bdb27a664d9e282d1d65c007ebf5838b196a..1cacee703dca30f9c4af6a4964839bb9fa4b0140 100644 --- a/tensorflow/compiler/xla/g3doc/_project.yaml +++ b/tensorflow/compiler/xla/g3doc/_project.yaml @@ -8,3 +8,4 @@ use_site_branding: true hide_from_products_list: true content_license: cc3-apache2 buganizer_id: 171704 +include: /_project_included.yaml diff --git a/tensorflow/compiler/xla/g3doc/operation_semantics.md b/tensorflow/compiler/xla/g3doc/operation_semantics.md index db90d184b5218614ac49363ebf2a7e25fffe44de..b3fdd36b113ca879befbbb7b954b81c0e18456a4 100644 --- a/tensorflow/compiler/xla/g3doc/operation_semantics.md +++ b/tensorflow/compiler/xla/g3doc/operation_semantics.md @@ -322,6 +322,37 @@ Invokes a computation with the given arguments. The arity and types of the `args` must match the parameters of the `computation`. It is allowed to have no `args`. +## Cholesky + +See also +[`XlaBuilder::Cholesky`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Computes the +[Cholesky decomposition](https://en.wikipedia.org/wiki/Cholesky_decomposition) +of a batch of symmetric (Hermitian) positive definite matrices. + + `Cholesky(a, lower)` + +Arguments | Type | Semantics +--------- | ------- | ----------------------------------------------------- +`a` | `XlaOp` | a rank > 2 array of a complex or floating-point type. +`lower` | `bool` | whether to use the upper or lower triangle of `a`. + +If `lower` is `true`, computes lower-triangular matrices `l` such that $$ a = l +. l^T $$. If `lower` is `false`, computes upper-triangular matrices `u` such +that $$ a = u^T . u $$. + +Input data is read only from the lower/upper triangle of `a`, depending on the +value of `lower`. Values from the other triangle are ignored. Output data is +returned in the same triangle; the values in the other triangle are +implementation-defined and may be anything. + +If the rank of `a` is greater than 2, `a` is treated as a batch of matrices, +where all except the minor 2 dimensions are batch dimensions. + +If `a` is not symmetric (Hermitian) positive definite, the result is +implementation-defined. + ## Clamp See also @@ -510,25 +541,49 @@ See also false_computation)` Arguments | Type | Semantics -------------------- | ---------------- | --------------------------------- +------------------- | ---------------- | -------------------------------------- `pred` | `XlaOp` | Scalar of type `PRED` -`true_operand` | `XlaOp` | Argument of type `T_0` -`true_computation` | `XlaComputation` | XlaComputation of type `T_0 -> S` -`false_operand` | `XlaOp` | Argument of type `T_1` -`false_computation` | `XlaComputation` | XlaComputation of type `T_1 -> S` +`true_operand` | `XlaOp` | Argument of type $$ T_0 $$ +`true_computation` | `XlaComputation` | XlaComputation of type $$ T_0 \to S$$ +`false_operand` | `XlaOp` | Argument of type $$ T_1 $$ +`false_computation` | `XlaComputation` | XlaComputation of type $$ T_1 \to S $$ Executes `true_computation` if `pred` is `true`, `false_computation` if `pred` is `false`, and returns the result. -The `true_computation` must take in a single argument of type `T_0` and will be -invoked with `true_operand` which must be of the same type. The -`false_computation` must take in a single argument of type `T_1` and will be +The `true_computation` must take in a single argument of type $$ T_0 $$ and will +be invoked with `true_operand` which must be of the same type. The +`false_computation` must take in a single argument of type $$ T_1 $$ and will be invoked with `false_operand` which must be of the same type. The type of the returned value of `true_computation` and `false_computation` must be the same. Note that only one of `true_computation` and `false_computation` will be executed depending on the value of `pred`. + `Conditional(branch_index, branch_computations, branch_operands)` + +| Arguments | Type | Semantics | +| --------------------- | --------------------- | ---------------------------- | +| `branch_index` | `XlaOp` | Scalar of type `PRED` or | +: : : `S32` : +| `branch_computations` | sequence of N | XlaComputations of type $$ | +: : `XlaComputation` : T_0 \to S , T_1 \to S , ..., : +: : : T_{N-1} \to S $$ : +| `branch_operands` | sequence of N `XlaOp` | Arguments of type $$ T_0 , | +: : : T_1 , ..., T_{N-1} $$ : + +Executes `branch_computations[branch_index]`, and returns the result. If +`branch_index` is a `PRED`, then the `true` branch is in position 0 and the +`false` branch is in position 1. If `branch_index` is an `S32` which is < 0 +or >= N, then `branch_computations[N-1]` is executed as the default branch. + +Each `branch_computations[b]` must take in a single argument of type `T_b` and +will be invoked with `branch_operands[b]` which must be of the same type. The +type of the returned value of each `branch_computations[b]` must be the same. + +Note that only one of the `branch_computations` will be executed depending on +the value of `branch_index`. + ## Conv (convolution) See also @@ -2439,6 +2494,46 @@ Permutes the operand dimensions with the given permutation, so This is the same as Reshape(operand, permutation, Permute(permutation, operand.shape.dimensions)). +## TriangularSolve + +See also +[`XlaBuilder::TriangularSolve`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Solves systems of linear equations with lower or upper triangular coefficient +matrices by forward- or back-substitution. Broadcasting along leading +dimensions, this routine solves one of the matrix systems `op(a) * x = +b`, or `x * op(a) = b`, for the variable `x`, given `a` and `b`, where `op(a)` is +either `op(a) = a`, or `op(a) = Transpose(a)`, or `op(a) = Conj(Transpose(a))`. + + `TriangularSolve(a, b, left_side, lower, unit_diagonal, transpose_a)` + +| Arguments | Type | Semantics | +| --------------- | ----------- | -------------------------------------------- | +| `a` | `XlaOp` | a rank > 2 array of a complex or | +: : : floating-point type with shape `[..., M, : +: : : M]`. : +| `b` | `XlaOp` | a rank > 2 array of the same type with shape | +: : : `[..., M, K]` if `left_side` is true, `[..., : +: : : K, M]` otherwise. : +| `left_side` | `bool` | indicates whether to solve a system of the | +: : : form `op(a) * x = b` (`true`) or `x * : +: : : op(a) = b` (`false`). : +| `lower` | `bool` | whether to use the upper or lower triangle | +: : : of `a`. : +| `unit_diagonal` | `bool` | if `true`, the diagonal elements of `a` are | +: : : assumed to be `1` and not accessed. : +| `transpose_a` | `Transpose` | whether to use `a` as is, transpose it or | +: : : take its conjugate transpose. : + +Input data is read only from the lower/upper triangle of `a`, depending on the +value of `lower`. Values from the other triangle are ignored. Output data is +returned in the same triangle; the values in the other triangle are +implementation-defined and may be anything. + +If the rank of `a` and `b` are greater than 2, they are treated as batches of +matrices, where all except the minor 2 dimensions are batch dimensions. `a` and +`b` must have equal batch dimensions. + ## Tuple See also diff --git a/tensorflow/compiler/xla/python/BUILD b/tensorflow/compiler/xla/python/BUILD index 55eacc1c16a76522215d27ac7cf4e801e69c9740..68128a3097f478084b380591ccc3fbd69ca8351b 100644 --- a/tensorflow/compiler/xla/python/BUILD +++ b/tensorflow/compiler/xla/python/BUILD @@ -13,8 +13,6 @@ py_library( visibility = ["//visibility:public"], deps = [ ":pywrap_xla", - "//tensorflow/compiler/xla:xla_data_proto_py", - "//tensorflow/compiler/xla/service:hlo_proto_py", ], ) @@ -33,6 +31,7 @@ py_test( deps = [ ":custom_call_for_test", ":xla_client", + "//tensorflow/compiler/xla:xla_data_proto_py", "//tensorflow/python:platform_test", ], ) @@ -70,7 +69,6 @@ cc_library( "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:xla_builder", "//tensorflow/compiler/xla/client:xla_computation", - "//tensorflow/compiler/xla/client/lib:cholesky", "//tensorflow/compiler/xla/client/lib:math", "//tensorflow/compiler/xla/client/lib:qr", "//tensorflow/compiler/xla/service:computation_placer", diff --git a/tensorflow/compiler/xla/python/local_computation_builder.cc b/tensorflow/compiler/xla/python/local_computation_builder.cc index c14a01a858af414fc78a5f727372e8fa64cad4b8..5cfbb2c20df726529d555cb46d68c4a52b2067e1 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.cc +++ b/tensorflow/compiler/xla/python/local_computation_builder.cc @@ -21,7 +21,6 @@ limitations under the License. #include "absl/memory/memory.h" #include "tensorflow/compiler/xla/client/client_library.h" -#include "tensorflow/compiler/xla/client/lib/cholesky.h" #include "tensorflow/compiler/xla/client/lib/math.h" #include "tensorflow/compiler/xla/client/lib/qr.h" #include "tensorflow/compiler/xla/client/xla_builder.h" @@ -443,6 +442,8 @@ StatusOr ComputationBuilder::GetReturnValueShape() { return program_shape.result(); } +LocalOp ComputationBuilder::ReplicaId() { return xla::ReplicaId(&builder_); } + LocalOp ComputationBuilder::Infeed(const Shape& shape) { return xla::Infeed(&builder_, shape); } @@ -496,7 +497,8 @@ LocalOp ComputationBuilder::Collapse(const LocalOp& operand, LocalOp ComputationBuilder::AllToAll( const LocalOp& operand, int64 split_dimension, int64 concat_dimension, int64 split_count, absl::Span replica_groups) { - std::vector rg(replica_groups.size()); + std::vector rg; + rg.reserve(replica_groups.size()); for (int i = 0; i < replica_groups.size(); ++i) { rg.push_back(replica_groups[i]); } @@ -711,8 +713,8 @@ LocalOp ComputationBuilder::SortKeyVal(const LocalOp& keys, return xla::Sort(keys.op(), {values.op()}, dimension); } -LocalOp ComputationBuilder::Cholesky(const LocalOp& a) { - return xla::Cholesky(a.op()); +LocalOp ComputationBuilder::Cholesky(const LocalOp& a, bool lower) { + return xla::Cholesky(a.op(), lower); } LocalOp ComputationBuilder::QR(const LocalOp& a, bool full_matrices) { diff --git a/tensorflow/compiler/xla/python/local_computation_builder.h b/tensorflow/compiler/xla/python/local_computation_builder.h index 66b1cce7fb598388af40940ea2ed52ac2f8ee8e1..fa878501aba729c65206ce7ba2e2ff56dce53bd8 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.h +++ b/tensorflow/compiler/xla/python/local_computation_builder.h @@ -87,7 +87,7 @@ class LocalShapedBuffer { // analogous to std::unique_ptr::release(). ShapedBuffer Release(); - // Destructures a tuple-valued LocalShapedBuffer into its constitutent + // Destructures a tuple-valued LocalShapedBuffer into its constituent // elements in LocalShapedBufferTuple form. StatusOr DestructureTuple(); @@ -224,6 +224,8 @@ class ComputationBuilder { // Returns the shape of the current return value for the computation. StatusOr GetReturnValueShape(); + LocalOp ReplicaId(); + LocalOp Infeed(const Shape& shape); void Outfeed(const LocalOp& operand, const Shape& shape, @@ -356,7 +358,7 @@ class ComputationBuilder { LocalOp QR(const LocalOp& a, bool full_matrices); - LocalOp Cholesky(const LocalOp& a); + LocalOp Cholesky(const LocalOp& a, bool lower); // `transpose_a` is the integer value of a TriangularSolveOptions::Transpose // enum. We use an integer here so we don't have to teach SWIG about the diff --git a/tensorflow/compiler/xla/python/local_computation_builder.i b/tensorflow/compiler/xla/python/local_computation_builder.i index 7d7a860baa03e99cc254b7596fb5f9d41acbef20..9fcb4822c7feebf329ab367d09acac5c32eab4cd 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.i +++ b/tensorflow/compiler/xla/python/local_computation_builder.i @@ -16,92 +16,6 @@ limitations under the License. // SWIG typemaps and declarations for building, compiling, and // executing XLA computations, wrapping most of what is declared in // local_computation_builder.h. -// -// The typemaps below implement/assert the following correspondences -// (with elaborations below): -// -// C++ Python -// -------------------------------------+--------------------------------------- -// Span <- sequence of int -// vector -> sequence of int -// Span <- sequence of LocalOp -// Literal <-> (nested tuple of) numpy ndarray -// std::vector <- sequence of (nested tuple of) ndarray -// Shape -> pair holding (dtype, dimensions) -// <- object duck-typed as xla_client.Shape -// ProgramShape -> pair of ([arg_shapes], ret_shape) -// std::vector <- sequence of xla_client.Shape objects -// PrimitiveType <- int -// Span> <- sequence of int pairs -// PaddingConfig proto <- corresponding Python proto -// ConvolutionDimensionNumbers proto <- corresponding Python proto -// DotDimensionNumbers proto <- corresponding Python proto -// GatherDimensionNumbers proto <- corresponding Python proto -// ScatterDimensionNumbers proto <- corresponding Python proto -// Span <- sequence of ReplicaGroup Python proto -// -// Arrows indicate whether a conversion only ever occurs in one -// direction, or whether it is maintained bidirectionally. -// -// The Python objects corresponding to C++ Literals have the type: -// -// T = ndarray | (T, ...) -// -// where a terminal numpy ndarray translates to a Literal with a -// non-tuple Shape, an XLA primitive element type corresponding to the -// ndarray's dtype. Meanwhile, a non-terminal "tuple of T" translates -// to a tuple-shaped Literal whose tuple components are translated -// recursively. For example, if x is a numpy ndarray in Python, with -// shape (2, 3) and dtype of dtype('float32'), then x translates to a -// Literal with rank 2, dimension 2 and 3, and XLA primitive type -// F32. Meanwhile, -// -// (x, (x, x), (x,)), -// -// translates to a tuple-shaped XLA Literal, whose component subshapes -// are a 2x3 F32-shaped literal followed by two tuple-shaped literals. -// -// Shapes output by C++ become Python objects with the type: -// -// T = (dtype, S) -// S = DIMENSIONS | TUPLE_SHAPES -// DIMENSIONS = (int, ...) -// TUPLE_SHAPES = (T, ...) -// -// In the pair described by the T rule, the terminal dtype determines -// whether S expands as DIMENSIONS or TUPLE_SHAPES. Namely if it is -// dtype('O'), numpy's object dtype, the structure represents a tuple -// shape and the expansion of the non-terminal S is -// TUPLE_SHAPES. Otherwise, dtype describes a primitive element type -// and S expands into DIMENSIONS giving dimension sizes. For example: -// -// (dtype('float32'), (3, 5, 7)) -// -// describes a 3x5x7 array of F32s, and -// -// (dtype('O'), ((dtype('float32'), (2, 3)), -// (dtype('float64'), (4, 5)))) -// -// describes a tuple shape with two subshapes: the first a 2x3 F32, -// and the other a 4x5 F64. -// -// The Python int corresponding to a PrimitiveType enum must be valid -// per xla_data.proto (e.g. xla_data.PRED, xla_data.F32). -// -// The SWIG object wrappers generated by this file are not intended -// for end use, but rather for internal use in the Python XLA client, -// xla_client.py. -// -// One central reason for the Python-side indirection is that the -// Python-side objects produced by the typemaps in this file are -// further packaged up by xla_client before being passed on. For -// instance, the Python pair produced for a C++ Shape is further -// wrapped in a Python class (xla_client.Shape) so as not to expose -// the raw pair externally. -// -// Other SWIG object wrappers (e.g. of Computation) are further -// wrapped by xla_client in order to set up a custom destructor that -// triggers memory deallocation on the C++ side. %module(threads="1") local_computation_builder @@ -379,6 +293,7 @@ tensorflow::ImportNumpy(); %unignore xla::swig::ComputationBuilder::Parameter; %unignore xla::swig::ComputationBuilder::GetShape; %unignore xla::swig::ComputationBuilder::GetReturnValueShape; +%unignore xla::swig::ComputationBuilder::ReplicaId; %unignore xla::swig::ComputationBuilder::Infeed; %unignore xla::swig::ComputationBuilder::Outfeed; %unignore xla::swig::ComputationBuilder::ConstantLiteral; diff --git a/tensorflow/compiler/xla/python/xla_client.py b/tensorflow/compiler/xla/python/xla_client.py index 9019a979a61c6ebb62adaa5503560c604e2b30f8..38458d7f090f24ea2132298f2f7bd5201efbfed1 100644 --- a/tensorflow/compiler/xla/python/xla_client.py +++ b/tensorflow/compiler/xla/python/xla_client.py @@ -30,9 +30,12 @@ import numpy as np import six from six.moves import xrange -from tensorflow.compiler.xla import xla_data_pb2 +# Note this module does *not* depend on any Python protocol buffers. The XLA +# Python bindings are currently packaged both as part of jaxlib and as part +# of TensorFlow. If we use protocol buffers here, then importing both jaxlib +# and TensorFlow may fail with duplicate protocol buffer message definitions. + from tensorflow.compiler.xla.python import pywrap_xla as c_api -from tensorflow.compiler.xla.service import hlo_pb2 # Import the XRT backend, if available. try: @@ -228,15 +231,6 @@ def BackendSpec(backend, target): raise ValueError('Unknown backend {}'.format(backend)) -def OpMetadataToProto(pyobj): - proto = xla_data_pb2.OpMetadata() - for field in _OP_METADATA_FIELDS: - attr = getattr(pyobj, field) - if attr is not None: - setattr(proto, field, attr) - return proto - - def CurrentSourceInfoMetadata(op_type=None, op_name=None, skip_frames=1): """Helper for use in source mapping that returns an OpMetadata object.""" full_filename, lineno = inspect.stack()[skip_frames][1:3] @@ -276,8 +270,7 @@ def _convert_padding_type_to_pad_values(padding_type, lhs_dims, rhs_dims, pad_sizes = [max((out_size - 1) * stride + filter_size - in_size, 0) for out_size, stride, filter_size, in_size in zip(out_shape, window_strides, rhs_dims, lhs_dims)] - return [(pad_size // 2, pad_size - pad_size // 2) - for pad_size in pad_sizes] + return [(pad_size // 2, pad_size - pad_size // 2) for pad_size in pad_sizes] else: msg = 'Unexpected PaddingType value: {}' raise ValueError(msg.format(padding_type)) @@ -349,29 +342,56 @@ _BINARY_OPS = [ ] +class PrimitiveType(enum.IntEnum): + """Python copy of the XLA PrimitiveType enum. + + Must match the corresponding protocol buffer. + """ + PRIMITIVE_TYPE_INVALID = 0 + PRED = 1 + S8 = 2 + S16 = 3 + S32 = 4 + S64 = 5 + U8 = 6 + U16 = 7 + U32 = 8 + U64 = 9 + BF16 = 16 + F16 = 10 + F32 = 11 + F64 = 12 + C64 = 15 + C128 = 18 + TUPLE = 13 + OPAQUE = 14 + TOKEN = 17 + + XLA_ELEMENT_TYPE_TO_DTYPE = { - xla_data_pb2.PRED: np.dtype('bool'), - xla_data_pb2.S8: np.dtype('int8'), - xla_data_pb2.S16: np.dtype('int16'), - xla_data_pb2.S32: np.dtype('int32'), - xla_data_pb2.S64: np.dtype('int64'), - xla_data_pb2.U8: np.dtype('uint8'), - xla_data_pb2.U16: np.dtype('uint16'), - xla_data_pb2.U32: np.dtype('uint32'), - xla_data_pb2.U64: np.dtype('uint64'), - xla_data_pb2.F16: np.dtype('float16'), - xla_data_pb2.F32: np.dtype('float32'), - xla_data_pb2.F64: np.dtype('float64'), - xla_data_pb2.C64: np.dtype('complex64'), - xla_data_pb2.C128: np.dtype('complex128'), - xla_data_pb2.TUPLE: np.dtype(np.object), + PrimitiveType.PRED: np.dtype('bool'), + PrimitiveType.S8: np.dtype('int8'), + PrimitiveType.S16: np.dtype('int16'), + PrimitiveType.S32: np.dtype('int32'), + PrimitiveType.S64: np.dtype('int64'), + PrimitiveType.U8: np.dtype('uint8'), + PrimitiveType.U16: np.dtype('uint16'), + PrimitiveType.U32: np.dtype('uint32'), + PrimitiveType.U64: np.dtype('uint64'), + PrimitiveType.F16: np.dtype('float16'), + PrimitiveType.F32: np.dtype('float32'), + PrimitiveType.F64: np.dtype('float64'), + PrimitiveType.C64: np.dtype('complex64'), + PrimitiveType.C128: np.dtype('complex128'), + PrimitiveType.TUPLE: np.dtype(np.object), } # Note the conversion on the key. Numpy has a known issue wherein dtype hashing # doesn't work as expected (https://github.com/numpy/numpy/issues/7242). Thus, # when keying by dtype in this dict, we use the string form of dtypes. -DTYPE_TO_XLA_ELEMENT_TYPE = {str(dt): et - for et, dt in XLA_ELEMENT_TYPE_TO_DTYPE.items()} +DTYPE_TO_XLA_ELEMENT_TYPE = { + str(dt): et for et, dt in XLA_ELEMENT_TYPE_TO_DTYPE.items() +} def dtype_to_etype(dtype): @@ -430,6 +450,13 @@ class LocalBuffer(object): self.delete() +class Format(enum.IntEnum): + """Python copy of the Format protocol buffer enum.""" + INVALID_FORMAT = 0 + DENSE = 1 + SPARSE = 2 + + class Shape(object): """Represents an XLA shape. @@ -459,8 +486,8 @@ class Shape(object): if (not isinstance(dimensions, tuple) or not all(isinstance(i, int) for i in dimensions)): dimensions = tuple(int(i) for i in dimensions) - return Shape(dimensions, np.dtype(element_type), - minor_to_major=minor_to_major) + return Shape( + dimensions, np.dtype(element_type), minor_to_major=minor_to_major) @staticmethod def from_pyval(pyval): @@ -539,8 +566,8 @@ class Shape(object): """Map f over each leaf-level array subshape. Args: - f: The function to apply. Whenever f returns None, the identity is - applied instead. + f: The function to apply. Whenever f returns None, the identity is applied + instead. Returns: A new Shape with the mapped leaves. @@ -565,8 +592,8 @@ class Shape(object): raise ValueError('not an array shape') if not isinstance(minor_to_major, tuple): raise TypeError('minor_to_major must be a tuple') - updated = Shape.array_shape( - self.element_type(), self.dimensions(), minor_to_major) + updated = Shape.array_shape(self.element_type(), self.dimensions(), + minor_to_major) updated._check_minor_to_major() # pylint: disable=protected-access return updated @@ -583,7 +610,7 @@ class Shape(object): def serialize(self, proto): """Serializes 'shape' into proto.""" if self.is_tuple(): - proto.element_type = xla_data_pb2.TUPLE + proto.element_type = PrimitiveType.TUPLE for shape in self.tuple_shapes(): shape.serialize(proto.tuple_shapes.add()) else: @@ -591,7 +618,7 @@ class Shape(object): proto.dimensions.extend(self.dimensions()) proto.is_dynamic_dimension.extend([False for _ in self.dimensions()]) if self.minor_to_major(): - proto.layout.format = xla_data_pb2.DENSE + proto.layout.format = Format.DENSE proto.layout.minor_to_major.extend(self.minor_to_major()) @@ -602,7 +629,7 @@ ProgramShape = collections.namedtuple('ProgramShape', def _wrap_shape(shape_info): dtype, dims = shape_info element_type = DTYPE_TO_XLA_ELEMENT_TYPE[str(dtype)] - if element_type == xla_data_pb2.TUPLE: + if element_type == PrimitiveType.TUPLE: shapes = tuple(_wrap_shape(subshape_info) for subshape_info in dims) return Shape.tuple_shape(shapes) else: @@ -690,15 +717,14 @@ class Computation(object): def computation(self): return self._c_computation - def GetProto(self): - """Get the HloModuleProto proto object in this computation. + def GetSerializedProto(self): + """Gets the serialized HloModuleProto proto object in this computation. Returns: - An HloModuleProto proto object that has the whole-graph information. + A string containing a serialized HloModuleProto proto containing the + computation and its dependencies. """ - serialized = self.computation.GetSerializedProto() - proto = hlo_pb2.HloModuleProto.FromString(serialized) - return proto + return self.computation.GetSerializedProto() def GetHloText(self): """Get the textual HLO representation of this computation. @@ -867,12 +893,6 @@ class Executable(object): self._backend.delete_executable(self._c_executable) -def _make_replica_group_proto(replica_group): - replica_group_proto = xla_data_pb2.ReplicaGroup() - replica_group_proto.replica_ids.extend(replica_group) - return replica_group_proto - - class ComputationBuilder(object): """XLA computation builder. @@ -899,6 +919,7 @@ class ComputationBuilder(object): root: if not None, the operator containing the return value of the computation. backend: deprecated. Pass a `backend` to `Computation.Compile` instead. + Returns: A `Computation`. """ @@ -938,8 +959,8 @@ class ComputationBuilder(object): """Enqueues a constant op onto the computation. Args: - value: value for the constant, as a np.array with an explicit dtype set - to one of the supported types. + value: value for the constant, as a np.array with an explicit dtype set to + one of the supported types. Returns: A LocalOp. @@ -1008,9 +1029,9 @@ class ComputationBuilder(object): Args: shape: the parameter's shape as a Shape object. name: optional string name for the parameter. - parameter_num: parameter number in the computation function. If None, - the next linear parameter number is used. The default value capability - can be used for auto-numbering. If you're using auto-numbering for some + parameter_num: parameter number in the computation function. If None, the + next linear parameter number is used. The default value capability can + be used for auto-numbering. If you're using auto-numbering for some parameters, use it for *all* parameters to avoid clashes. Returns: @@ -1027,8 +1048,8 @@ class ComputationBuilder(object): """Enqueues a Parameter op onto the computation. Args: - value: a Numpy array, or a nested tuple thereof, from which the - shape is inferred. + value: a Numpy array, or a nested tuple thereof, from which the shape is + inferred. name: as in ParameterWithShape. parameter_num: as in ParameterWithShape. @@ -1083,8 +1104,8 @@ class ComputationBuilder(object): Args: operand: the operand LocalOp to broadcast. shape: tuple of integers, the expected output shape. - broadcast_dimensions: tuple of integers identifying which dimensions - of the output are to be broadcast into. + broadcast_dimensions: tuple of integers identifying which dimensions of + the output are to be broadcast into. Returns: A LocalOp representing the added broadcast-in-dimensions op. @@ -1136,20 +1157,28 @@ class ComputationBuilder(object): def GetComputationStats(self): raise NotImplementedError() + def ReplicaId(self): + """Enqueues a ReplicaId operation onto the computation. + + Returns: + A LocalOp representing the replica id. + """ + return self._client.ReplicaId() + def Pad(self, operand, padding_value, padding_config): """Enqueues a Pad operation onto the computation. Args: operand: LocalOp representing the array to pad. padding_value: LocalOp representing the scalar pad value. - padding_config: either an xla_data_pb2.PaddingConfig or a list of integer - triples (edge_padding_low, edge_padding_high, interior_padding) - representing the configuration of the padding operation. + padding_config: either a PaddingConfig or a list of integer triples + (edge_padding_low, edge_padding_high, interior_padding) representing the + configuration of the padding operation. Returns: A LocalOp representing the added Pad op. """ - if not isinstance(padding_config, xla_data_pb2.PaddingConfig): + if isinstance(padding_config, tuple) or isinstance(padding_config, list): padding_config = GetPaddingConfigFromTriples(padding_config) return self._client.Pad(operand, padding_value, padding_config) @@ -1195,7 +1224,8 @@ class ComputationBuilder(object): else: replica_groups = list(replica_groups) replica_groups_protos = [ - _make_replica_group_proto(group) for group in replica_groups] + _make_replica_group_proto(group) for group in replica_groups + ] if not replica_groups: split_count = get_replica_count() else: @@ -1222,7 +1252,8 @@ class ComputationBuilder(object): replica_groups = [] # special value for XLA API else: replica_groups = [ - _make_replica_group_proto(group) for group in replica_groups] + _make_replica_group_proto(group) for group in replica_groups + ] return self._client.CrossReplicaSum(operand, replica_groups) def Collapse(self, operand, dimensions): @@ -1250,8 +1281,8 @@ class ComputationBuilder(object): """Select and scatter op, used by the gradient of ReduceWindow. Args: - operand: LocalOp for array of dimension N and type T over - which the windows slide. + operand: LocalOp for array of dimension N and type T over which the + windows slide. select: Computation of type (T, T) -> Pred to apply to the elements of each window to indicate which element is selected. window_dimensions: sequence of N integers for dimensions of the window. @@ -1266,8 +1297,8 @@ class ComputationBuilder(object): A LocalOp representing the added SelectAndScatter op. """ pads = _convert_padding_type_to_pad_values( - padding, self.GetShape(operand).dimensions(), - window_dimensions, window_strides) + padding, self.GetShape(operand).dimensions(), window_dimensions, + window_strides) return self._client.SelectAndScatterWithGeneralPadding( operand, select.computation, window_dimensions, window_strides, pads, source, init_value, scatter.computation) @@ -1321,8 +1352,8 @@ class ComputationBuilder(object): Args: operand: LocalOp for the N dimensional array to be sliced. - start_indices: LocalOp for the 1D array of N integers - containing the starting indices of the slice. + start_indices: LocalOp for the 1D array of N integers containing the + starting indices of the slice. slice_sizes: iterable of N integers containing the slice sizes in each dimension. @@ -1339,6 +1370,7 @@ class ComputationBuilder(object): update: N dimensional array comprising the slice update. start_indices: Rank-1 array of N integers comprising the starting indices of the slice along each dimension. + Returns: A LocalOp representing the added DynamicUpdateSlice op. """ @@ -1372,8 +1404,8 @@ class ComputationBuilder(object): Args: computation_to_apply: a Computation object. - operands: an iterable of LocalOp. The number and types of - operands must match the arity of computation_to_apply. + operands: an iterable of LocalOp. The number and types of operands must + match the arity of computation_to_apply. Returns: A LocalOp representing the added call op. @@ -1450,8 +1482,8 @@ class ComputationBuilder(object): A LocalOp representing the added ReduceWindow op. """ pads = _convert_padding_type_to_pad_values( - padding, self.GetShape(operand).dimensions(), window_dimensions, - window_strides) + padding, + self.GetShape(operand).dimensions(), window_dimensions, window_strides) return self._client.ReduceWindowWithGeneralPadding( operand, init_value, computation_to_apply.computation, window_dimensions, window_strides, (), (), pads) @@ -1484,10 +1516,8 @@ class ComputationBuilder(object): Args: mu: A LocalOp to an F32 scalar specifying the mean. - sigma: A LocalOp to an F32 scalar specifying the standard - deviation. + sigma: A LocalOp to an F32 scalar specifying the standard deviation. dims: A 1D array-like of nonnegative integers specifying the dimensions. - Returns: a LocalOp to the generated array of F32 values. """ shape = Shape.array_shape(self.GetShape(mu).element_type(), dims) @@ -1497,16 +1527,15 @@ class ComputationBuilder(object): """Enqueues an RngUniform operation onto the computation. Args: - a: a LocalOp to an F32, S32, or U32 scalar (consistent with - the type of b) specifying the low end of the interval [a, b) over which - values are generated. - b: a LocalOp to an F32, S32, or U32 scalar (consistent with - the type of a) specifying the high end of the interval [a, b) over which - values are generated. + a: a LocalOp to an F32, S32, or U32 scalar (consistent with the type of b) + specifying the low end of the interval [a, b) over which values are + generated. + b: a LocalOp to an F32, S32, or U32 scalar (consistent with the type of a) + specifying the high end of the interval [a, b) over which values are + generated. dims: A 1D array-like of nonnegative integers specifying the dimensions. - - Returns: a LocalOp to the generated array of values with the - same numeric type (F32, S32, or U32) as the arguments a and b. + Returns: a LocalOp to the generated array of values with the same numeric + type (F32, S32, or U32) as the arguments a and b. """ shape = Shape.array_shape(self.GetShape(a).element_type(), dims) return self._client.RngUniform(a, b, shape) @@ -1518,7 +1547,6 @@ class ComputationBuilder(object): cond: a Computation for the loop condition, which has type T -> PRED body: a Computation for the loop body, which has type T -> T init: a LocalOp for the initial parameter, which has type T - Returns: a LocalOp representing the While operation. """ return self._client.While(cond.computation, body.computation, init) @@ -1533,19 +1561,17 @@ class ComputationBuilder(object): true_computation: a Computation to apply to true_operand, type T_0 -> S false_operand: a ComputationDatahandle of type T_1 false_computation: a Computation to apply to false_operand, type T_1 -> S - Returns: a LocalOp representing the Conditional operation. """ - return self._client.Conditional( - pred, true_operand, true_computation.computation, false_operand, - false_computation.computation) + return self._client.Conditional(pred, true_operand, + true_computation.computation, false_operand, + false_computation.computation) def IsConstant(self, operand): """Checks whether the given operand is a compile-time constant. Args: operand: a ComputationDataHandle to test. - Returns: bool indicating whether `operand` is a compile-time constant, meaning its value does not depend on any parametersor, or on stateful operators such as `RngNormal` or `Infeed`. @@ -1568,7 +1594,6 @@ class ComputationBuilder(object): Args: lhs: LocalOp for the rank 1 or rank 2 left-hand-side array. rhs: LocalOp for the rank 1 or rank 2 right-hand-side array. - Returns: a LocalOp representing the Dot operation. """ return self._client.Dot(lhs, rhs) @@ -1579,14 +1604,13 @@ class ComputationBuilder(object): Args: lhs: LocalOp for the left-hand-side array. rhs: LocalOp for the right-hand-side array. - dimension_numbers: either an xla_data_pb2.DotDimensionNumbers or a nested - tuple ((lhs_contract, rhs_contract), (lhs_batch, rhs_batch)) of lists of + dimension_numbers: either a DotDimensionNumbers or a nested tuple + ((lhs_contract, rhs_contract), (lhs_batch, rhs_batch)) of lists of integers representing the dimensions to treat as contracting dimensions and batch dimensions on each input operand. - Returns: a LocalOp representing the DotGeneral operation. """ - if not isinstance(dimension_numbers, xla_data_pb2.DotDimensionNumbers): + if isinstance(dimension_numbers, tuple): dimension_numbers = GetDotDimensionsFromLists(dimension_numbers) return self._client.DotGeneral(lhs, rhs, dimension_numbers) @@ -1599,15 +1623,15 @@ class ComputationBuilder(object): window_strides: length-N array-like of integer kernel strides. padding: PaddingType representing either 'SAME' or 'VALID' padding. feature_group_count: number of feature groups for grouped convolution. - Returns: a LocalOp representing the Conv operation. """ pads = _convert_padding_type_to_pad_values( - padding, self.GetShape(lhs).dimensions()[2:], + padding, + self.GetShape(lhs).dimensions()[2:], self.GetShape(rhs).dimensions()[2:], window_strides) return self.ConvGeneralDilated( - lhs, rhs, window_strides, pads, (), (), - dimension_numbers=None, feature_group_count=feature_group_count) + lhs, rhs, window_strides, pads, (), (), dimension_numbers=None, + feature_group_count=feature_group_count) def ConvWithGeneralPadding(self, lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, feature_group_count=1): @@ -1632,7 +1656,7 @@ class ComputationBuilder(object): def _GetConvDimensionNumbers(self, num_spatial_dims): """Create ConvolutionDimensionNumbers proto for convolutions.""" nd = num_spatial_dims - dimension_numbers = xla_data_pb2.ConvolutionDimensionNumbers() + dimension_numbers = ConvolutionDimensionNumbers() dimension_numbers.input_batch_dimension = 0 dimension_numbers.input_feature_dimension = 1 dimension_numbers.output_batch_dimension = 0 @@ -1656,35 +1680,33 @@ class ComputationBuilder(object): padding: length-N array-like of pairs of integers of (low, high) padding. lhs_dilation: length-N array-like of integer dilation factors. rhs_dilation: length-N array-like of integer dilation factors. - dimension_numbers: optional, either an - xla_data_pb2.ConvolutionDimensionNumbers proto instance or a tuple - (lhs_spec, rhs_spec, out_spec) where each element is a string of length - N+2 identifying by position (1) batch dimensions in lhs, rhs, and the - output with the character 'N', (2) feature dimensions in lhs and the - output with the character 'C', (3) input and output feature dimensions - in rhs with the characters 'I' and 'O' respectively, and (4) spatial - dimension correspondences between lhs, rhs, and the output using any - distinct characters. For example, to indicate dimension numbers - consistent with the Conv operation with two spatial dimensions, one - could use ('NCHW', 'OIHW', 'NCHW'). As another example, to indicate - dimension numbers consistent with the TensorFlow Conv2D operation, one - could use ('NHWC', 'HWIO', 'NHWC'). When using the latter form of - convolution dimension specification, window strides are associated with - spatial dimension character labels according to the order in which the - labels appear in the rhs_spec string, so that window_strides[0] is - matched with the dimension corresponding to the first character - appearing in rhs_spec that is not 'I' or 'O'. By default, use the same - dimension numbering as Conv and ConvWithGeneralPadding. + dimension_numbers: optional, either a ConvolutionDimensionNumbers object + or a tuple (lhs_spec, rhs_spec, out_spec). Each element is a string of + length N+2 identifying by position: (1) batch dimensions in lhs, rhs, + and the output with the character 'N', (2) feature dimensions in lhs + and the output with the character 'C', (3) input and output feature + dimensions in rhs with the characters 'I' and 'O' respectively, and + (4) spatial dimension correspondences between lhs, rhs, and the output + using any distinct characters. For example, to indicate dimension + numbers consistent with the Conv operation with two spatial + dimensions, one could use ('NCHW', 'OIHW', 'NCHW'). As another + example, to indicate dimension numbers consistent with the TensorFlow + Conv2D operation, one could use ('NHWC', 'HWIO', 'NHWC'). When using + the latter form of convolution dimension specification, window strides + are associated with spatial dimension character labels according to + the order in which the labels appear in the rhs_spec string, so that + window_strides[0] is matched with the dimension corresponding to the + first character appearing in rhs_spec that is not 'I' or 'O'. By + default, use the same dimension numbering as Conv and + ConvWithGeneralPadding. feature_group_count: number of feature groups for grouped convolution. - Returns: a LocalOp representing the ConvGenralDilated operation. """ if dimension_numbers is None: dimension_numbers = self._GetConvDimensionNumbers(len(window_strides)) - elif not isinstance(dimension_numbers, - xla_data_pb2.ConvolutionDimensionNumbers): + elif isinstance(dimension_numbers, tuple): lhs_spec, rhs_spec, out_spec = dimension_numbers - dimension_numbers = xla_data_pb2.ConvolutionDimensionNumbers() + dimension_numbers = ConvolutionDimensionNumbers() dimension_numbers.input_batch_dimension = lhs_spec.index('N') dimension_numbers.input_feature_dimension = lhs_spec.index('C') @@ -1701,10 +1723,9 @@ class ComputationBuilder(object): dimension_numbers.output_spatial_dimensions.extend( sorted((i for i, c in enumerate(out_spec) if c not in {'N', 'C'}), key=lambda i: rhs_spec.index(out_spec[i]))) - return self._client.ConvGeneralDilated(lhs, rhs, window_strides, padding, - lhs_dilation, rhs_dilation, - dimension_numbers, - feature_group_count) + return self._client.ConvGeneralDilated( + lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, + dimension_numbers, feature_group_count) def Sort(self, operand, dimension=-1): """Enqueues a sort operation onto the computation.""" @@ -1714,9 +1735,9 @@ class ComputationBuilder(object): """Enqueues a key-value sort operation onto the computation.""" return self._client.SortKeyVal(keys, values, dimension) - def Cholesky(self, a): + def Cholesky(self, a, lower=True): """Enqueues a Cholesky decomposition onto the computation.""" - return self._client.Cholesky(a) + return self._client.Cholesky(a, lower) def QR(self, a, full_matrices=True): """Enqueues a QR decomposition onto the computation.""" @@ -1742,15 +1763,14 @@ class ComputationBuilder(object): def Gather(self, a, start_indices, dimension_numbers, slice_sizes): """Enqueues a Gather operation onto the computation.""" - return self._client.Gather(a, start_indices, dimension_numbers, - slice_sizes) + return self._client.Gather(a, start_indices, dimension_numbers, slice_sizes) def Scatter(self, a, scatter_indices, updates, update_computation, dimension_numbers): """Enqueues a Scatter operation onto the computation.""" return self._client.Scatter( a, scatter_indices, updates, update_computation.computation, - dimension_numbers,) + dimension_numbers) def _forward_methods_to_local_builder(): @@ -1791,7 +1811,6 @@ def _forward_methods_to_local_builder(): _forward_methods_to_local_builder() - _default_replica_count = 1 @@ -1844,22 +1863,111 @@ def register_cpu_custom_call_target(name, fn): c_api.RegisterCpuCustomCallTarget(name, fn) +class PaddingConfigDimension(object): + """Python representation of a xla.PaddingConfigDimension protobuf.""" + __slots__ = ('edge_padding_low', 'edge_padding_high', 'interior_padding') + + def __init__(self): + self.edge_padding_low = [] + self.edge_padding_high = [] + self.interior_padding = [] + + +class PaddingConfig(object): + """Python representation of a xla.PaddingConfig protobuf.""" + __slots__ = ('dimensions',) + + def __init__(self): + self.dimensions = [] + + def GetPaddingConfigFromTriples(triples): """Create PaddingConfig proto from list of triples of integers.""" - padding_config = xla_data_pb2.PaddingConfig() + padding_config = PaddingConfig() for lo, hi, interior in triples: - dimension = padding_config.dimensions.add() + dimension = PaddingConfigDimension() dimension.edge_padding_low = lo dimension.edge_padding_high = hi dimension.interior_padding = interior + padding_config.dimensions.append(dimension) return padding_config +class DotDimensionNumbers(object): + """Python representation of a xla.DotDimensionNumbers protobuf.""" + __slots__ = ('lhs_contracting_dimensions', 'rhs_contracting_dimensions', + 'lhs_batch_dimensions', 'rhs_batch_dimensions') + + def __init__(self): + self.lhs_contracting_dimensions = [] + self.rhs_contracting_dimensions = [] + self.lhs_batch_dimensions = [] + self.rhs_batch_dimensions = [] + + def GetDotDimensionsFromLists(dimension_numbers): (lhs_contract, rhs_contract), (lhs_batch, rhs_batch) = dimension_numbers - dot_dims_proto = xla_data_pb2.DotDimensionNumbers() + dot_dims_proto = DotDimensionNumbers() dot_dims_proto.lhs_contracting_dimensions.extend(lhs_contract) dot_dims_proto.rhs_contracting_dimensions.extend(rhs_contract) dot_dims_proto.lhs_batch_dimensions.extend(lhs_batch) dot_dims_proto.rhs_batch_dimensions.extend(rhs_batch) return dot_dims_proto + + +class ConvolutionDimensionNumbers(object): + """Python representation of a xla.ConvolutionDimensionNumbers protobuf.""" + __slots__ = ('input_batch_dimension', 'input_feature_dimension', + 'input_spatial_dimensions', 'kernel_input_feature_dimension', + 'kernel_output_feature_dimension', 'kernel_spatial_dimensions', + 'output_batch_dimension', 'output_feature_dimension', + 'output_spatial_dimensions') + + def __init__(self): + self.input_batch_dimension = 0 + self.input_feature_dimension = 0 + self.input_spatial_dimensions = [] + self.kernel_input_feature_dimension = 0 + self.kernel_output_feature_dimension = 0 + self.kernel_spatial_dimensions = [] + self.output_batch_dimension = 0 + self.output_feature_dimension = 0 + self.output_spatial_dimensions = [] + + +class GatherDimensionNumbers(object): + """Python representation of a xla.GatherDimensionNumbers protobuf.""" + __slots__ = ('offset_dims', 'collapsed_slice_dims', 'start_index_map', + 'index_vector_dim') + + def __init__(self): + self.offset_dims = [] + self.collapsed_slice_dims = [] + self.start_index_map = [] + self.index_vector_dim = 0 + + +class ScatterDimensionNumbers(object): + """Python representation of a xla.ScatterDimensionNumbers protobuf.""" + __slots__ = ('update_window_dims', 'inserted_window_dims', + 'scatter_dims_to_operand_dims', 'index_vector_dim') + + def __init__(self): + self.update_window_dims = [] + self.inserted_window_dims = [] + self.scatter_dims_to_operand_dims = [] + self.index_vector_dim = 0 + + +class ReplicaGroup(object): + """Python representation of a xla.ReplicaGroup protobuf.""" + __slots__ = ('replica_ids',) + + def __init__(self): + self.replica_ids = [] + + +def _make_replica_group_proto(replica_group): + replica_group_proto = ReplicaGroup() + replica_group_proto.replica_ids.extend(replica_group) + return replica_group_proto diff --git a/tensorflow/compiler/xla/python/xla_client_test.py b/tensorflow/compiler/xla/python/xla_client_test.py index 51ef7d7f3a17f341e955f48615b05a886813430b..65594f5669daa10dc3fcafa5ac865132fb265dfe 100644 --- a/tensorflow/compiler/xla/python/xla_client_test.py +++ b/tensorflow/compiler/xla/python/xla_client_test.py @@ -24,11 +24,24 @@ import threading import numpy as np +from tensorflow.compiler.xla import xla_data_pb2 from tensorflow.compiler.xla.python import custom_call_for_test from tensorflow.compiler.xla.python import xla_client import unittest +class EnumTest(unittest.TestCase): + """Verifies Python enumerations match their protocol buffer equivalents.""" + + def testPrimitiveType(self): + for name, value in xla_client.PrimitiveType.__members__.items(): + self.assertEqual(value, getattr(xla_data_pb2, name)) + + def testFormat(self): + for name, value in xla_client.Format.__members__.items(): + self.assertEqual(value, getattr(xla_data_pb2, name)) + + class ComputationTest(unittest.TestCase): """Base class for running an XLA Computation through the local client.""" @@ -230,16 +243,6 @@ class ComputationsWithConstantsTest(ComputationTest): c.Constant(NumpyArrayS32([1]))) self._ExecuteAndCompareClose(c, expected=[2**31 - 1]) - def testGetProto(self): - c = self._NewComputation() - c.Add( - c.Constant(NumpyArrayF32([[1, 2, 3], [4, 5, 6]])), - c.Constant(NumpyArrayF32([[1, -1, 1], [-1, 1, -1]]))) - built = c.Build() - proto = built.GetProto() # HloModuleProto - self.assertTrue(len(proto.computations) == 1) - self.assertTrue(len(proto.computations[0].instructions) == 3) - def testSum2DF64(self): c = self._NewComputation() c.Add( @@ -528,11 +531,11 @@ class SingleOpTest(ComputationTest): def testConvertElementType(self): xla_types = { - np.bool: xla_client.xla_data_pb2.PRED, - np.int32: xla_client.xla_data_pb2.S32, - np.int64: xla_client.xla_data_pb2.S64, - np.float32: xla_client.xla_data_pb2.F32, - np.float64: xla_client.xla_data_pb2.F64, + np.bool: xla_client.PrimitiveType.PRED, + np.int32: xla_client.PrimitiveType.S32, + np.int64: xla_client.PrimitiveType.S64, + np.float32: xla_client.PrimitiveType.F32, + np.float64: xla_client.PrimitiveType.F64, } def _ConvertAndTest(template, src_dtype, dst_dtype): @@ -553,13 +556,13 @@ class SingleOpTest(ComputationTest): def testBitcastConvertType(self): xla_x32_types = { - np.int32: xla_client.xla_data_pb2.S32, - np.float32: xla_client.xla_data_pb2.F32, + np.int32: xla_client.PrimitiveType.S32, + np.float32: xla_client.PrimitiveType.F32, } xla_x64_types = { - np.int64: xla_client.xla_data_pb2.S64, - np.float64: xla_client.xla_data_pb2.F64, + np.int64: xla_client.PrimitiveType.S64, + np.float64: xla_client.PrimitiveType.F64, } def _ConvertAndTest(template, src_dtype, dst_dtype, dst_etype): @@ -579,7 +582,7 @@ class SingleOpTest(ComputationTest): for src_dtype, dst_dtype in itertools.product(xla_types, xla_types): _ConvertAndTest(x, src_dtype, dst_dtype, xla_types[dst_dtype]) - # TODO(b/123523486): re-enable when shape check is resolved + # TODO(b/123523486) implement AllToAll on CPU def DISABLED_testAllToAllOneReplica(self): samples = [ NumpyArrayF32([97.0]), @@ -603,6 +606,11 @@ class SingleOpTest(ComputationTest): c.CrossReplicaSum(c.Constant(lhs)) self._ExecuteAndCompareExact(c, expected=lhs) + def testReplicaId(self): + c = self._NewComputation() + _ = c.ReplicaId() + self._ExecuteAndCompareExact(c, expected=0) + def testCrossReplicaSumOneReplicaWithSingletonGroup(self): samples = [ NumpyArrayF32(42.0), @@ -658,7 +666,7 @@ class SingleOpTest(ComputationTest): lhs = NumpyArrayF32(rng.randn(10, 3, 4)) rhs = NumpyArrayF32(rng.randn(10, 4, 5)) - dimension_numbers = xla_client.xla_data_pb2.DotDimensionNumbers() + dimension_numbers = xla_client.DotDimensionNumbers() dimension_numbers.lhs_contracting_dimensions.append(2) dimension_numbers.rhs_contracting_dimensions.append(1) dimension_numbers.lhs_batch_dimensions.append(0) @@ -970,12 +978,13 @@ class SingleOpTest(ComputationTest): def testPadWithPaddingConfig(self): c = self._NewComputation() - padding_config = xla_client.xla_data_pb2.PaddingConfig() + padding_config = xla_client.PaddingConfig() for lo, hi, interior in [(1, 2, 1), (0, 1, 0)]: - dimension = padding_config.dimensions.add() + dimension = xla_client.PaddingConfigDimension() dimension.edge_padding_low = lo dimension.edge_padding_high = hi dimension.interior_padding = interior + padding_config.dimensions.append(dimension) c.Pad( c.Constant(NumpyArrayF32([[1.0, 2.0], [3.0, 4.0]])), c.Constant(NumpyArrayF32(0.0)), @@ -1018,14 +1027,13 @@ class SingleOpTest(ComputationTest): c.Constant(NumpyArrayF32(2))) self._ExecuteAndCompareExact(c, expected=[-1, -1, 0, 1, 2, 2]) - # TODO(b/72689392): re-enable when bug S32 resolved - def DISABLED_testClampS32(self): + def testClampS32(self): c = self._NewComputation() c.Clamp( c.Constant(NumpyArrayS32(-1)), c.Constant(NumpyArrayS32([-2, -1, 0, 1, 2, 3])), c.Constant(NumpyArrayS32(2))) - self._ExecuteAndCompareExact(c, expected=[-1, 0, 1, 2, 2]) + self._ExecuteAndCompareExact(c, expected=[-1, -1, 0, 1, 2, 2]) def testSelect(self): c = self._NewComputation() @@ -1188,7 +1196,7 @@ class SingleOpTest(ComputationTest): def testGather(self): a = np.arange(9).astype(np.int32).reshape((3, 3)) indices = np.array([[[0, 2], [2, 1]], [[1, 2], [2, 0]]], dtype=np.int32) - dnums = xla_client.xla_data_pb2.GatherDimensionNumbers() + dnums = xla_client.GatherDimensionNumbers() dnums.offset_dims.append(1) dnums.offset_dims.append(2) dnums.start_index_map.append(0) @@ -1652,7 +1660,7 @@ class EmbeddedComputationsTest(ComputationTest): scatter_indices = np.array([0, 2], dtype=np.int32) updates = np.array([[10, 20, 30], [70, 80, 90]], dtype=np.int32) - dnums = xla_client.xla_data_pb2.ScatterDimensionNumbers() + dnums = xla_client.ScatterDimensionNumbers() dnums.update_window_dims.append(1) dnums.inserted_window_dims.append(0) dnums.scatter_dims_to_operand_dims.append(0) diff --git a/tensorflow/compiler/xla/python/xla_data.i b/tensorflow/compiler/xla/python/xla_data.i index 974f314af24f61c0015a8d51c16dff1bfc84c7cc..b18583c64d400bdb7b3bc50b3548df23f4a8c469 100644 --- a/tensorflow/compiler/xla/python/xla_data.i +++ b/tensorflow/compiler/xla/python/xla_data.i @@ -13,9 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// SWIG typemaps and declarations for building, compiling, and -// executing XLA computations, wrapping most of what is declared in -// xla_data.h. +// SWIG typemaps for building, compiling, and executing XLA computations. // // The typemaps below implement/assert the following correspondences // (with elaborations below): @@ -33,11 +31,11 @@ limitations under the License. // std::vector <- sequence of xla_client.Shape objects // PrimitiveType <- int // Span> <- sequence of int pairs -// PaddingConfig proto <- corresponding Python proto -// ConvolutionDimensionNumbers proto <- corresponding Python proto -// DotDimensionNumbers proto <- corresponding Python proto -// GatherDimensionNumbers proto <- corresponding Python proto -// ScatterDimensionNumbers proto <- corresponding Python proto +// PaddingConfig proto <- ducktyped Python proto +// ConvolutionDimensionNumbers proto <- ducktyped Python proto +// DotDimensionNumbers proto <- ducktyped Python proto +// GatherDimensionNumbers proto <- ducktyped Python proto +// ScatterDimensionNumbers proto <- ducktyped Python proto // Span <- sequence of ReplicaGroup Python proto // // Arrows indicate whether a conversion only ever occurs in one @@ -102,6 +100,8 @@ limitations under the License. // Other SWIG object wrappers (e.g. of Computation) are further // wrapped by xla_client in order to set up a custom destructor that // triggers memory deallocation on the C++ side. +// + %module(threads="1") xla_data diff --git a/tensorflow/compiler/xla/python/xrt.h b/tensorflow/compiler/xla/python/xrt.h index dd5bba6d5c9641dadc323f70745e870c14543321..710c3af3fa6b407127643797dbabad201cf076d4 100644 --- a/tensorflow/compiler/xla/python/xrt.h +++ b/tensorflow/compiler/xla/python/xrt.h @@ -69,7 +69,7 @@ class XrtAllocationTuple { std::vector elements_; }; -// Destructures a tuple-valued XrtAllocation into its constitutent elements +// Destructures a tuple-valued XrtAllocation into its constituent elements // in XrtAllocationTuple form. // // Accepts a `session_target` argument, used in constructing the diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD index 8d8394cb43ee013b9396a54e3a4d037445fcc0e1..3f47793f3d3d84c7532011ad4fa9cb56434ea694 100644 --- a/tensorflow/compiler/xla/service/BUILD +++ b/tensorflow/compiler/xla/service/BUILD @@ -1581,6 +1581,29 @@ cc_library( ], ) +cc_library( + name = "cholesky_expander", + srcs = ["cholesky_expander.cc"], + hdrs = ["cholesky_expander.h"], + deps = [ + ":op_expander_pass", + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/client:xla_computation", + "//tensorflow/compiler/xla/client/lib:constants", + "//tensorflow/compiler/xla/client/lib:loops", + "//tensorflow/compiler/xla/client/lib:math", + "//tensorflow/compiler/xla/client/lib:matrix", + "//tensorflow/compiler/xla/client/lib:slicing", + "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", + ], +) + tf_cc_test( name = "batchnorm_expander_test", size = "small", @@ -1788,6 +1811,8 @@ cc_library( deps = [ ":hlo", ":hlo_evaluator", + ":pattern_matcher", + "@com_google_absl//absl/base", "@com_google_absl//absl/types:optional", ], ) @@ -1847,6 +1872,36 @@ tf_cc_test( ], ) +cc_library( + name = "while_loop_trip_count_annotator", + srcs = ["while_loop_trip_count_annotator.cc"], + hdrs = ["while_loop_trip_count_annotator.h"], + deps = [ + ":hlo", + ":hlo_pass", + ":while_loop_analysis", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:xla_data_proto", + ], +) + +tf_cc_test( + name = "while_loop_trip_count_annotator_test", + srcs = ["while_loop_trip_count_annotator_test.cc"], + deps = [ + ":pattern_matcher", + ":while_loop_simplifier", + ":while_loop_trip_count_annotator", + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", # fixdeps: keep + "//tensorflow/core:test", + ], +) + cc_library( name = "defuser", srcs = ["defuser.cc"], @@ -2005,9 +2060,11 @@ tf_cc_test( srcs = ["dynamic_padder_test.cc"], deps = [ ":dynamic_padder", + ":hlo_parser", "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:xla_data_proto", @@ -3438,10 +3495,13 @@ cc_library( srcs = ["hlo_runner.cc"], hdrs = ["hlo_runner.h"], deps = [ + ":backend", + ":compiler", ":computation_placer", ":executable", ":hlo", ":hlo_module_group", + ":hlo_parser", ":transfer_manager", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -3449,11 +3509,9 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/service:backend", - "//tensorflow/compiler/xla/service:compiler", - "//tensorflow/compiler/xla/service:hlo_parser", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", "//tensorflow/core:stream_executor_no_cuda", "//third_party/eigen3", "@com_google_absl//absl/memory", @@ -3723,11 +3781,13 @@ cc_library( "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/types:span", "@com_google_absl//absl/types:variant", ], ) diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/tensorflow/compiler/xla/service/algebraic_simplifier.cc index bd17e96106abd9de0dd3bbf418439b0fb3edb746..9429b7fb1a24903a386d10b8fabb3e3ab2d22f8d 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier.cc @@ -109,7 +109,7 @@ bool IsAllFpConstantPowerOf2(const HloInstruction* op) { int exp; double mantissa = std::frexp(*val, &exp); - // frexp returns a value in the range (-1; -0.5] U [0.5, 1). A return value + // frexp returns a value in the range (-1, -0.5] U [0.5, 1). A return value // of +/-0.5 therefore indicates that the floating point value is a power of // 2. return mantissa == 0.5 || mantissa == -0.5; @@ -1031,6 +1031,24 @@ Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { divide->shape(), HloOpcode::kMultiply, a, new_power)); } + // A/sqrt(B) => A*rsqrt(X). + if (Match(divide, m::Divide(m::Op(&a), m::Sqrt(m::Op(&b))))) { + auto* rsqrt = computation_->AddInstruction( + HloInstruction::CreateUnary(divide->shape(), HloOpcode::kRsqrt, b)); + return ReplaceWithNewInstruction( + divide, HloInstruction::CreateBinary(rsqrt->shape(), + HloOpcode::kMultiply, a, rsqrt)); + } + + // A/rsqrt(B) => A*sqrt(B). + if (Match(divide, m::Divide(m::Op(&a), m::Rsqrt(m::Op(&b))))) { + auto* sqrt = computation_->AddInstruction( + HloInstruction::CreateUnary(divide->shape(), HloOpcode::kSqrt, b)); + return ReplaceWithNewInstruction( + divide, HloInstruction::CreateBinary(sqrt->shape(), + HloOpcode::kMultiply, a, sqrt)); + } + // Simplifying integral division would produce unexpected results. if (ShapeUtil::ElementIsIntegral(divide->shape())) { return Status::OK(); @@ -2820,6 +2838,27 @@ Status AlgebraicSimplifierVisitor::HandleSlice(HloInstruction* slice) { new_slice_starts, new_slice_limits, slice->slice_strides())); } + auto only_broadcast_dims_sliced = [&] { + if (slice->operand(0)->opcode() != HloOpcode::kBroadcast) { + return false; + } + for (int64 dim : slice->operand(0)->dimensions()) { + if (slice->slice_starts(dim) != 0 || slice->slice_strides(dim) != 1 || + slice->slice_limits(dim) != + slice->operand(0)->shape().dimensions(dim)) { + return false; + } + } + return true; + }; + if (only_broadcast_dims_sliced()) { + return ReplaceWithNewInstruction( + slice, + HloInstruction::CreateBroadcast( + slice->shape(), slice->mutable_operand(0)->mutable_operand(0), + slice->mutable_operand(0)->dimensions())); + } + TF_ASSIGN_OR_RETURN(bool replaced, TrySimplifyScalarSlice(slice)); if (replaced) { return Status::OK(); diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc index af03fcb100813e8942efcaefc296b971c01a6aaa..d04e40b04f7ac3f2bc9391d04331320b3c44482e 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc @@ -2694,6 +2694,33 @@ TEST_F(AlgebraicSimplifierTest, SliceOfSliceToSlice) { EXPECT_EQ(computation->root_instruction()->slice_limits(1), dim1 - 4); } +TEST_F(AlgebraicSimplifierTest, SliceOfBroadcastToBroadcast) { + HloComputation::Builder builder(TestName()); + const int64 dim0 = 11; + const int64 dim1 = 12; + HloInstruction* param = + builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {dim0}), "param")); + HloInstruction* broadcast = + builder.AddInstruction(HloInstruction::CreateBroadcast( + ShapeUtil::MakeShape(F32, {dim0, dim1}), param, {0})); + builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {dim0, dim1 - 9}), broadcast, + /*start_indices=*/{0, 3}, + /*limit_indices=*/{dim0, dim1 - 6}, /*strides=*/{1, 1})); + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Slice(m::Broadcast(m::Parameter(0))))); + + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Broadcast(m::Parameter(0)))); +} + TEST_F(AlgebraicSimplifierTest, SliceOfReshapeToReshapeOfSlice) { HloComputation::Builder builder(TestName()); const int64 dim0 = 11; @@ -4879,5 +4906,41 @@ TEST_F(AlgebraicSimplifierTest, DividedByConstantInstructionWithoutLayout) { EXPECT_THAT(root, GmockMatch(m::Multiply())); } +// Test that 1/sqrt(X) is simplified to rsqrt(X). +TEST_F(AlgebraicSimplifierTest, RecipSqrt) { + const char* kModuleStr = R"( + HloModule m + test { + p0 = f32[] parameter(0) + p1 = f32[] parameter(1) + sqrt = f32[] sqrt(p0) + ROOT div = f32[] divide(p1, sqrt) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).ValueOrDie()); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::MultiplyAnyOrder(m::Parameter(1), + m::Rsqrt(m::Parameter(0))))); +} + +// Test that 1/rsqrt(X) is simplified to sqrt(X). +TEST_F(AlgebraicSimplifierTest, RecipRsqrt) { + const char* kModuleStr = R"( + HloModule m + test { + p0 = f32[] parameter(0) + p1 = f32[] parameter(1) + rsqrt = f32[] rsqrt(p0) + ROOT div = f32[] divide(p1, rsqrt) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).ValueOrDie()); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::MultiplyAnyOrder(m::Parameter(1), + m::Sqrt(m::Parameter(0))))); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/batchnorm_expander.cc b/tensorflow/compiler/xla/service/batchnorm_expander.cc index e5f5c3edb2ac0c217317fbf809463aa31af9af59..d14e803be6ad6d0b7a7e22442de7e6da77f93577 100644 --- a/tensorflow/compiler/xla/service/batchnorm_expander.cc +++ b/tensorflow/compiler/xla/service/batchnorm_expander.cc @@ -95,15 +95,8 @@ class BatchNormExpanderVisitor : public DfsHloVisitorWithDefault { HloInstruction* operand, const std::function)>& add_instruction) { - HloInstruction* exponent = add_instruction(HloInstruction::CreateBroadcast( - operand->shape(), - add_instruction(HloInstruction::CreateConvert( - ShapeUtil::MakeShape(operand->shape().element_type(), {}), - add_instruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR0(-0.5f))))), - {})); - return HloInstruction::CreateBinary(operand->shape(), HloOpcode::kPower, - operand, exponent); + return HloInstruction::CreateUnary(operand->shape(), HloOpcode::kRsqrt, + operand); } std::unique_ptr Mean( @@ -524,7 +517,7 @@ Status BatchNormExpanderVisitor::HandleBatchNormGrad( activation_shape, HloOpcode::kSubtract, activation, mean_broadcasted); // Grad[Y] * (X - E[X]). - auto grad_output_times_activiation_minus_mean = + auto grad_output_times_activation_minus_mean = add_binary(activation_shape, HloOpcode::kMultiply, grad_output, activation_minus_mean); @@ -532,9 +525,9 @@ Status BatchNormExpanderVisitor::HandleBatchNormGrad( GetOrCreateScalarAddComputation(ptype); // sum(Grad[Y] * (X - E[X])). - auto sum_grad_output_times_activiation_minus_mean = + auto sum_grad_output_times_activation_minus_mean = add(HloInstruction::CreateReduce( - feature_shape, grad_output_times_activiation_minus_mean, zero, + feature_shape, grad_output_times_activation_minus_mean, zero, dimensions_without_feature, add_reduce_computation)); // Grad[beta] = Sum(Grad[Y]). @@ -544,7 +537,7 @@ Status BatchNormExpanderVisitor::HandleBatchNormGrad( // Grad[scale] = Sum(Grad[Y] * (X - E[X]) * rsqrt[Var[X] + epsilon]). auto grad_scale = add_binary(feature_shape, HloOpcode::kMultiply, - sum_grad_output_times_activiation_minus_mean, + sum_grad_output_times_activation_minus_mean, rsqrt_var_add_epsilon); // I2 = Sum(Grad[Y]) @@ -553,7 +546,7 @@ Status BatchNormExpanderVisitor::HandleBatchNormGrad( // I3 = Sum(Grad[Y] * (X - E[X])) auto i3 = add(HloInstruction::CreateBroadcast( - activation_shape, sum_grad_output_times_activiation_minus_mean, + activation_shape, sum_grad_output_times_activation_minus_mean, {feature_index})); // I4 = (X - E[X]) * I3 diff --git a/tensorflow/compiler/xla/service/batchnorm_expander_test.cc b/tensorflow/compiler/xla/service/batchnorm_expander_test.cc index 8e8fbbd935b154e5a77d68e60d861601d740bf03..34b516184fa861bd71f99f70a32782d242f11914 100644 --- a/tensorflow/compiler/xla/service/batchnorm_expander_test.cc +++ b/tensorflow/compiler/xla/service/batchnorm_expander_test.cc @@ -60,7 +60,7 @@ TEST_F(BatchNormExpanderTest, BatchNormTraining) { HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( - HloInstruction::CreateParameter(0, input_shape, "activiation")); + HloInstruction::CreateParameter(0, input_shape, "activation")); HloInstruction* param1 = builder.AddInstruction( HloInstruction::CreateParameter(1, scale_shape, "scale")); diff --git a/tensorflow/compiler/xla/service/bfloat16_normalization.cc b/tensorflow/compiler/xla/service/bfloat16_normalization.cc index d1b14d604f0559b6b18f7d1fba127669c241c8a3..72459961485f77b690eed6b8bde2cd03ebe770f1 100644 --- a/tensorflow/compiler/xla/service/bfloat16_normalization.cc +++ b/tensorflow/compiler/xla/service/bfloat16_normalization.cc @@ -84,7 +84,12 @@ Status BFloat16NormalizationVisitor::InsertConvertAfterOutput( auto convert = computation->AddInstruction( HloInstruction::CreateConvert(hlo->shape(), hlo)); for (auto* user : materialized_users) { - TF_RETURN_IF_ERROR(hlo->ReplaceUseWith(user, convert)); + if (user->opcode() == HloOpcode::kConvert && + user->shape().element_type() == F32) { + TF_RETURN_IF_ERROR(user->ReplaceAllUsesWith(hlo)); + } else { + TF_RETURN_IF_ERROR(hlo->ReplaceUseWith(user, convert)); + } } if (is_root) { computation->set_root_instruction(convert); @@ -205,6 +210,28 @@ Status BFloat16NormalizationVisitor::HandleMultipleOutputs( return Status::OK(); } + std::vector bf16_called_comps; + for (auto* comp : hlo->called_computations()) { + bool comp_has_bf16 = false; + if (comp->root_instruction()->shape().element_type() == F32) { + f32_count += 1; + } else if (comp->root_instruction()->shape().element_type() == BF16) { + bf16_count += 1; + comp_has_bf16 = true; + } + for (auto* param : comp->parameter_instructions()) { + if (param->shape().element_type() == F32) { + f32_count += 1; + } else if (param->shape().element_type() == BF16) { + bf16_count += 1; + comp_has_bf16 = true; + } + } + if (comp_has_bf16) { + bf16_called_comps.push_back(comp); + } + } + std::vector materialized_users = hlo->users(); std::vector output_elements(hlo->operand_count()); auto original_shape = hlo->shape(); @@ -236,7 +263,7 @@ Status BFloat16NormalizationVisitor::HandleMultipleOutputs( computation_->set_root_instruction(tuple); } *tuple->mutable_shape() = original_shape; - return Status::OK(); + return ConvertCalledComputations(hlo, bf16_called_comps); } Status BFloat16NormalizationVisitor::HandleInstruction(HloInstruction* hlo) { diff --git a/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc b/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc index 2caa979745b3b40817acb1b6951e1de5ffa294a4..7dd46ca4e048210843e227c79f639be1bd34fe30 100644 --- a/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc +++ b/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc @@ -326,6 +326,14 @@ TEST_F(BFloat16NormalizationTest, ResolveMixedPrecisionTupleSortRoot) { EXPECT_EQ(ShapeUtil::GetSubshape(sort->shape(), {0}).element_type(), F32); EXPECT_NE(computation->root_instruction(), sort); EXPECT_EQ(computation->root_instruction()->opcode(), HloOpcode::kTuple); + EXPECT_EQ(sort->to_apply()->parameter_instruction(1)->shape().element_type(), + F32); + // Make sure that no convert to BF16 was added to the 'to_apply' comparison + // computation. + auto users = sort->to_apply()->parameter_instruction(1)->users(); + for (auto user : users) { + EXPECT_NE(user->opcode(), HloOpcode::kConvert); + } } // Tests that the normalization should not cause unsupported mixed precision due diff --git a/tensorflow/compiler/xla/service/buffer_assignment.cc b/tensorflow/compiler/xla/service/buffer_assignment.cc index cbebbdc8a2d7d0b65f12accbe424bea383ff5355..cb682f49a5c8097b2fa5ce15ea9fdbbcf46668b4 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment.cc +++ b/tensorflow/compiler/xla/service/buffer_assignment.cc @@ -1620,62 +1620,46 @@ void BufferAssigner::BuildColocatedBufferSets( AddSetToColocatedBufferSets(colocated_set, colocated_buffer_sets); }); } else if (opcode == HloOpcode::kConditional) { - const HloInstruction* conditional_hlo = instruction; + const HloInstruction* conditional = instruction; ShapeUtil::ForEachSubshape( - conditional_hlo->shape(), - [this, conditional_hlo, &points_to_analysis, colocated_buffer_sets]( + conditional->shape(), + [this, conditional, &points_to_analysis, colocated_buffer_sets]( const Shape& /*subshape*/, const ShapeIndex& index) { std::vector colocated_set; - // Add conditional.result. - AddBufferToColocatedSet(conditional_hlo, index, - points_to_analysis, &colocated_set); - // Add conditional.true_computation.root. - AddBufferToColocatedSet( - conditional_hlo->true_computation()->root_instruction(), - index, points_to_analysis, &colocated_set); - // Add conditional.false_computation.root. - AddBufferToColocatedSet( - conditional_hlo->false_computation()->root_instruction(), - index, points_to_analysis, &colocated_set); + // Add cond.result. + AddBufferToColocatedSet(conditional, index, points_to_analysis, + &colocated_set); + for (int j = 0; j < conditional->branch_count(); ++j) { + // Add each cond.branch_computation[j].root. + AddBufferToColocatedSet( + conditional->branch_computation(j)->root_instruction(), + index, points_to_analysis, &colocated_set); + } AddSetToColocatedBufferSets(colocated_set, colocated_buffer_sets); }); - // Add true_operand and conditional.true_computation.parameter(0) as a - // colocated buffer set. Note that this has to be done for each subshape - // in the true_operand of the conditional. - ShapeUtil::ForEachSubshape( - conditional_hlo->operand(1)->shape(), - [this, conditional_hlo, &points_to_analysis, colocated_buffer_sets]( - const Shape& /*subshape*/, const ShapeIndex& index) { - std::vector true_set; - // Add conditional.true_operand. - AddBufferToColocatedSet(conditional_hlo->operand(1), index, - points_to_analysis, &true_set); - // Add conditional.true_computation.parameter_instruction(0). - AddBufferToColocatedSet( - conditional_hlo->true_computation()->parameter_instruction(0), - index, points_to_analysis, &true_set); - AddSetToColocatedBufferSets(true_set, colocated_buffer_sets); - }); - - // Add false_operand and conditional.false_computation.parameter(0) as a - // colocated buffer set. Note that this has to be done for each subshape - // in the false_operand of the conditional. - ShapeUtil::ForEachSubshape( - conditional_hlo->operand(2)->shape(), - [this, conditional_hlo, &points_to_analysis, colocated_buffer_sets]( - const Shape& /*subshape*/, const ShapeIndex& index) { - std::vector false_set; - // Add conditional.false_operand. - AddBufferToColocatedSet(conditional_hlo->operand(2), index, - points_to_analysis, &false_set); - // Add conditional.false_computation.parameter_instruction(0). - AddBufferToColocatedSet( - conditional_hlo->false_computation()->parameter_instruction( - 0), - index, points_to_analysis, &false_set); - AddSetToColocatedBufferSets(false_set, colocated_buffer_sets); - }); + for (int j = 0; j < conditional->branch_count(); ++j) { + // Add branch_operand[j] (which is operand[j+1]) and + // cond.branch_computation[j].parameter(0) as a colocated + // buffer set. Note that this has to be done for each subshape in the + // branch_operand of the case. + ShapeUtil::ForEachSubshape( + conditional->operand(j + 1)->shape(), + [this, j, conditional, &points_to_analysis, + colocated_buffer_sets](const Shape& /*subshape*/, + const ShapeIndex& index) { + std::vector branch_set; + // Add cond.operand[j+1]. + AddBufferToColocatedSet(conditional->operand(j + 1), index, + points_to_analysis, &branch_set); + // Add cond.branch_computation[j].parameter_instruction(0). + AddBufferToColocatedSet( + conditional->branch_computation(j)->parameter_instruction( + 0), + index, points_to_analysis, &branch_set); + AddSetToColocatedBufferSets(branch_set, colocated_buffer_sets); + }); + } } } } diff --git a/tensorflow/compiler/xla/client/lib/cholesky.cc b/tensorflow/compiler/xla/service/cholesky_expander.cc similarity index 76% rename from tensorflow/compiler/xla/client/lib/cholesky.cc rename to tensorflow/compiler/xla/service/cholesky_expander.cc index bb41f9932d1cc62b62d37fea2c10fbfeaa0bd15e..1c39cf9bc0a093ec54715d4180b49094ca6266a0 100644 --- a/tensorflow/compiler/xla/client/lib/cholesky.cc +++ b/tensorflow/compiler/xla/service/cholesky_expander.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/xla/client/lib/cholesky.h" +#include "tensorflow/compiler/xla/service/cholesky_expander.h" #include #include @@ -29,6 +29,7 @@ limitations under the License. #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/lib/core/errors.h" namespace xla { @@ -134,10 +135,8 @@ XlaOp CholeskyUnblocked(XlaOp a, PrecisionConfig::Precision precision) { }); } -} // namespace - -XlaOp Cholesky(XlaOp a, int64 block_size, - PrecisionConfig::Precision precision) { +XlaOp BuildCholesky(XlaOp a, int64 block_size, + PrecisionConfig::Precision precision) { XlaBuilder* builder = a.builder(); return builder->ReportErrorOrReturn([&]() -> StatusOr { TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); @@ -206,4 +205,55 @@ XlaOp Cholesky(XlaOp a, int64 block_size, }); } +} // namespace + +bool CholeskyExpander::InstructionMatchesPattern(HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kCholesky; +} + +StatusOr CholeskyExpander::ExpandInstruction( + HloInstruction* instruction) { + const CholeskyOptions& options = instruction->cholesky_options(); + const string name = absl::StrFormat( + "xla.cholesky_%s_%s", instruction->operand(0)->shape().ToString(), + options.lower() ? "lower" : "upper"); + + HloModule* module = instruction->parent()->parent(); + + HloComputation*& computation = + computation_cache_.emplace(name, nullptr).first->second; + if (!computation) { + // Builds a new expansion. + // + // TODO(b/62327888): We do something unusual here: we build the computation + // using the XlaBuilder API, which is nominally an XLA client API. We do + // this because the external APIs for building complicated computations + // (XlaBuilder) are much more ergonomic than the internal ones. As it turns + // out, XlaBuilder isn't really a client API—what it does is build a + // HloModuleProto protocol buffer, that we can then deserialize and clone + // into our HloModule. Ideally we would avoid the protocol buffer step; + // that is left as an exercise for future work. + XlaBuilder builder(name); + XlaOp a = Parameter(&builder, 0, instruction->operand(0)->shape(), "a"); + XlaOp l = BuildCholesky(MaybeTransposeInMinorDims(a, !options.lower()), + /*block_size=*/128, + /*precision=*/PrecisionConfig::HIGHEST); + MaybeTransposeInMinorDims(l, !options.lower()); + + TF_ASSIGN_OR_RETURN(XlaComputation xla_computation, builder.Build()); + + TF_ASSIGN_OR_RETURN(ProgramShape program_shape, + xla_computation.GetProgramShape()); + HloModuleConfig config(program_shape); + TF_ASSIGN_OR_RETURN(auto new_module, HloModule::CreateFromProto( + xla_computation.proto(), config)); + HloCloneContext context(module); + computation = + module->DeepCloneComputation(new_module->entry_computation(), &context); + } + + return instruction->parent()->AddInstruction(HloInstruction::CreateCall( + instruction->shape(), instruction->operands(), computation)); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/cholesky_expander.h b/tensorflow/compiler/xla/service/cholesky_expander.h new file mode 100644 index 0000000000000000000000000000000000000000..d2958db1b8ca676f3872016ac6a62b872a6b6649 --- /dev/null +++ b/tensorflow/compiler/xla/service/cholesky_expander.h @@ -0,0 +1,41 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CHOLESKY_EXPANDER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CHOLESKY_EXPANDER_H_ + +#include "absl/container/flat_hash_map.h" +#include "tensorflow/compiler/xla/service/op_expander_pass.h" + +namespace xla { + +class CholeskyExpander : public OpExpanderPass { + public: + absl::string_view name() const override { return "cholesky_expander"; } + + protected: + bool InstructionMatchesPattern(HloInstruction* instruction) override; + + StatusOr ExpandInstruction( + HloInstruction* instruction) override; + + private: + // Mapping from op signatures to existing computations. + absl::flat_hash_map computation_cache_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CHOLESKY_EXPANDER_H_ diff --git a/tensorflow/compiler/xla/service/conditional_simplifier.cc b/tensorflow/compiler/xla/service/conditional_simplifier.cc index 4ea3a13f2835c5fef99c274f14d7d683c9ff5fc8..f1d0ca44f08688ccda5b4385d65eabc0fc2fc5e6 100644 --- a/tensorflow/compiler/xla/service/conditional_simplifier.cc +++ b/tensorflow/compiler/xla/service/conditional_simplifier.cc @@ -33,8 +33,8 @@ limitations under the License. namespace xla { // Tries to replace a conditional with a call operation of the corresponding -// computation. If the given conditional has a constant predicate, tries to -// replace it with a call to its true/false computation as appropriate and then +// computation. If the given conditional has a constant branch_index, tries to +// replace it with a call to its corresponding branch computation and then // inline that computation. // // Returns true if it made a change to the graph. @@ -50,24 +50,30 @@ static StatusOr TryRemoveConditional(HloInstruction* conditional) { return false; } - if (conditional->operand(0)->opcode() != HloOpcode::kConstant) { - VLOG(2) << "Not attempting to remove conditional as its predicate is not a " - "compile-time constant: " - << conditional->ToShortString(); - return false; - } + // We can always inline a 1-branch conditional due to default branch fallback. + int branch_index = 0; + if (conditional->branch_count() > 1) { + if (conditional->operand(0)->opcode() != HloOpcode::kConstant) { + VLOG(2) << "Not attempting to remove conditional as its branch_index is " + "not a compile-time constant: " + << conditional->ToShortString(); + return false; + } + if (conditional->operand(0)->shape().element_type() == PRED) { + branch_index = conditional->operand(0)->literal().Get({}) ? 0 : 1; + } else { + branch_index = conditional->operand(0)->literal().Get({}); + if (branch_index < 0 || branch_index >= conditional->branch_count()) { + branch_index = conditional->branch_count() - 1; + } + } + } auto computation = conditional->parent(); HloInstruction* call_op; - if (conditional->operand(0)->literal().Get({})) { - call_op = computation->AddInstruction(HloInstruction::CreateCall( - conditional->shape(), {conditional->mutable_operand(1)}, - conditional->true_computation())); - } else { - call_op = computation->AddInstruction(HloInstruction::CreateCall( - conditional->shape(), {conditional->mutable_operand(2)}, - conditional->false_computation())); - } + call_op = computation->AddInstruction(HloInstruction::CreateCall( + conditional->shape(), {conditional->mutable_operand(branch_index + 1)}, + conditional->branch_computation(branch_index))); conditional->SetupDerivedInstruction(call_op); TF_RETURN_IF_ERROR(computation->ReplaceInstruction(conditional, call_op)); TF_RETURN_IF_ERROR(CallInliner::Inline(call_op).status()); diff --git a/tensorflow/compiler/xla/service/copy_insertion.cc b/tensorflow/compiler/xla/service/copy_insertion.cc index 5e26a63cebfa9b2e50f4b13335c10c246999d4df..fa0db40b065881e31eb0b04ecc3f99a03887caf5 100644 --- a/tensorflow/compiler/xla/service/copy_insertion.cc +++ b/tensorflow/compiler/xla/service/copy_insertion.cc @@ -319,8 +319,7 @@ Status AddCopiesForConditional(const HloAliasAnalysis& alias_analysis, << conditional->name(); TF_RET_CHECK(conditional->opcode() == HloOpcode::kConditional); - for (HloComputation* computation : - {conditional->true_computation(), conditional->false_computation()}) { + for (HloComputation* computation : conditional->branch_computations()) { HloInstruction* root = computation->root_instruction(); std::vector users = root->users(); TF_ASSIGN_OR_RETURN(HloInstruction * deep_copy, diff --git a/tensorflow/compiler/xla/service/cpu/BUILD b/tensorflow/compiler/xla/service/cpu/BUILD index 42672bc3875af2d732d80691df6bf85b9d8080cd..c8fef147b857aae1b15ff5fdc976a45d42cc6d12 100644 --- a/tensorflow/compiler/xla/service/cpu/BUILD +++ b/tensorflow/compiler/xla/service/cpu/BUILD @@ -111,6 +111,7 @@ cc_library( "//tensorflow/compiler/xla/service:buffer_assignment", "//tensorflow/compiler/xla/service:buffer_liveness", "//tensorflow/compiler/xla/service:call_inliner", + "//tensorflow/compiler/xla/service:cholesky_expander", "//tensorflow/compiler/xla/service:conditional_simplifier", "//tensorflow/compiler/xla/service:convolution_group_converter", "//tensorflow/compiler/xla/service:dot_decomposer", diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc index cc9489daaee2ea9127897a36c9cb868193ab7f8e..eb5d843fe8bf78dc040fd4d7e2c6b4b366762bc9 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc @@ -50,6 +50,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/buffer_liveness.h" #include "tensorflow/compiler/xla/service/call_inliner.h" +#include "tensorflow/compiler/xla/service/cholesky_expander.h" #include "tensorflow/compiler/xla/service/conditional_simplifier.h" #include "tensorflow/compiler/xla/service/convolution_group_converter.h" #include "tensorflow/compiler/xla/service/cpu/buffer_info_util.h" @@ -106,6 +107,7 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/platform/dynamic_annotations.h" namespace xla { namespace cpu { @@ -256,6 +258,7 @@ Status CpuCompiler::RunHloPassesThroughLayoutAssn( pipeline.AddPass(); + pipeline.AddPass(); pipeline.AddPass(); // TODO(b/65775800): Fix wrong output bug in Call and remove the CallInliner @@ -636,7 +639,13 @@ StatusOr> CpuCompiler::RunBackend( IrEmitter ir_emitter(*module, *assignment, llvm_module.get(), std::move(instruction_to_profile_idx), std::move(computation_to_profile_idx), - &target_machine_features); + &target_machine_features, +#ifdef MEMORY_SANITIZER + /*emit_code_for_msan=*/true +#else + /*emit_code_for_msan=*/false +#endif + ); TF_RETURN_IF_ERROR(ir_emitter.EmitConstantGlobals()); @@ -834,7 +843,9 @@ CpuCompiler::CompileAheadOfTime(std::unique_ptr module_group, IrEmitter ir_emitter(*module, *assignment, &llvm_module, std::move(instruction_to_profile_idx), std::move(computation_to_profile_idx), - &target_machine_features); + &target_machine_features, + // TODO(b/66051036): Run full msan for AOT. + /*emit_code_for_msan=*/false); TF_RETURN_IF_ERROR(ir_emitter.EmitConstantGlobals()); diff --git a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc index 6f79ad7c1468f27c74d84770ec6358fbcd1c1f09..2fbd60eb7262d403a11e9f036d832cb92d907dfa 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc @@ -42,9 +42,10 @@ bool CanBeLoopFused(const HloInstruction& hlo) { hlo.opcode() == HloOpcode::kTranspose; } -bool IsMatrixVectorDot(const HloInstruction* hlo) { +bool IsNonComplexMatrixVectorDot(const HloInstruction* hlo) { const Shape& hlo_shape = hlo->shape(); - return hlo->opcode() == HloOpcode::kDot && hlo_shape.dimensions_size() == 2 && + return !ShapeUtil::ElementIsComplex(hlo_shape) && + hlo->opcode() == HloOpcode::kDot && hlo_shape.dimensions_size() == 2 && (hlo_shape.dimensions(0) == 1 || hlo_shape.dimensions(1) == 1); } @@ -55,7 +56,8 @@ bool HasExactlyOneUse(const HloInstruction& hlo_instr) { bool CanBeOutputFused(const HloInstruction* producer, const HloInstruction* consumer) { - return consumer->opcode() == HloOpcode::kAdd && IsMatrixVectorDot(producer) && + return consumer->opcode() == HloOpcode::kAdd && + IsNonComplexMatrixVectorDot(producer) && HasExactlyOneUse(*producer) == 1; } diff --git a/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc index 0028fbaed895becad8da496aa8acdf7dc173a2a0..368ae8bffc548526101d75aaca3d7dafe21c7e18 100644 --- a/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc @@ -27,6 +27,8 @@ limitations under the License. #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" +using xla::llvm_ir::IrArray; + namespace xla { namespace cpu { @@ -105,21 +107,44 @@ StatusOr CpuElementalIrEmitter::EmitTanh(PrimitiveType prim_type, llvm_ir::ElementGenerator CpuElementalIrEmitter::MakeElementGenerator( const HloInstruction* hlo, const HloToElementGeneratorMap& operand_to_generator) { - if (hlo->opcode() == HloOpcode::kMap) { - return [this, hlo, &operand_to_generator]( - const llvm_ir::IrArray::Index& index) -> StatusOr { - std::vector operands; - for (int i = 0; i < hlo->operand_count(); i++) { - TF_ASSIGN_OR_RETURN(llvm::Value * operand_value, - operand_to_generator.at(hlo->operand(i))( - ElementwiseSourceIndex(index, *hlo, i))); - operands.push_back(operand_value); - } - return ir_emitter_->EmitElementalMap(*Cast(hlo), - operands, llvm_ir::IrName(hlo)); - }; + switch (hlo->opcode()) { + case HloOpcode::kMap: + return [this, hlo, &operand_to_generator]( + const IrArray::Index& index) -> StatusOr { + std::vector operands; + for (int i = 0; i < hlo->operand_count(); i++) { + TF_ASSIGN_OR_RETURN(llvm::Value * operand_value, + operand_to_generator.at(hlo->operand(i))( + ElementwiseSourceIndex(index, *hlo, i))); + operands.push_back(operand_value); + } + return ir_emitter_->EmitElementalMap(*Cast(hlo), + operands, llvm_ir::IrName(hlo)); + }; + case HloOpcode::kReduceWindow: + return [this, hlo, &operand_to_generator](const IrArray::Index& index) { + return ir_emitter_->EmitElementalReduceWindow( + Cast(hlo), + operand_to_generator.at(hlo->operand(0)), index); + }; + case HloOpcode::kConvolution: + return [this, hlo, &operand_to_generator](const IrArray::Index& index) { + return ir_emitter_->EmitElementalConvolution( + Cast(hlo), + operand_to_generator.at(hlo->operand(0)), + operand_to_generator.at(hlo->operand(1)), index); + }; + case HloOpcode::kReduce: + return [this, hlo, &operand_to_generator](const IrArray::Index& index) { + return ir_emitter_->EmitElementalReduce( + Cast(hlo), + operand_to_generator.at(hlo->operand(0)), + operand_to_generator.at(hlo->operand(1)), index); + }; + default: + return ElementalIrEmitter::MakeElementGenerator(hlo, + operand_to_generator); } - return ElementalIrEmitter::MakeElementGenerator(hlo, operand_to_generator); } } // namespace cpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/cpu/ir_emission_utils.cc index a8b139aec9e96b6bb580baf74789df7c998cebf8..2cc618e430215e26cb41c0a24a9c01b1ae33cec1 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emission_utils.cc +++ b/tensorflow/compiler/xla/service/cpu/ir_emission_utils.cc @@ -72,7 +72,8 @@ bool PotentiallyImplementedAsEigenConvolution( CHECK( ShapeUtil::SameElementTypeIgnoringFpPrecision(input_shape, kernel_shape)); // TODO(b/65408531): Explore using Eigen dot for complex64 type. - if (ShapeUtil::ElementIsComplex(input_shape)) { + PrimitiveType primitive_type = input_shape.element_type(); + if (primitive_type != F16 && primitive_type != F32) { return false; } if (window_util::HasWindowReversal(convolution.window())) { diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc index 9967cf28ee2389f9bef9780d2c986140f9bf2682..d75c33f486956cc521fa25525b395086595f289d 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc @@ -86,7 +86,8 @@ IrEmitter::IrEmitter( llvm::Module* llvm_module, std::unordered_map instruction_to_profile_idx, std::unordered_map computation_to_profile_idx, - const TargetMachineFeatures* target_machine_features) + const TargetMachineFeatures* target_machine_features, + bool emit_code_for_msan) : assignment_(assignment), module_(llvm_module), arch_type_(llvm::Triple(llvm_module->getTargetTriple()).getArch()), @@ -96,7 +97,8 @@ IrEmitter::IrEmitter( alias_analysis_(hlo_module, assignment, &llvm_module->getContext()), hlo_module_config_(hlo_module.config()), is_top_level_computation_(false), - target_machine_features_(*target_machine_features) { + target_machine_features_(*target_machine_features), + emit_code_for_msan_(emit_code_for_msan) { b_.setFastMathFlags(llvm_ir::GetFastMathFlags( /*fast_math_enabled=*/hlo_module_config_.debug_options() .xla_cpu_enable_fast_math())); @@ -642,8 +644,9 @@ llvm::Value* IrEmitter::EmitElementalMap( return EmitThreadLocalCall(*map_instr.to_apply(), elemental_operands, name); } -StatusOr IrEmitter::EmitTargetElementLoopBodyForReduceWindow( - HloReduceWindowInstruction* reduce_window, +StatusOr IrEmitter::EmitElementalReduceWindow( + const HloReduceWindowInstruction* reduce_window, + const llvm_ir::ElementGenerator& input_generator, const llvm_ir::IrArray::Index& index) { const HloInstruction* operand = reduce_window->operand(0); const Window& window = reduce_window->window(); @@ -714,8 +717,8 @@ StatusOr IrEmitter::EmitTargetElementLoopBodyForReduceWindow( SetToFirstInsertPoint(if_data.true_block, &b_); // We are not in the padding, so carry out the computation. - llvm_ir::IrArray input_array(GetIrArrayFor(operand)); - llvm::Value* input_value = input_array.EmitReadArrayElement(input_index, &b_); + TF_ASSIGN_OR_RETURN(llvm::Value* const input_value, + input_generator(input_index)); llvm::Value* result = EmitThreadLocalCall( *reduce_window->to_apply(), {Load(accumulator_address), input_value}, "reducer_function"); @@ -739,11 +742,7 @@ Status IrEmitter::HandleReduceWindow(HloInstruction* reduce_window) { // // This is completely un-optimized and just here to have something // that works. - return EmitTargetElementLoop( - reduce_window, [&](const llvm_ir::IrArray::Index& index) { - return EmitTargetElementLoopBodyForReduceWindow( - Cast(reduce_window), index); - }); + return DefaultAction(reduce_window); } Status IrEmitter::HandleSelectAndScatter(HloInstruction* select_and_scatter) { @@ -946,8 +945,10 @@ Status IrEmitter::HandleDot(HloInstruction* dot) { hlo_module_config_, target_machine_features_); } -StatusOr IrEmitter::EmitTargetElementLoopBodyForConvolution( - HloConvolutionInstruction* convolution, +StatusOr IrEmitter::EmitElementalConvolution( + const HloConvolutionInstruction* convolution, + const llvm_ir::ElementGenerator& input_generator, + const llvm_ir::ElementGenerator& kernel_generator, const llvm_ir::IrArray::Index& index) { const HloInstruction* lhs = convolution->operand(0); const HloInstruction* rhs = convolution->operand(1); @@ -1059,7 +1060,6 @@ StatusOr IrEmitter::EmitTargetElementLoopBodyForConvolution( input_index[dnums.input_feature_dimension()] = input_feature; input_index[dnums.input_batch_dimension()] = batch; - llvm_ir::IrArray kernel_array(GetIrArrayFor(rhs)); llvm_ir::IrArray::Index kernel_index(b_.getInt64Ty(), num_dims); for (int i = 0; i < num_spatial_dims; ++i) { kernel_index[dnums.kernel_spatial_dimensions(i)] = @@ -1072,10 +1072,11 @@ StatusOr IrEmitter::EmitTargetElementLoopBodyForConvolution( kernel_index[dnums.kernel_input_feature_dimension()] = input_feature; kernel_index[dnums.kernel_output_feature_dimension()] = output_feature; - llvm_ir::IrArray input_array(GetIrArrayFor(lhs)); - llvm::Value* product = - FMul(input_array.EmitReadArrayElement(input_index, &b_), - kernel_array.EmitReadArrayElement(kernel_index, &b_)); + TF_ASSIGN_OR_RETURN(llvm::Value* const input_value, + input_generator(input_index)); + TF_ASSIGN_OR_RETURN(llvm::Value* const kernel_value, + kernel_generator(kernel_index)); + llvm::Value* product = FMul(input_value, kernel_value); llvm::Value* sum = FAdd(Load(sum_address), product); Store(sum, sum_address); @@ -1088,7 +1089,7 @@ Status IrEmitter::HandleConvolution(HloInstruction* convolution) { auto rhs = convolution->operand(1); TF_RETURN_IF_ERROR(ElementTypesSameAndSupported( /*instruction=*/*convolution, /*operands=*/{lhs, rhs}, - /*supported_types=*/{F16, F32, C64, C128})); + /*supported_types=*/{F16, F32, F64, C64, C128})); // TODO(tonywy): Add PotentiallyImplementedAsMKLCovolution to support // different data layouts. @@ -1243,11 +1244,7 @@ Status IrEmitter::HandleConvolution(HloInstruction* convolution) { // // See the description of convolution in the XLA documentation for the pseudo // code for convolution. - return EmitTargetElementLoop( - convolution, [&](const llvm_ir::IrArray::Index& index) { - return EmitTargetElementLoopBodyForConvolution( - Cast(convolution), index); - }); + return DefaultAction(convolution); } Status IrEmitter::HandleFft(HloInstruction* fft) { @@ -1803,10 +1800,12 @@ StatusOr IrEmitter::EmitVectorizedReduce( return true; } -StatusOr IrEmitter::EmitTargetElementLoopBodyForReduce( - HloReduceInstruction* reduce, const llvm_ir::IrArray::Index& index) { - const HloInstruction* arg = reduce->mutable_operand(0); - const HloInstruction* init_value = reduce->mutable_operand(1); +StatusOr IrEmitter::EmitElementalReduce( + const HloReduceInstruction* reduce, + const llvm_ir::ElementGenerator& input_generator, + const llvm_ir::ElementGenerator& initial_value_generator, + const llvm_ir::IrArray::Index& index) { + const HloInstruction* arg = reduce->operand(0); absl::Span dimensions(reduce->dimensions()); // Initialize an accumulator with init_value. @@ -1814,9 +1813,10 @@ StatusOr IrEmitter::EmitTargetElementLoopBodyForReduce( llvm::AllocaInst* accumulator_addr = llvm_ir::EmitAllocaAtFunctionEntry( llvm_ir::PrimitiveTypeToIrType(accumulator_type, module_), "accumulator", &b_, MinimumAlignmentForPrimitiveType(accumulator_type)); - llvm::Value* init_value_addr = GetEmittedValueFor(init_value); - llvm::Value* load_init_value = Load(init_value_addr); - Store(load_init_value, accumulator_addr); + TF_ASSIGN_OR_RETURN( + llvm::Value* const init_value, + initial_value_generator(llvm_ir::IrArray::Index(index.GetType()))); + Store(init_value, accumulator_addr); // The enclosing loops go over all the target elements. Now we have to compute // the actual target element. For this, we build a new loop nest to iterate @@ -1835,7 +1835,6 @@ StatusOr IrEmitter::EmitTargetElementLoopBodyForReduce( // fill in the rest of the dimensions with induction Value*s taken from // 'index' which iterates over the target array. See the high-level // description in the XLA documentation for details. - llvm_ir::IrArray arg_array(GetIrArrayFor(arg)); llvm_ir::IrArray::Index input_index = reduced_dims_index; llvm_ir::IrArray::Index::const_iterator it = index.begin(); @@ -1847,7 +1846,8 @@ StatusOr IrEmitter::EmitTargetElementLoopBodyForReduce( CHECK(index.end() == it); // Apply the reduction function to the loaded value. - llvm::Value* input_element = arg_array.EmitReadArrayElement(input_index, &b_); + TF_ASSIGN_OR_RETURN(llvm::Value* const input_element, + input_generator(input_index)); llvm::Value* result = EmitThreadLocalCall( *reduce->to_apply(), {Load(accumulator_addr), input_element}, "reduce_function"); @@ -1882,11 +1882,11 @@ Status IrEmitter::HandleReduce(HloInstruction* reduce) { } } - return EmitTargetElementLoop(reduce, - [&](const llvm_ir::IrArray::Index& index) { - return EmitTargetElementLoopBodyForReduce( - Cast(reduce), index); - }); + return DefaultAction(reduce); +} + +Status IrEmitter::HandleAllToAll(HloInstruction*) { + return Unimplemented("AllToAll is not implemented on CPU."); } Status IrEmitter::HandleSend(HloInstruction* send) { @@ -2225,6 +2225,25 @@ Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) { InBoundsGEP(operands_alloca, {b_.getInt64(i)}); Store(operand_as_i8ptr, slot_in_operands_alloca); } + if (emit_code_for_msan_) { + // Mark the alloca as initialized for msan. The buffer gets read by the + // custom callee, which might be msan-instrumented. + // TODO(b/66051036): Run the msan instrumentation pass instead. + const llvm::DataLayout& dl = module_->getDataLayout(); + llvm::Type* intptr_type = b_.getIntPtrTy(dl); + auto* msan_unpoison_ir_function = llvm::cast( + module_ + ->getOrInsertFunction( + "__msan_unpoison", + llvm::FunctionType::get( + /*Result=*/b_.getVoidTy(), + /*Params=*/{i8_ptr_type, intptr_type}, /*isVarArg=*/false)) + .getCallee()); + Call(msan_unpoison_ir_function, + {PointerCast(operands_alloca, i8_ptr_type), + llvm::ConstantInt::get( + intptr_type, *operands_alloca->getAllocationSizeInBits(dl) / 8)}); + } auto* custom_call_ir_function = llvm::dyn_cast( module_ ->getOrInsertFunction( @@ -2494,53 +2513,109 @@ Status IrEmitter::HandleConcatenate(HloInstruction* concatenate) { } Status IrEmitter::HandleConditional(HloInstruction* conditional) { - auto pred = conditional->operand(0); - TF_RET_CHECK(ShapeUtil::IsScalar(pred->shape()) && - pred->shape().element_type() == PRED) - << "Predicate on a Conditional must be bool; got: " - << ShapeUtil::HumanString(pred->shape()); - - HloComputation* true_computation = conditional->true_computation(); - HloComputation* false_computation = conditional->false_computation(); - TF_RET_CHECK(ShapeUtil::Equal(conditional->shape(), - true_computation->root_instruction()->shape())) - << "Shape of conditional should be same as the shape of the true " - << "computation; got: " << ShapeUtil::HumanString(conditional->shape()) - << " and " - << ShapeUtil::HumanString(true_computation->root_instruction()->shape()); - - TF_RET_CHECK(ShapeUtil::Equal(conditional->shape(), - false_computation->root_instruction()->shape())) - << "Shape of conditional should be same as the shape of the false " - << "computation; got: " << ShapeUtil::HumanString(conditional->shape()) - << " and " - << ShapeUtil::HumanString(false_computation->root_instruction()->shape()); + auto branch_index = conditional->operand(0); + int num_branches = conditional->branch_count(); + TF_RET_CHECK(ShapeUtil::IsScalar(branch_index->shape()) && + (branch_index->shape().element_type() == PRED || + branch_index->shape().element_type() == S32)) + << "Branch index on a conditional must be scalar bool or int32; got: " + << ShapeUtil::HumanString(branch_index->shape()); + + for (int b = 0; b < num_branches; ++b) { + HloComputation* br_computation = conditional->branch_computation(b); + TF_RET_CHECK(ShapeUtil::Equal(conditional->shape(), + br_computation->root_instruction()->shape())) + << "Shape of conditional should be same as the shape of the " << b + << "th branch computation; got: " + << ShapeUtil::HumanString(conditional->shape()) << " and " + << ShapeUtil::HumanString(br_computation->root_instruction()->shape()); + } TF_RETURN_IF_ERROR(EmitTargetAddressForOp(conditional)); - // Generating: - // if (pred) - // cond_result = true_computation(true_operand) - // else - // cond_result = false_computation(false_operand) - llvm::LoadInst* pred_value = - Load(GetIrArrayFor(pred).GetBasePointer(), "load_predicate_value"); - llvm::Value* pred_cond = ICmpNE( - pred_value, - llvm::ConstantInt::get(llvm_ir::PrimitiveTypeToIrType(PRED, module_), 0), - "boolean_predicate"); - llvm_ir::LlvmIfData if_data = - llvm_ir::EmitIfThenElse(pred_cond, "conditional", &b_); - - SetToFirstInsertPoint(if_data.true_block, &b_); - EmitGlobalCall(*conditional->true_computation(), - IrName(conditional, "_true")); - - SetToFirstInsertPoint(if_data.false_block, &b_); - EmitGlobalCall(*conditional->false_computation(), - IrName(conditional, "_false")); - - SetToFirstInsertPoint(if_data.after_block, &b_); + if (branch_index->shape().element_type() == PRED) { + // Emit an if-else to LLVM: + // if (pred) + // cond_result = true_computation(true_operand) + // else + // cond_result = false_computation(false_operand) + llvm::LoadInst* pred_value = Load( + GetIrArrayFor(branch_index).GetBasePointer(), "load_predicate_value"); + llvm::Value* pred_cond = + ICmpNE(pred_value, + llvm::ConstantInt::get( + llvm_ir::PrimitiveTypeToIrType(PRED, module_), 0), + "boolean_predicate"); + llvm_ir::LlvmIfData if_data = + llvm_ir::EmitIfThenElse(pred_cond, "conditional", &b_); + + SetToFirstInsertPoint(if_data.true_block, &b_); + EmitGlobalCall(*conditional->branch_computation(0), + IrName(conditional, "_true")); + + SetToFirstInsertPoint(if_data.false_block, &b_); + EmitGlobalCall(*conditional->branch_computation(1), + IrName(conditional, "_false")); + + SetToFirstInsertPoint(if_data.after_block, &b_); + return Status::OK(); + } + // We emit a switch statement to LLVM: + // switch (branch_index) { + // default: + // result = branch_computations[num_branches-1](operands[num_branches-1]); + // break; + // case 0: + // result = branch_computations[0](operands[0]); break; + // case 1: + // result = branch_computations[1](operands[1]); break; + // ... + // case [[num_branches-2]]: + // result = branch_computations[num_branches-2](operands[num_branches-2]); + // break; + // } + llvm::LoadInst* branch_index_value = Load( + GetIrArrayFor(branch_index).GetBasePointer(), "load_branch_index_value"); + + auto case_block = b_.GetInsertBlock(); + llvm::BasicBlock* after_block; + // Add a terminator to the case block, if necessary. + if (case_block->getTerminator() == nullptr) { + after_block = llvm_ir::CreateBasicBlock(nullptr, "case-after", &b_); + b_.SetInsertPoint(case_block); + b_.CreateBr(after_block); + } else { + after_block = + case_block->splitBasicBlock(b_.GetInsertPoint(), "case-after"); + } + // Our basic block should now end with an unconditional branch. Remove it; + // we're going to replace it with a switch based branch. + case_block->getTerminator()->eraseFromParent(); + + // Lower the default branch computation. + auto default_block = llvm_ir::CreateBasicBlock(nullptr, "case-default", &b_); + b_.SetInsertPoint(default_block); + EmitGlobalCall(*conditional->branch_computation(num_branches - 1), + IrName(conditional, "_default")); + b_.CreateBr(after_block); + + // Prepare the switch (branch_index) { ... } instruction. + b_.SetInsertPoint(case_block); + llvm::SwitchInst* case_inst = + b_.CreateSwitch(branch_index_value, default_block, num_branches - 1); + // Lower each branch's computation. + for (int b = 0; b < num_branches - 1; ++b) { // last branch is default + // Lower the case b: { ... ; break; } computation. + auto branch_block = + llvm_ir::CreateBasicBlock(nullptr, absl::StrCat("case-branch", b), &b_); + b_.SetInsertPoint(branch_block); + EmitGlobalCall(*conditional->branch_computation(b), + IrName(conditional, absl::StrCat("_branch", b))); + b_.CreateBr(after_block); + case_inst->addCase(b_.getInt32(b), branch_block); + } + + SetToFirstInsertPoint(after_block, &b_); return Status::OK(); } diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.h b/tensorflow/compiler/xla/service/cpu/ir_emitter.h index 974dd7cd3f2254bfbc86fffae02c06c481af8902..e183ae01070e7d42701a3a32d5ddb8667e163663 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.h +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.h @@ -72,13 +72,15 @@ class IrEmitter : public DfsHloVisitorWithDefault, // index in the profiling array. // computation_to_profile_idx: the mapping from HLO computations to their // index in the profiling array. + // emit_code_for_msan: whether emitted code should be compatible with msan. IrEmitter(const HloModule& hlo_module, const BufferAssignment& assignment, llvm::Module* llvm_module, std::unordered_map instruction_to_profile_idx, std::unordered_map computation_to_profile_idx, - const TargetMachineFeatures* target_machine); + const TargetMachineFeatures* target_machine, + bool emit_code_for_msan); ~IrEmitter() override; // Emit and return the given HLO computation as an LLVM IR @@ -116,6 +118,23 @@ class IrEmitter : public DfsHloVisitorWithDefault, const HloMapInstruction& map_instr, absl::Span elemental_operands, absl::string_view name); + // Emit code to emit the element at `index` for a reduce window instruction. + StatusOr EmitElementalReduceWindow( + const HloReduceWindowInstruction* reduce_window, + const llvm_ir::ElementGenerator& input_generator, + const llvm_ir::IrArray::Index& index); + // Emit code to emit the element at `index` for a convolution instruction. + StatusOr EmitElementalConvolution( + const HloConvolutionInstruction* convolution, + const llvm_ir::ElementGenerator& input_generator, + const llvm_ir::ElementGenerator& kernel_generator, + const llvm_ir::IrArray::Index& index); + // Emit code to emit the element at `index` for a reduce instruction. + StatusOr EmitElementalReduce( + const HloReduceInstruction* reduce, + const llvm_ir::ElementGenerator& input_generator, + const llvm_ir::ElementGenerator& initial_value_generator, + const llvm_ir::IrArray::Index& index); protected: // @@ -125,6 +144,7 @@ class IrEmitter : public DfsHloVisitorWithDefault, // special in some way are handled explicitly in HandleFoo methods. Status DefaultAction(HloInstruction* hlo) override; + Status HandleAllToAll(HloInstruction* instruction) override; Status HandleBitcast(HloInstruction* bitcast) override; Status HandleConstant(HloInstruction* constant) override; Status HandleCopy(HloInstruction* copy) override; @@ -524,17 +544,6 @@ class IrEmitter : public DfsHloVisitorWithDefault, // Returns the number of bytes within the shape. int64 ByteSizeOf(const Shape& shape) const; - StatusOr EmitTargetElementLoopBodyForMap( - HloMapInstruction* map, const llvm_ir::IrArray::Index& index); - StatusOr EmitTargetElementLoopBodyForReduceWindow( - HloReduceWindowInstruction* reduce_window, - const llvm_ir::IrArray::Index& index); - StatusOr EmitTargetElementLoopBodyForConvolution( - HloConvolutionInstruction* convolution, - const llvm_ir::IrArray::Index& index); - StatusOr EmitTargetElementLoopBodyForReduce( - HloReduceInstruction* reduce, const llvm_ir::IrArray::Index& index); - enum class XfeedKind { kInfeed, kOutfeed, @@ -574,6 +583,8 @@ class IrEmitter : public DfsHloVisitorWithDefault, std::vector thread_local_computations_; std::vector global_computations_; + bool emit_code_for_msan_; + TF_DISALLOW_COPY_AND_ASSIGN(IrEmitter); }; diff --git a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc index 9c2685674fbc133de1220caef81ac3b60a1c0f7c..f7b64738b7b314b56f4ae60336d9c85c90287219 100644 --- a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc +++ b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc @@ -337,6 +337,11 @@ bool RegisterKnownJITSymbols() { reinterpret_cast(memset_pattern16)); #endif +#ifdef MEMORY_SANITIZER + registry->Register("__msan_unpoison", + reinterpret_cast(__msan_unpoison)); +#endif + return true; } diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h index af039776b7f157a407f8f4cf3d7cabdbee45b014..246f2af09b5539612ef0e75929833f532dfa4083 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h @@ -103,9 +103,16 @@ class DfsHloVisitorBase { virtual Status HandlePower(HloInstructionPtr hlo) { return HandleElementwiseBinary(hlo); } + virtual Status HandleSqrt(HloInstructionPtr hlo) { + return HandleElementwiseUnary(hlo); + } + virtual Status HandleRsqrt(HloInstructionPtr hlo) { + return HandleElementwiseUnary(hlo); + } virtual Status HandleConvolution(HloInstructionPtr hlo) = 0; virtual Status HandleFft(HloInstructionPtr fft) = 0; virtual Status HandleTriangularSolve(HloInstructionPtr hlo) = 0; + virtual Status HandleCholesky(HloInstructionPtr hlo) = 0; virtual Status HandleAllReduce(HloInstructionPtr hlo) = 0; virtual Status HandleAllToAll(HloInstructionPtr hlo) = 0; virtual Status HandleCollectivePermute(HloInstructionPtr hlo) = 0; diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h b/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h index 341bb37b8355e9987a0331d0a66bb8fe87f019cf..79ce3f82e8c1fe91d590ea7c47fa219ce8e8a80f 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h @@ -94,6 +94,9 @@ class DfsHloVisitorWithDefaultBase Status HandleTriangularSolve(HloInstructionPtr hlo) override { return DefaultAction(hlo); } + Status HandleCholesky(HloInstructionPtr hlo) override { + return DefaultAction(hlo); + } Status HandleAllReduce(HloInstructionPtr crs) override { return DefaultAction(crs); } diff --git a/tensorflow/compiler/xla/service/dynamic_padder.cc b/tensorflow/compiler/xla/service/dynamic_padder.cc index 4db280f817141bd52e3a5b9564600a618f81aeac..4f732a2546e7f09002186347a56d5456bdb2ed5e 100644 --- a/tensorflow/compiler/xla/service/dynamic_padder.cc +++ b/tensorflow/compiler/xla/service/dynamic_padder.cc @@ -80,6 +80,22 @@ StatusOr ChooseIdentityValue(HloInstruction* inst) { } } +bool ShouldSkipPadOnOperand(const HloInstruction* inst, int64 operand_num, + int64 dimension) { + if ((inst->opcode() == HloOpcode::kReduceWindow || + inst->opcode() == HloOpcode::kSelectAndScatter) && + operand_num == 0 && inst->window().dimensions(dimension).size() == 1) { + return true; + } + + if (operand_num == 0 && inst->opcode() == HloOpcode::kConvolution && + inst->convolution_dimension_numbers().input_batch_dimension() == + dimension) { + return true; + } + return false; +} + } // namespace StatusOr DynamicPadder::Run(HloModule* module) { @@ -105,6 +121,11 @@ StatusOr DynamicPadder::Run(HloModule* module) { } VLOG(1) << "Has dynamic dimension of operand" << operand_num << " @" << dim; + + if (ShouldSkipPadOnOperand(inst, operand_num, dim)) { + continue; + } + TF_ASSIGN_OR_RETURN(HloInstruction * identity_value, ChooseIdentityValue(inst)); if (identity_value == nullptr) { diff --git a/tensorflow/compiler/xla/service/dynamic_padder_test.cc b/tensorflow/compiler/xla/service/dynamic_padder_test.cc index 55a11286e4596d87c330315322cae704fc5cd707..fda806bbf81519cfa0babb240fcde517dedcf284 100644 --- a/tensorflow/compiler/xla/service/dynamic_padder_test.cc +++ b/tensorflow/compiler/xla/service/dynamic_padder_test.cc @@ -22,8 +22,10 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/hlo_parser.h" #include "tensorflow/compiler/xla/service/hlo_runner.h" #include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" @@ -133,19 +135,84 @@ TEST_F(DynamicPadderTest, ConvolutionTest) { module_->AddEntryComputation(builder.Build()); + // Set up binding for contracting dimensions. + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{2, {}}, + DynamicParameterBinding::DynamicDimension{0, {}, 1})); + + TF_ASSERT_OK(RunPadder().status()); + + ExpectPadded(conv->operand(0)); +} + +TEST_F(DynamicPadderTest, ConvolutionNoPad) { + auto builder = HloComputation::Builder(TestName()); + constexpr int xdim = 3; + constexpr int ydim = 2; + constexpr int zdim = 1; + auto xy_shape = ShapeUtil::MakeShape(F32, {xdim, ydim}); + auto yz_shape = ShapeUtil::MakeShape(F32, {ydim, zdim}); + auto zx_shape = ShapeUtil::MakeShape(F32, {zdim, xdim}); + + auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/0, xy_shape, "A")); + auto* b_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/1, yz_shape, "B")); + builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/2, scalar_shape_, "size_param")); + + auto dnums = XlaBuilder::CreateDefaultConvDimensionNumbers(0); + + dnums.set_kernel_input_feature_dimension(0); + dnums.set_kernel_output_feature_dimension(1); + dnums.set_input_batch_dimension(0); + dnums.set_output_batch_dimension(1); + dnums.set_output_feature_dimension(0); + + Window window; + + auto* conv = builder.AddInstruction(HloInstruction::CreateConvolve( + zx_shape, a_param, b_param, /*feature_group_count=*/1, + /*batch_group_count=*/1, window, dnums, + HloTestBase::DefaultPrecisionConfig(2))); + + module_->AddEntryComputation(builder.Build()); + // Set up dynamic parameter binding for non-contracting dimension. TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( DynamicParameterBinding::DynamicParameter{2, {}}, DynamicParameterBinding::DynamicDimension{0, {}, 0})); - // Set up binding for contracting dimensions. + TF_ASSERT_OK(RunPadder().status()); + + EXPECT_THAT(conv->operand(0), op::Parameter()); +} + +TEST_F(DynamicPadderTest, ReduceWindowNoPadForTrivialWindow) { + auto builder = HloComputation::Builder(TestName()); + auto input_shape = ShapeUtil::MakeShape(F32, {4, 5}); + auto reduce_shape = ShapeUtil::MakeShape(F32, {3, 5}); + + auto input = builder.AddInstruction( + HloInstruction::CreateParameter(0, input_shape, "input")); + builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "size_param")); + auto init = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0))); + TF_ASSERT_OK_AND_ASSIGN(Window window, ParseWindow("size=2x1 pad=0_0x0_0")); + auto output = builder.AddInstruction(HloInstruction::CreateReduceWindow( + reduce_shape, input, init, window, GetScalarAddComputation())); + + module_->AddEntryComputation(builder.Build()); + + // Set up dynamic parameter binding. TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( - DynamicParameterBinding::DynamicParameter{2, {}}, + DynamicParameterBinding::DynamicParameter{1, {}}, DynamicParameterBinding::DynamicDimension{0, {}, 1})); TF_ASSERT_OK(RunPadder().status()); - ExpectPadded(conv->operand(0)); + EXPECT_THAT(output->operand(0), op::Parameter()); } } // namespace diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc index 808929be75ec6fd0cfb15418a231431b8d51e089..abd8ead52abd89541f375d9fd2e81158c78bda03 100644 --- a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc @@ -423,6 +423,10 @@ StatusOr ElementalIrEmitter::EmitFloatUnaryOp( return EmitSin(op->shape().element_type(), operand_value); case HloOpcode::kTanh: return EmitTanh(op->shape().element_type(), operand_value); + case HloOpcode::kSqrt: + return EmitSqrt(op->shape().element_type(), operand_value); + case HloOpcode::kRsqrt: + return EmitRsqrt(op->shape().element_type(), operand_value); case HloOpcode::kFloor: return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::floor, {operand_value}, @@ -436,9 +440,7 @@ StatusOr ElementalIrEmitter::EmitFloatUnaryOp( {operand_value}, {operand_value->getType()}, b_); case HloOpcode::kRoundNearestAfz: - return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::round, - {operand_value}, - {operand_value->getType()}, b_); + return EmitRoundNearestAfz(op->shape().element_type(), operand_value); case HloOpcode::kSign: { auto type = operand_value->getType(); auto zero = llvm::ConstantFP::get(type, 0.0); @@ -655,6 +657,20 @@ StatusOr ElementalIrEmitter::EmitComplexUnaryOp( EmitComposeComplex(op, FDiv(EmitExtractReal(operand_value), cplx_abs), FDiv(EmitExtractImag(operand_value), cplx_abs))); } + case HloOpcode::kSqrt: { + auto a = EmitExtractReal(operand_value); + auto b = EmitExtractImag(operand_value); + auto c = llvm::ConstantFP::get(a->getType(), 0.5); + auto d = llvm::ConstantFP::get(b->getType(), 0.0); + return EmitComplexPower(op, a, b, c, d); + } + case HloOpcode::kRsqrt: { + auto a = EmitExtractReal(operand_value); + auto b = EmitExtractImag(operand_value); + auto c = llvm::ConstantFP::get(a->getType(), -0.5); + auto d = llvm::ConstantFP::get(b->getType(), 0.0); + return EmitComplexPower(op, a, b, c, d); + } case HloOpcode::kNegate: return EmitComposeComplex(op, FNeg(EmitExtractReal(operand_value)), FNeg(EmitExtractImag(operand_value))); @@ -738,6 +754,43 @@ StatusOr ElementalIrEmitter::EmitFloatBinaryOp( } } +// (a+bi)^(c+di) = +// (a*a+b*b)^(0.5c) * exp(-d*atan2(b,a)) * (cos(q) + i*sin(q)), +// where q = c*atan2(b,a)+0.5d*ln(a*a+b*b) +StatusOr ElementalIrEmitter::EmitComplexPower( + const HloInstruction* op, llvm::Value* a, llvm::Value* b, llvm::Value* c, + llvm::Value* d) { + PrimitiveType component_type = + primitive_util::ComplexComponentType(op->shape().element_type()); + auto aa_p_bb = FAdd(FMul(a, a), FMul(b, b)); + auto zero = llvm::ConstantFP::get(a->getType(), 0); + auto one_half = llvm::ConstantFP::get(a->getType(), 0.5); + auto one = llvm::ConstantFP::get(a->getType(), 1); + auto half_c = FMul(one_half, c); + + TF_ASSIGN_OR_RETURN(auto aa_p_bb_to_half_c, + EmitPow(component_type, aa_p_bb, half_c)); + + auto neg_d = FNeg(d); + TF_ASSIGN_OR_RETURN(auto arg_lhs, EmitAtan2(component_type, b, a)); + auto neg_d_arg_lhs = FMul(neg_d, arg_lhs); + TF_ASSIGN_OR_RETURN(auto e_to_neg_d_arg_lhs, + EmitExp(component_type, neg_d_arg_lhs)); + auto coeff = FMul(aa_p_bb_to_half_c, e_to_neg_d_arg_lhs); + TF_ASSIGN_OR_RETURN(auto ln_aa_p_bb, EmitLog(component_type, aa_p_bb)); + auto half_d = FMul(one_half, d); + auto q = FAdd(FMul(c, arg_lhs), FMul(half_d, ln_aa_p_bb)); + TF_ASSIGN_OR_RETURN(auto cos_q, EmitCos(component_type, q)); + TF_ASSIGN_OR_RETURN(auto sin_q, EmitSin(component_type, q)); + // d^c is 0 if d is 0 and c > 0. 0^0 is defined to be 1.0, see + // Branch Cuts for Complex Elementary Functions or Much Ado About + // Nothing's Sign Bit, W. Kahan, Section 10. + return Select( + And(And(FCmpOEQ(aa_p_bb, zero), FCmpOEQ(d, zero)), FCmpOLE(zero, c)), + EmitComposeComplex(op, Select(FCmpOEQ(zero, c), one, zero), zero), + EmitComposeComplex(op, FMul(coeff, cos_q), FMul(coeff, sin_q))); +} + StatusOr ElementalIrEmitter::EmitComplexBinaryOp( const HloInstruction* op, llvm::Value* lhs_value, llvm::Value* rhs_value) { switch (op->opcode()) { @@ -804,42 +857,11 @@ StatusOr ElementalIrEmitter::EmitComplexBinaryOp( EmitExtractImag(rhs_value), b_)); case HloOpcode::kPower: { - // (a+bi)^(c+di) = - // (a*a+b*b)^(0.5c) * exp(-d*atan2(b,a)) * (cos(q) + i*sin(q)), - // where q = c*atan2(b,a)+0.5d*ln(a*a+b*b) - PrimitiveType component_type = - primitive_util::ComplexComponentType(op->shape().element_type()); auto a = EmitExtractReal(lhs_value); auto b = EmitExtractImag(lhs_value); auto c = EmitExtractReal(rhs_value); auto d = EmitExtractImag(rhs_value); - auto aa_p_bb = FAdd(FMul(a, a), FMul(b, b)); - auto zero = llvm::ConstantFP::get(a->getType(), 0); - auto one_half = llvm::ConstantFP::get(a->getType(), 0.5); - auto one = llvm::ConstantFP::get(a->getType(), 1); - auto half_c = FMul(one_half, c); - - TF_ASSIGN_OR_RETURN(auto aa_p_bb_to_half_c, - EmitPow(component_type, aa_p_bb, half_c)); - - auto neg_d = FNeg(d); - TF_ASSIGN_OR_RETURN(auto arg_lhs, EmitAtan2(component_type, b, a)); - auto neg_d_arg_lhs = FMul(neg_d, arg_lhs); - TF_ASSIGN_OR_RETURN(auto e_to_neg_d_arg_lhs, - EmitExp(component_type, neg_d_arg_lhs)); - auto coeff = FMul(aa_p_bb_to_half_c, e_to_neg_d_arg_lhs); - TF_ASSIGN_OR_RETURN(auto ln_aa_p_bb, EmitLog(component_type, aa_p_bb)); - auto half_d = FMul(one_half, d); - auto q = FAdd(FMul(c, arg_lhs), FMul(half_d, ln_aa_p_bb)); - TF_ASSIGN_OR_RETURN(auto cos_q, EmitCos(component_type, q)); - TF_ASSIGN_OR_RETURN(auto sin_q, EmitSin(component_type, q)); - // 0^c is 0 if d is 0 and c > 0. 0^0 is defined to be 1.0, see - // Branch Cuts for Complex Elementary Functions or Much Ado About - // Nothing's Sign Bit, W. Kahan, Section 10. - return Select( - And(And(FCmpOEQ(aa_p_bb, zero), FCmpOEQ(d, zero)), FCmpOLE(zero, c)), - EmitComposeComplex(op, Select(FCmpOEQ(zero, c), one, zero), zero), - EmitComposeComplex(op, FMul(coeff, cos_q), FMul(coeff, sin_q))); + return EmitComplexPower(op, a, b, c, d); } default: return Unimplemented("binary complex op '%s'", @@ -1052,6 +1074,18 @@ StatusOr ElementalIrEmitter::EmitLog1p(PrimitiveType prim_type, return Select(x_is_small, for_small_x, for_large_x); } +StatusOr ElementalIrEmitter::EmitSqrt(PrimitiveType prim_type, + llvm::Value* value) { + return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::sqrt, {value}, + {value->getType()}, b_); +} + +StatusOr ElementalIrEmitter::EmitRsqrt(PrimitiveType prim_type, + llvm::Value* value) { + TF_ASSIGN_OR_RETURN(auto sqrt, EmitSqrt(prim_type, value)); + return FDiv(llvm::ConstantFP::get(sqrt->getType(), 1.0), sqrt); +} + StatusOr ElementalIrEmitter::EmitSin(PrimitiveType prim_type, llvm::Value* value) { return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::sin, {value}, @@ -1093,6 +1127,12 @@ StatusOr ElementalIrEmitter::EmitExpm1(PrimitiveType prim_type, return Select(x_is_small, for_small_x, for_large_x); } +StatusOr ElementalIrEmitter::EmitRoundNearestAfz( + PrimitiveType /*prim_type*/, llvm::Value* value) { + return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::round, {value}, + {value->getType()}, b_); +} + StatusOr ElementalIrEmitter::EmitPow(PrimitiveType prim_type, llvm::Value* lhs, llvm::Value* rhs) { @@ -2133,7 +2173,11 @@ StatusOr ElementalIrEmitter::EmitElementalDot( } lhs_index.InsertAt(lhs_contracting_dim, inner_loop->GetIndVarValue()); - for (int64 i = 0; i < rhs_dims - 1; i++) { + int64 num_batch_dims = dim_numbers.rhs_batch_dimensions_size(); + for (int64 i = 0; i < num_batch_dims; i++) { + rhs_index.push_back(dot_result_index[dim_numbers.rhs_batch_dimensions(i)]); + } + for (int64 i = 0; i < rhs_dims - 1 - num_batch_dims; i++) { rhs_index.push_back(dot_result_index[lhs_dims - 1 + i]); } rhs_index.InsertAt(rhs_contracting_dim, inner_loop->GetIndVarValue()); @@ -2188,8 +2232,10 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( case HloOpcode::kNegate: case HloOpcode::kNot: case HloOpcode::kReal: + case HloOpcode::kRsqrt: case HloOpcode::kSign: case HloOpcode::kSin: + case HloOpcode::kSqrt: case HloOpcode::kTanh: return [this, hlo, &operand_to_generator]( const IrArray::Index& index) -> StatusOr { @@ -2398,6 +2444,15 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( -> StatusOr { return EmitElementalDot(hlo, operand_to_generator, dot_result_index); }; + case HloOpcode::kReplicaId: + return [this, hlo](const IrArray::Index&) -> StatusOr { + if (hlo_module_config_.replica_count() != 1) { + return Unimplemented("Replication is not implemented on CPU/GPU."); + } + llvm::Type* type = llvm_ir::PrimitiveTypeToIrType( + hlo->shape().element_type(), module_); + return llvm::ConstantInt::getNullValue(type); + }; default: return [hlo](const IrArray::Index& index) { return Unimplemented("Unhandled opcode for elemental IR emission: %s", diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.h b/tensorflow/compiler/xla/service/elemental_ir_emitter.h index 7d360fe38cfeda17878c363253c41883ec9fd64f..7afecbbd31808e988e2ba08004d7598b363ac49e 100644 --- a/tensorflow/compiler/xla/service/elemental_ir_emitter.h +++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.h @@ -119,6 +119,12 @@ class ElementalIrEmitter : public IrBuilderMixin { virtual StatusOr EmitLog(PrimitiveType prim_type, llvm::Value* value); + virtual StatusOr EmitSqrt(PrimitiveType prim_type, + llvm::Value* value); + + virtual StatusOr EmitRsqrt(PrimitiveType prim_type, + llvm::Value* value); + virtual StatusOr EmitLog1p(PrimitiveType prim_type, llvm::Value* value); @@ -140,6 +146,9 @@ class ElementalIrEmitter : public IrBuilderMixin { virtual StatusOr EmitTanh(PrimitiveType prim_type, llvm::Value* value); + virtual StatusOr EmitRoundNearestAfz(PrimitiveType prim_type, + llvm::Value* value); + virtual StatusOr EmitReducePrecision(const HloInstruction* hlo, llvm::Value* x); @@ -211,6 +220,11 @@ class ElementalIrEmitter : public IrBuilderMixin { const HloModuleConfig& hlo_module_config_; private: + // Computes the complex power function, returns (a + i*b)^(c + i*d). + StatusOr EmitComplexPower(const HloInstruction* op, + llvm::Value* a, llvm::Value* b, + llvm::Value* c, llvm::Value* d); + // Returns a ElementGenerator for an RNG HloInstruction using the Philox // random number generation algorithm. llvm_ir::ElementGenerator MakePhiloxRngElementGenerator( diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter_test.cc b/tensorflow/compiler/xla/service/elemental_ir_emitter_test.cc index 852f34e06df35242b13110ae4411b8c969c26019..9e9d3daf25e0caadbc6b14bd1b00c467674d7fc7 100644 --- a/tensorflow/compiler/xla/service/elemental_ir_emitter_test.cc +++ b/tensorflow/compiler/xla/service/elemental_ir_emitter_test.cc @@ -60,5 +60,30 @@ ENTRY main { Literal rhs = LiteralUtil::CreateR3({{{3}, {4}}}); RunTest(hlo_text, {&lhs, &rhs}); } + +XLA_TEST_F(ElementalIrEmitterExecutionTest, BatchDot) { + const char* hlo_text = R"( +HloModule BatchDot + +fused_computation.1 { + param_0 = f64[1,1,8]{2,1,0} parameter(0) + r.1 = f64[2,4]{1,0} reshape(param_0) + param_1 = f64[1,2,2,2,1]{4,3,2,1,0} parameter(1) + r.2 = f64[2,4,1]{2,1,0} reshape(param_1) + ROOT dot = f64[2,1]{1,0} dot(r.1, r.2), lhs_batch_dims={0}, + lhs_contracting_dims={1}, + rhs_batch_dims={0}, + rhs_contracting_dims={1} +} + +ENTRY resampler_Resampler.49 { + p0 = f64[1,1,8]{2,1,0} parameter(0) + p1 = f64[1,2,2,2,1]{4,3,2,1,0} parameter(1) + ROOT f = f64[2,1]{1,0} fusion(p0, p1), kind=kLoop, calls=fused_computation.1 +} +)"; + + EXPECT_TRUE(RunAndCompare(hlo_text, ErrorSpec{4e-3, 4e-3})); +} } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/flatten_call_graph.cc b/tensorflow/compiler/xla/service/flatten_call_graph.cc index 85409b330b11537158059dcce8c2a96c98d38f30..f16a4485550a4262be8089c7d6c7c8252830dc1b 100644 --- a/tensorflow/compiler/xla/service/flatten_call_graph.cc +++ b/tensorflow/compiler/xla/service/flatten_call_graph.cc @@ -26,7 +26,7 @@ namespace xla { namespace { -// Helper to replace the called computation at a while-, call-, or +// Helper to replace the called computation at a while-, call-, case-, or // conditional-instruction. This function replaces exactly one instance of // 'computation' with 'new_computation' even if 'instruction' calls // 'computation' more than once. @@ -49,11 +49,14 @@ void ReplaceCalledComputation(HloInstruction* instruction, break; } case HloOpcode::kConditional: { - if (computation == instruction->true_computation()) { - instruction->set_true_computation(new_computation); - } else { - CHECK_EQ(computation, instruction->false_computation()); - instruction->set_false_computation(new_computation); + for (int b = 0; b < instruction->branch_count(); ++b) { + if (b == instruction->branch_count() - 1) { + CHECK_EQ(computation, instruction->branch_computation(b)); + } + if (computation == instruction->branch_computation(b)) { + instruction->set_branch_computation(b, new_computation); + break; + } } break; } diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD index 25c4f70d89b4ebc483a61f1e28c7a55eb31f4bdf..3bc0daf9e708a9472ee1724ea7a6161e42adfd0c 100644 --- a/tensorflow/compiler/xla/service/gpu/BUILD +++ b/tensorflow/compiler/xla/service/gpu/BUILD @@ -7,7 +7,7 @@ load( "//tensorflow/core:platform/default/build_config_root.bzl", "tf_cuda_tests_tags", ) -load("//tensorflow:tensorflow.bzl", "tf_cc_test") +load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_cuda_library", "if_cuda") licenses(["notice"]) # Apache 2.0 @@ -156,7 +156,6 @@ cc_library( "ir_emitter_unnested.h", ], deps = [ - ":backend_configs", ":buffer_allocations", ":cudnn_conv_runner", ":elemental_ir_emitter", @@ -164,8 +163,10 @@ cc_library( ":gpu_executable", ":hlo_to_ir_bindings", ":ir_emission_utils", + ":nccl_all_reduce_thunk", ":parallel_loop_emitter", ":partition_assignment", + ":thunk", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -179,6 +180,7 @@ cc_library( "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_casting_utils", "//tensorflow/compiler/xla/service:name_uniquer", + "//tensorflow/compiler/xla/service:pattern_matcher", "//tensorflow/compiler/xla/service:while_loop_analysis", "//tensorflow/compiler/xla/service/llvm_ir:buffer_assignment_util", "//tensorflow/compiler/xla/service/llvm_ir:dynamic_update_slice_util", @@ -287,9 +289,44 @@ cc_library( ], ) +cc_library( + name = "thunk", + srcs = ["thunk.cc"], + hdrs = ["thunk.h"], + deps = [ + ":buffer_allocations", + ":hlo_execution_profiler", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/core:lib", + "//tensorflow/core:stream_executor_no_cuda", + ], +) + +tf_cuda_library( + name = "nccl_all_reduce_thunk", + srcs = ["nccl_all_reduce_thunk.cc"], + hdrs = ["nccl_all_reduce_thunk.h"], + deps = [ + ":buffer_allocations", + ":hlo_execution_profiler", + ":thunk", + "@com_google_absl//absl/synchronization", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla/service:buffer_assignment", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/core:lib", + "//tensorflow/core:stream_executor_no_cuda", + "//tensorflow/stream_executor/cuda:cuda_activation", + "//tensorflow/stream_executor/cuda:cuda_gpu_executor", + ] + if_cuda([ + "@local_config_nccl//:nccl", + ]), +) + cc_library( name = "gpu_executable", srcs = [ + "cholesky_thunk.cc", "conditional_thunk.cc", "convolution_thunk.cc", "copy_thunk.cc", @@ -303,13 +340,13 @@ cc_library( "memset_thunk.cc", "outfeed_thunk.cc", "sequential_thunk.cc", - "thunk.cc", "thunk_schedule.cc", "triangular_solve_thunk.cc", "tuple_thunk.cc", "while_thunk.cc", ], hdrs = [ + "cholesky_thunk.h", "conditional_thunk.h", "convolution_thunk.h", "copy_thunk.h", @@ -323,7 +360,6 @@ cc_library( "memset_thunk.h", "outfeed_thunk.h", "sequential_thunk.h", - "thunk.h", "thunk_schedule.h", "triangular_solve_thunk.h", "tuple_thunk.h", @@ -332,12 +368,15 @@ cc_library( deps = [ ":buffer_allocations", ":cudnn_conv_runner", + ":cusolver_context", ":hlo_execution_profiler", ":infeed_manager", ":ir_emission_utils", + ":nccl_all_reduce_thunk", ":outfeed_manager", ":partition_assignment", ":stream_assignment", + ":thunk", "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_tree", @@ -368,6 +407,7 @@ cc_library( "//tensorflow/stream_executor", "//tensorflow/stream_executor:blas", "//tensorflow/stream_executor:device_memory", + "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/memory", @@ -406,6 +446,7 @@ cc_library( ":cudnn_conv_runner", ":gpu_executable", ":ir_emission_utils", + ":scratch_allocator", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:protobuf_util", "//tensorflow/compiler/xla/service:compiler", @@ -423,6 +464,18 @@ cc_library( ], ) +cc_library( + name = "scratch_allocator", + srcs = ["scratch_allocator.cc"], + hdrs = ["scratch_allocator.h"], + deps = [ + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla/service:device_memory_allocator", + "//tensorflow/core:stream_executor_no_cuda", + ], +) + cc_library( name = "cudnn_conv_runner", srcs = ["cudnn_conv_runner.cc"], @@ -479,6 +532,43 @@ tf_cc_test( ], ) +cc_library( + name = "cusolver_context", + srcs = ["cusolver_context.cc"], + hdrs = ["cusolver_context.h"], + deps = [ + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + "//tensorflow/core:stream_executor_no_cuda", + "//tensorflow/stream_executor:blas", + "@local_config_cuda//cuda:cuda_headers", + "@local_config_cuda//cuda:cusolver", + ], +) + +cc_library( + name = "cusolver_rewriter", + srcs = ["cusolver_rewriter.cc"], + hdrs = ["cusolver_rewriter.h"], + deps = [ + ":cusolver_context", + ":ir_emission_utils", + ":scratch_allocator", + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service:device_memory_allocator", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_pass", + "//tensorflow/core:lib", + "//tensorflow/core:stream_executor_no_cuda", + "//tensorflow/stream_executor:blas", + "@com_google_absl//absl/types:optional", + ], +) + cc_library( name = "instruction_fusion", srcs = ["instruction_fusion.cc"], @@ -712,11 +802,13 @@ cc_library( srcs = ["nvptx_compiler.cc"], hdrs = ["nvptx_compiler.h"], deps = [ + ":cudnn_batchnorm_rewriter", ":cudnn_conv_algorithm_picker", ":cudnn_conv_pad_for_tensor_cores", ":cudnn_conv_padding_legalization", ":cudnn_conv_rewriter", ":cudnn_fused_conv_rewriter", + ":cusolver_rewriter", ":fusion_merger", ":gpu_constants", ":gpu_copy_insertion", @@ -770,8 +862,8 @@ cc_library( "//tensorflow/compiler/xla/service:tuple_simplifier", "//tensorflow/compiler/xla/service:while_loop_constant_sinking", "//tensorflow/compiler/xla/service:while_loop_simplifier", + "//tensorflow/compiler/xla/service:while_loop_trip_count_annotator", "//tensorflow/compiler/xla/service:zero_sized_hlo_elimination", - "//tensorflow/compiler/xla/service/gpu:cudnn_batchnorm_rewriter", "//tensorflow/compiler/xla/service/gpu/llvm_gpu_backend", "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", "//tensorflow/core:cuda_libdevice_path", diff --git a/tensorflow/compiler/xla/service/gpu/cholesky_thunk.cc b/tensorflow/compiler/xla/service/gpu/cholesky_thunk.cc new file mode 100644 index 0000000000000000000000000000000000000000..7daef16cb62338cfa5b027136ecd4262288eec8d --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/cholesky_thunk.cc @@ -0,0 +1,119 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/gpu/cholesky_thunk.h" + +#include + +#include "absl/strings/str_cat.h" +#include "absl/strings/str_format.h" +#include "tensorflow/compiler/xla/service/gpu/hlo_execution_profiler.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" +#include "tensorflow/stream_executor/blas.h" +#include "tensorflow/stream_executor/device_memory.h" + +namespace xla { +namespace gpu { + +CholeskyThunk::CholeskyThunk(const CholeskyOptions& options, + BufferAllocation::Slice a_buffer, + BufferAllocation::Slice workspace_buffer, + BufferAllocation::Slice info_buffer, + PrimitiveType type, int64 batch_size, int64 n, + const HloInstruction* hlo) + : Thunk(Kind::kCholesky, hlo), + uplo_(options.lower() ? se::blas::UpperLower::kLower + : se::blas::UpperLower::kUpper), + a_buffer_(a_buffer), + workspace_buffer_(workspace_buffer), + info_buffer_(info_buffer), + type_(type), + batch_size_(batch_size), + a_batch_stride_(n * n * + ShapeUtil::ByteSizeOfPrimitiveType( + hlo->operand(0)->shape().element_type())), + n_(n) {} + +Status CholeskyThunk::ExecuteOnStream( + const BufferAllocations& buffer_allocations, se::Stream* stream, + HloExecutionProfiler* profiler) { + VLOG(3) << "type=" << PrimitiveType_Name(type_) + << " uplo=" << se::blas::UpperLowerString(uplo_) + << " batch_size=" << batch_size_ << " n=" << n_ + << " a=" << a_buffer_.ToString() + << " workspace=" << workspace_buffer_.ToString() + << " info=" << info_buffer_.ToString(); + + CusolverContext* context; + { + tensorflow::mutex_lock lock(mu_); + auto result = contexts_.emplace(stream, CusolverContext()); + if (result.second) { + TF_ASSIGN_OR_RETURN(result.first->second, + CusolverContext::Create(stream)); + } + context = &result.first->second; + } + + char* a_base = static_cast( + buffer_allocations.GetDeviceAddress(a_buffer_).opaque()); + int* info_base = static_cast( + buffer_allocations.GetDeviceAddress(info_buffer_).opaque()); + se::DeviceMemoryBase workspace_data = + buffer_allocations.GetDeviceAddress(workspace_buffer_); + for (int64 i = 0; i < batch_size_; ++i) { + se::DeviceMemoryBase a_data = + se::DeviceMemoryBase(a_base + i * a_batch_stride_, a_batch_stride_); + se::DeviceMemory info_data( + se::DeviceMemoryBase(info_base + i, sizeof(int))); + switch (type_) { + case F32: { + TF_RETURN_IF_ERROR( + context->Potrf(uplo_, n_, se::DeviceMemory(a_data), n_, + info_data, se::DeviceMemory(workspace_data))); + break; + } + case F64: { + TF_RETURN_IF_ERROR(context->Potrf( + uplo_, n_, se::DeviceMemory(a_data), n_, info_data, + se::DeviceMemory(workspace_data))); + break; + } + case C64: { + TF_RETURN_IF_ERROR(context->Potrf( + uplo_, n_, se::DeviceMemory>(a_data), n_, + info_data, se::DeviceMemory>(workspace_data))); + break; + } + case C128: { + TF_RETURN_IF_ERROR(context->Potrf( + uplo_, n_, se::DeviceMemory>(a_data), n_, + info_data, se::DeviceMemory>(workspace_data))); + break; + } + default: + return InvalidArgument("Invalid type for cholesky %s", + PrimitiveType_Name(type_)); + } + } + return Status::OK(); +} + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/cholesky_thunk.h b/tensorflow/compiler/xla/service/gpu/cholesky_thunk.h new file mode 100644 index 0000000000000000000000000000000000000000..cde245a7e8bc0909059d4643cae3de138bddcdec --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/cholesky_thunk.h @@ -0,0 +1,77 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CHOLESKY_THUNK_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CHOLESKY_THUNK_H_ + +#include "absl/base/thread_annotations.h" +#include "absl/types/optional.h" +#include "tensorflow/compiler/xla/service/buffer_assignment.h" +#include "tensorflow/compiler/xla/service/gpu/buffer_allocations.h" +#include "tensorflow/compiler/xla/service/gpu/cusolver_context.h" +#include "tensorflow/compiler/xla/service/gpu/gpu_executable.h" +#include "tensorflow/compiler/xla/service/gpu/hlo_execution_profiler.h" +#include "tensorflow/compiler/xla/service/gpu/thunk.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" +#include "tensorflow/stream_executor/blas.h" + +namespace xla { +namespace gpu { + +// This class stores everything that StreamExecutor needs to launch a Cholesky +// decomposition (LAPACK potrf). It is generated by IrEmitter. +// +// Thread-compatible. +class CholeskyThunk : public Thunk { + public: + static StatusOr ScratchBufferSize(int64 n); + CholeskyThunk(const CholeskyOptions& options, + BufferAllocation::Slice a_buffer, + BufferAllocation::Slice workspace_buffer, + BufferAllocation::Slice info_buffer, + PrimitiveType type, + int64 batch_size, int64 n, const HloInstruction* hlo); + + CholeskyThunk(const CholeskyThunk&) = delete; + CholeskyThunk& operator=(const CholeskyThunk&) = delete; + + Status ExecuteOnStream(const BufferAllocations& buffer_allocations, + se::Stream* stream, + HloExecutionProfiler* profiler) override; + + private: + se::blas::UpperLower uplo_; + + const BufferAllocation::Slice a_buffer_; + const BufferAllocation::Slice workspace_buffer_; + const BufferAllocation::Slice info_buffer_; + + const PrimitiveType type_; + const int64 batch_size_; + const int64 a_batch_stride_; + const int64 n_; + + tensorflow::mutex mu_; + absl::flat_hash_map contexts_ GUARDED_BY(mu_); +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CHOLESKY_THUNK_H_ diff --git a/tensorflow/compiler/xla/service/gpu/conditional_thunk.cc b/tensorflow/compiler/xla/service/gpu/conditional_thunk.cc index 9ed523998bf07567133fdac0e40b12b8ce4ea3b0..ea6392498264f25d53bec2309bfdf7bdcf6a2a2e 100644 --- a/tensorflow/compiler/xla/service/gpu/conditional_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/conditional_thunk.cc @@ -24,25 +24,35 @@ namespace xla { namespace gpu { ConditionalThunk::ConditionalThunk( - const BufferAllocation::Slice& predicate_buffer_index, - const BufferAllocation::Slice& true_operand_buffer_index, - const BufferAllocation::Slice& false_operand_buffer_index, - ThunkSequence true_thunk_sequence, ThunkSequence false_thunk_sequence, + const BufferAllocation::Slice& branch_index_buffer_index, + absl::Span branch_operand_buffer_indexes, + std::vector branch_thunk_sequences, const HloInstruction* hlo) : Thunk(Kind::kConditional, hlo), - predicate_buffer_index_(predicate_buffer_index), - true_operand_buffer_index_(true_operand_buffer_index), - false_operand_buffer_index_(false_operand_buffer_index), - // Pass nullptr as the HloInstruction* to the true_thunk_ and false_thunk_ - // constructors because these SequentialThunks are logically "part of" - // this ConditionalThunk, and shouldn't be profiled separately from it. - true_thunk_(std::move(true_thunk_sequence), nullptr), - false_thunk_(std::move(false_thunk_sequence), nullptr) {} + branch_index_is_bool_(hlo->operand(0)->shape().element_type() == PRED), + branch_index_buffer_index_(branch_index_buffer_index), + branch_operand_buffer_indexes_(branch_operand_buffer_indexes.begin(), + branch_operand_buffer_indexes.end()) { + // Pass nullptr as the HloInstruction* to the branch_thunks_ + // constructors because these SequentialThunks are logically "part of" + // this ConditionalThunk, and shouldn't be profiled separately from it. + branch_thunks_.reserve(branch_thunk_sequences.size()); + for (auto& branch_thunk_sequence : branch_thunk_sequences) { + branch_thunks_.emplace_back( + new SequentialThunk(std::move(branch_thunk_sequence), nullptr)); + } +} Status ConditionalThunk::Initialize(const GpuExecutable& executable, se::StreamExecutor* executor) { - TF_RETURN_IF_ERROR(true_thunk_.Initialize(executable, executor)); - TF_RETURN_IF_ERROR(false_thunk_.Initialize(executable, executor)); + if (branch_index_is_bool_) { + TF_RET_CHECK(branch_thunks_.size() == 2); + } else { + TF_RET_CHECK(!branch_thunks_.empty()); + } + for (auto& branch_thunk : branch_thunks_) { + TF_RETURN_IF_ERROR(branch_thunk->Initialize(executable, executor)); + } return Status::OK(); } @@ -51,31 +61,38 @@ Status ConditionalThunk::ExecuteOnStream( HloExecutionProfiler* profiler) { auto op_profiler = profiler->MakeScopedInstructionProfiler(hlo_instruction()); // Copy the predicate value from device. - bool predicate; - se::DeviceMemoryBase predicate_address = - buffer_allocations.GetDeviceAddress(predicate_buffer_index_); - stream->ThenMemcpy(&predicate, predicate_address, sizeof(bool)); + int32 branch_index = -1; + bool pred = false; + se::DeviceMemoryBase branch_index_address = + buffer_allocations.GetDeviceAddress(branch_index_buffer_index_); + if (branch_index_is_bool_) { + stream->ThenMemcpy(&pred, branch_index_address, sizeof(bool)); + } else { + stream->ThenMemcpy(&branch_index, branch_index_address, sizeof(int32)); + } Status block_status = stream->BlockHostUntilDone(); if (!block_status.ok()) { - return InternalError("Failed to retrieve predicate value on stream %p: %s.", - stream, block_status.error_message()); + return InternalError( + "Failed to retrieve branch_index value on stream %p: %s.", stream, + block_status.error_message()); } - - // Execute the true or the false computation depending on the value of the - // predicate. - if (predicate) { - profiler->StartHloComputation(); - TF_RETURN_IF_ERROR( - true_thunk_.ExecuteOnStream(buffer_allocations, stream, profiler)); - profiler->FinishHloComputation(hlo_instruction()->true_computation()); + if (branch_index_is_bool_) { + branch_index = pred ? 0 : 1; } else { - profiler->StartHloComputation(); - TF_RETURN_IF_ERROR( - false_thunk_.ExecuteOnStream(buffer_allocations, stream, profiler)); - profiler->FinishHloComputation(hlo_instruction()->false_computation()); + // Handle default scenario for branch_index not in [0, num_branches). + if (branch_index < 0 || branch_index >= hlo_instruction()->branch_count()) { + branch_index = hlo_instruction()->branch_count() - 1; + } } + // Execute the branch computation corresponding to the value of branch_index. + profiler->StartHloComputation(); + TF_RETURN_IF_ERROR(branch_thunks_[branch_index]->ExecuteOnStream( + buffer_allocations, stream, profiler)); + profiler->FinishHloComputation( + hlo_instruction()->branch_computation(branch_index)); + return Status::OK(); } diff --git a/tensorflow/compiler/xla/service/gpu/conditional_thunk.h b/tensorflow/compiler/xla/service/gpu/conditional_thunk.h index aef24342c9fe182eb54b1c2beff840a76e7b8115..c0093ca6397e636bee953ddf0af8c48caaaadae0 100644 --- a/tensorflow/compiler/xla/service/gpu/conditional_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/conditional_thunk.h @@ -16,6 +16,10 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CONDITIONAL_THUNK_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CONDITIONAL_THUNK_H_ +#include +#include + +#include "absl/types/span.h" #include "tensorflow/compiler/xla/service/gpu/buffer_allocations.h" #include "tensorflow/compiler/xla/service/gpu/hlo_execution_profiler.h" #include "tensorflow/compiler/xla/service/gpu/sequential_thunk.h" @@ -38,12 +42,11 @@ namespace gpu { // false computation share the same allocation. class ConditionalThunk : public Thunk { public: - ConditionalThunk(const BufferAllocation::Slice& predicate_buffer_index, - const BufferAllocation::Slice& true_operand_buffer_index, - const BufferAllocation::Slice& false_operand_buffer_index, - ThunkSequence true_thunk_sequence, - ThunkSequence false_thunk_sequence, - const HloInstruction* hlo); + ConditionalThunk( + const BufferAllocation::Slice& branch_index_buffer_index, + absl::Span branch_operand_buffer_indexes, + std::vector branch_thunk_sequences, + const HloInstruction* hlo); ConditionalThunk(const ConditionalThunk&) = delete; ConditionalThunk& operator=(const ConditionalThunk&) = delete; @@ -55,11 +58,10 @@ class ConditionalThunk : public Thunk { HloExecutionProfiler* profiler) override; private: - BufferAllocation::Slice predicate_buffer_index_; - BufferAllocation::Slice true_operand_buffer_index_; - BufferAllocation::Slice false_operand_buffer_index_; - SequentialThunk true_thunk_; - SequentialThunk false_thunk_; + const bool branch_index_is_bool_; + BufferAllocation::Slice branch_index_buffer_index_; + std::vector branch_operand_buffer_indexes_; + std::vector> branch_thunks_; }; } // namespace gpu diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_batchnorm_rewriter.cc b/tensorflow/compiler/xla/service/gpu/cudnn_batchnorm_rewriter.cc index 60289506524759580dbb9b82147c78c4ce1cb25e..2cceb0422d08ff7951308b0727941f5437785447 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_batchnorm_rewriter.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_batchnorm_rewriter.cc @@ -188,13 +188,8 @@ Status Visitor::HandleBatchNormGrad(HloInstruction* batch_norm) { computation_->AddInstruction(HloInstruction::CreateBroadcast( batch_norm->operand(3)->shape(), epsilon, {})))); HloInstruction* inverse_stddev = - computation_->AddInstruction(HloInstruction::CreateBinary( - var_plus_epsilon->shape(), HloOpcode::kPower, var_plus_epsilon, - computation_->AddInstruction(HloInstruction::CreateBroadcast( - var_plus_epsilon->shape(), - computation_->AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR0(-.5))), - {})))); + computation_->AddInstruction(HloInstruction::CreateUnary( + var_plus_epsilon->shape(), HloOpcode::kRsqrt, var_plus_epsilon)); std::vector operands(batch_norm->operands().begin(), batch_norm->operands().end()); diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.cc b/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.cc index 603af5a654589e0b02c762b57d70a8b7628b1d0f..0c4980f6549714572682d7d3da94f57347f24587 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.cc @@ -24,6 +24,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/buffer_comparator.h" #include "tensorflow/compiler/xla/service/gpu/convolution_thunk.h" #include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h" +#include "tensorflow/compiler/xla/service/gpu/scratch_allocator.h" #include "tensorflow/compiler/xla/service/hlo_casting_utils.h" #include "tensorflow/core/lib/strings/numbers.h" #include "tensorflow/core/platform/logger.h" @@ -37,47 +38,6 @@ using absl::optional; using se::DeviceMemoryBase; using se::dnn::AlgorithmDesc; -class ScratchAllocator : public se::ScratchAllocator { - public: - ScratchAllocator(int device_ordinal, DeviceMemoryAllocator* memory_allocator) - : device_ordinal_(device_ordinal), memory_allocator_(memory_allocator) {} - - int64 GetMemoryLimitInBytes(se::Stream* stream) override { - return 1LL << 32; // 4GB. TODO(jlebar): Tune this? - } - int64 TotalAllocatedBytes() { return total_allocated_bytes_; } - - StatusOr> AllocateBytes(se::Stream* stream, - int64 byte_size) override; - - private: - const int device_ordinal_; - DeviceMemoryAllocator* memory_allocator_; - std::vector allocated_buffers_; - int64 total_allocated_bytes_ = 0; -}; - -StatusOr> ScratchAllocator::AllocateBytes( - se::Stream* stream, int64 byte_size) { - CHECK_GE(byte_size, 0) << "byte_size must be positive."; - if (byte_size > GetMemoryLimitInBytes(stream)) { - return se::port::Status( - se::port::error::RESOURCE_EXHAUSTED, - absl::StrFormat( - "Allocating %d bytes exceeds the memory limit of %d bytes.", - byte_size, GetMemoryLimitInBytes(stream))); - } - - TF_ASSIGN_OR_RETURN(OwningDeviceMemory allocated_buffer, - memory_allocator_->Allocate(device_ordinal_, byte_size, - /*retry_on_failure=*/false)); - total_allocated_bytes_ += byte_size; - - se::DeviceMemoryBase buffer_addr = allocated_buffer.AsDeviceMemoryBase(); - allocated_buffers_.push_back(std::move(allocated_buffer)); - return se::DeviceMemory(buffer_addr); -} - std::vector GetAlgorithms(CudnnConvKind kind, se::StreamExecutor* stream_exec) { std::vector algorithms; diff --git a/tensorflow/compiler/xla/service/gpu/cusolver_context.cc b/tensorflow/compiler/xla/service/gpu/cusolver_context.cc new file mode 100644 index 0000000000000000000000000000000000000000..923b7bc452870f47505711e8abd4ce236be7815a --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/cusolver_context.cc @@ -0,0 +1,159 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/gpu/cusolver_context.h" + +#include "tensorflow/compiler/xla/util.h" + +namespace xla { +namespace gpu { + +namespace { + +// Type traits to get CUDA complex types from std::complex. +template +struct CUDAComplexT { + typedef T type; +}; +template <> +struct CUDAComplexT> { + typedef cuComplex type; +}; +template <> +struct CUDAComplexT> { + typedef cuDoubleComplex type; +}; + +template +inline typename CUDAComplexT::type* ToDevicePointer(se::DeviceMemory p) { + return static_cast::type*>(p.opaque()); +} + +cublasFillMode_t CUDABlasUpperLower(se::blas::UpperLower uplo) { + switch (uplo) { + case se::blas::UpperLower::kUpper: + return CUBLAS_FILL_MODE_UPPER; + case se::blas::UpperLower::kLower: + return CUBLAS_FILL_MODE_LOWER; + default: + LOG(FATAL) << "Invalid value of blas::UpperLower."; + } +} + +// Converts a cuSolver status to a Status. +Status CusolverStatusToStatus(cusolverStatus_t status) { + switch (status) { + case CUSOLVER_STATUS_SUCCESS: + return Status::OK(); + case CUSOLVER_STATUS_NOT_INITIALIZED: + return FailedPrecondition("cuSolver has not been initialized"); + case CUSOLVER_STATUS_ALLOC_FAILED: + return ResourceExhausted("cuSolver allocation failed"); + case CUSOLVER_STATUS_INVALID_VALUE: + return InvalidArgument("cuSolver invalid value error"); + case CUSOLVER_STATUS_ARCH_MISMATCH: + return FailedPrecondition("cuSolver architecture mismatch error"); + case CUSOLVER_STATUS_MAPPING_ERROR: + return Unknown("cuSolver mapping error"); + case CUSOLVER_STATUS_EXECUTION_FAILED: + return Unknown("cuSolver execution failed"); + case CUSOLVER_STATUS_INTERNAL_ERROR: + return Internal("cuSolver internal error"); + case CUSOLVER_STATUS_MATRIX_TYPE_NOT_SUPPORTED: + return Unimplemented("cuSolver matrix type not supported error"); + case CUSOLVER_STATUS_NOT_SUPPORTED: + return Unimplemented("cuSolver not supported error"); + case CUSOLVER_STATUS_ZERO_PIVOT: + return InvalidArgument("cuSolver zero pivot error"); + case CUSOLVER_STATUS_INVALID_LICENSE: + return FailedPrecondition("cuSolver invalid license error"); + default: + return Unknown("Unknown cuSolver error"); + } +} + +} // namespace + +StatusOr CusolverContext::Create(se::Stream* stream) { + cusolverDnHandle_t handle; + TF_RETURN_IF_ERROR(CusolverStatusToStatus(cusolverDnCreate(&handle))); + CusolverContext context(stream, handle); + + // StreamExecutor really should just expose the Cuda stream to clients... + const cudaStream_t* cuda_stream = + CHECK_NOTNULL(reinterpret_cast( + stream->implementation()->GpuStreamMemberHack())); + TF_RETURN_IF_ERROR( + CusolverStatusToStatus(cusolverDnSetStream(handle, *cuda_stream))); + + return std::move(context); +} + +CusolverContext::CusolverContext(se::Stream* stream, cusolverDnHandle_t handle) + : stream_(stream), handle_(handle) {} + +CusolverContext::CusolverContext(CusolverContext&& other) { + handle_ = other.handle_; + stream_ = other.stream_; + other.handle_ = nullptr; + other.stream_ = nullptr; +} + +CusolverContext& CusolverContext::operator=(CusolverContext&& other) { + std::swap(handle_, other.handle_); + std::swap(stream_, other.stream_); + return *this; +} + +CusolverContext::~CusolverContext() { + if (handle_) { + Status status = CusolverStatusToStatus(cusolverDnDestroy(handle_)); + if (!status.ok()) { + LOG(ERROR) << "cusolverDnDestroy failed: " << status; + } + } +} + +#define CALL_LAPACK_TYPES(m) \ + m(float, S) m(double, D) m(std::complex, C) m(std::complex, Z) + +#define DN_SOLVER_FN(method, type_prefix) cusolverDn##type_prefix##method + +#define POTRF_BUFFER_SIZE_INSTANCE(T, type_prefix) \ + StatusOr CusolverContext::PotrfBufferSize( \ + se::blas::UpperLower uplo, int n, se::DeviceMemory A, int lda) { \ + int size = -1; \ + TF_RETURN_IF_ERROR(CusolverStatusToStatus(DN_SOLVER_FN( \ + potrf_bufferSize, type_prefix)(handle(), CUDABlasUpperLower(uplo), n, \ + ToDevicePointer(A), lda, &size))); \ + return size; \ + } + +CALL_LAPACK_TYPES(POTRF_BUFFER_SIZE_INSTANCE); + +#define POTRF_INSTANCE(T, type_prefix) \ + Status CusolverContext::Potrf( \ + se::blas::UpperLower uplo, int n, se::DeviceMemory A, int lda, \ + se::DeviceMemory lapack_info, se::DeviceMemory workspace) { \ + return CusolverStatusToStatus(DN_SOLVER_FN(potrf, type_prefix)( \ + handle(), CUDABlasUpperLower(uplo), n, ToDevicePointer(A), lda, \ + ToDevicePointer(workspace), workspace.ElementCount(), \ + ToDevicePointer(lapack_info))); \ + } + +CALL_LAPACK_TYPES(POTRF_INSTANCE); + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/cusolver_context.h b/tensorflow/compiler/xla/service/gpu/cusolver_context.h new file mode 100644 index 0000000000000000000000000000000000000000..fdd89c3a8d599e2291b60abcd67e267a96d3ac8f --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/cusolver_context.h @@ -0,0 +1,88 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUSOLVER_CONTEXT_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUSOLVER_CONTEXT_H_ + +#include + +#include "cuda/include/cublas_v2.h" +#include "cuda/include/cusolverDn.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" +#include "tensorflow/stream_executor/blas.h" + +namespace xla { +namespace gpu { + +class CusolverContext { + public: + static StatusOr Create(se::Stream* stream); + CusolverContext() = default; + ~CusolverContext(); + + CusolverContext(const CusolverContext&) = delete; + CusolverContext(CusolverContext&&); + CusolverContext& operator=(const CusolverContext&) = delete; + CusolverContext& operator=(CusolverContext&&); + + se::Stream* stream() const { return stream_; } + cusolverDnHandle_t handle() const { return handle_; } + + // Computes the Cholesky factorization A = L * L^T for a single matrix. + // Returns Status::OK() if the kernel was launched successfully. See: + // http://docs.nvidia.com/cuda/cusolver/#cuds-lt-t-gt-potrf + Status Potrf(se::blas::UpperLower uplo, int n, se::DeviceMemory dev_A, + int lda, se::DeviceMemory dev_lapack_info, + se::DeviceMemory workspace); + Status Potrf(se::blas::UpperLower uplo, int n, se::DeviceMemory dev_A, + int lda, se::DeviceMemory dev_lapack_info, + se::DeviceMemory workspace); + Status Potrf(se::blas::UpperLower uplo, int n, + se::DeviceMemory> dev_A, int lda, + se::DeviceMemory dev_lapack_info, + se::DeviceMemory> workspace); + Status Potrf(se::blas::UpperLower uplo, int n, + se::DeviceMemory> dev_A, int lda, + se::DeviceMemory dev_lapack_info, + se::DeviceMemory> workspace); + + // Returns the size of the `workspace` required by Potrf, in number of + // elements of size T. + StatusOr PotrfBufferSize(se::blas::UpperLower uplo, int n, + se::DeviceMemory dev_A, int lda); + StatusOr PotrfBufferSize(se::blas::UpperLower uplo, int n, + se::DeviceMemory dev_A, int lda); + StatusOr PotrfBufferSize(se::blas::UpperLower uplo, int n, + se::DeviceMemory> dev_A, + int lda); + StatusOr PotrfBufferSize(se::blas::UpperLower uplo, int n, + se::DeviceMemory> dev_A, + int lda); + + private: + CusolverContext(se::Stream* stream, cusolverDnHandle_t handle); + + se::Stream* stream_ = nullptr; + cusolverDnHandle_t handle_ = nullptr; +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUSOLVER_CONTEXT_H_ diff --git a/tensorflow/compiler/xla/service/gpu/cusolver_rewriter.cc b/tensorflow/compiler/xla/service/gpu/cusolver_rewriter.cc new file mode 100644 index 0000000000000000000000000000000000000000..7861eb1ef04d4fa5ba5690ee388b77a3f354f88e --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/cusolver_rewriter.cc @@ -0,0 +1,216 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/gpu/cusolver_rewriter.h" + +#include +#include +#include + +#include "absl/types/optional.h" +#include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" +#include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h" +#include "tensorflow/compiler/xla/service/gpu/scratch_allocator.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" +#include "tensorflow/stream_executor/blas.h" + +namespace xla { +namespace gpu { + +namespace { + +void SetFortranLayout(Shape* shape) { + LayoutUtil::SetToDefaultLayout(shape); + int n = shape->mutable_layout()->minor_to_major_size(); + CHECK_GE(n, 2); + std::swap(shape->mutable_layout()->mutable_minor_to_major()->at(0), + shape->mutable_layout()->mutable_minor_to_major()->at(1)); +} + +StatusOr CreateCholesky(CusolverContext* context, + ScratchAllocator* allocator, + HloInstruction* operand, + const CholeskyOptions& options, + const OpMetadata& metadata) { + HloComputation* computation = operand->parent(); + + Shape a_shape = operand->shape(); + int ndim = a_shape.dimensions_size(); + CHECK_GE(ndim, 2); + int64 n = a_shape.dimensions(ndim - 1); + + int64 batch_size = std::accumulate(a_shape.dimensions().begin(), + a_shape.dimensions().end() - 2, int64{1}, + [](int64 a, int64 b) { return a * b; }); + + // Find the workspace size. + se::blas::UpperLower uplo = options.lower() ? se::blas::UpperLower::kLower + : se::blas::UpperLower::kUpper; + int64 workspace_size; // Number of elements of size a_shape.element_type() + switch (a_shape.element_type()) { + case F32: { + TF_ASSIGN_OR_RETURN(auto a, + allocator->Allocate(context->stream(), n * n)); + TF_ASSIGN_OR_RETURN(workspace_size, + context->PotrfBufferSize(uplo, n, a, n)); + break; + } + case F64: { + TF_ASSIGN_OR_RETURN( + auto a, allocator->Allocate(context->stream(), n * n)); + TF_ASSIGN_OR_RETURN(workspace_size, + context->PotrfBufferSize(uplo, n, a, n)); + break; + } + case C64: { + TF_ASSIGN_OR_RETURN(auto a, allocator->Allocate>( + context->stream(), n * n)); + TF_ASSIGN_OR_RETURN(workspace_size, + context->PotrfBufferSize(uplo, n, a, n)); + break; + } + case C128: { + TF_ASSIGN_OR_RETURN(auto a, allocator->Allocate>( + context->stream(), n * n)); + TF_ASSIGN_OR_RETURN(workspace_size, + context->PotrfBufferSize(uplo, n, a, n)); + break; + } + default: + return InvalidArgument("Invalid type for cholesky decomposition: %s", + a_shape.ToString()); + } + + // TODO(phawkins): Ideally we would relax this constraint. What we actually + // want is that: + // a) the batch dimensions are major, in no particular order. + // b) the two minor dimensions are in fortran (column-major) order, + + SetFortranLayout(&a_shape); + + // This call returns a tuple of (cholesky_result, workspace, info) where: + // * cholesky_result is the result of the Cholesky decomposition, + // * workspace is temporary scratch memory used by cuSolver. + // * info contains the Potrf success/failure status. + // Currently we have no meaningful way to report an error, so we simply + // discard the success/failure information. Obviously this is suboptimal. + Shape call_shape = ShapeUtil::MakeTupleShape( + {a_shape, + ShapeUtil::MakeShape(operand->shape().element_type(), {workspace_size}), + ShapeUtil::MakeShape(S32, {batch_size})}); + + HloInstruction* custom_call = + computation->AddInstruction(HloInstruction::CreateCustomCall( + call_shape, {operand}, kCusolverCholeskyCallTarget, {a_shape})); + custom_call->set_metadata(metadata); + TF_RETURN_IF_ERROR(custom_call->set_backend_config(options)); + return custom_call; +} + +} // namespace + +// Tries to rewrite a single convolution into a call to cudnn. +StatusOr RunOnInstruction(CusolverContext* context, + ScratchAllocator* allocator, + HloInstruction* instruction) { + if (instruction->opcode() != HloOpcode::kCholesky) { + return false; + } + + TF_ASSIGN_OR_RETURN( + HloInstruction * custom_call, + CreateCholesky(context, allocator, instruction->mutable_operand(0), + instruction->cholesky_options(), instruction->metadata())); + + VLOG(1) << "Replacing " << instruction->ToString() << " with " + << custom_call->ToString(); + + // The CustomCall returns a tuple (conv_result, scratch_memory). Extract out + // the conv result and replace `conv` with it. + TF_RETURN_IF_ERROR(instruction->parent()->ReplaceWithNewInstruction( + instruction, HloInstruction::CreateGetTupleElement(instruction->shape(), + custom_call, 0))); + return true; +} + +// Rewrites the convolutions in the given computation into calls to cudnn. +// Returns true if it made any changes. +StatusOr CusolverRewriter::RunOnComputation(HloComputation* computation) { + std::vector cusolver_calls; + for (auto* hlo : computation->instructions()) { + if (hlo->opcode() == HloOpcode::kCholesky) { + cusolver_calls.push_back(hlo); + } + } + + if (cusolver_calls.empty()) { + return false; + } + + // Create a stream for us to do our work on. We don't really need to do any + // work, just allocate memory, but that's the cuSolver API. + se::Stream stream{stream_exec_}; + stream.Init(); + const auto device_ordinal = stream_exec_->device_ordinal(); + + // allocator either points to this->allocator_ or, if that's null, to a + // StreamExecutorMemoryAllocator for stream_exec_. + DeviceMemoryAllocator* allocator; + absl::optional se_allocator; + if (allocator_ != nullptr) { + allocator = allocator_; + } else { + se_allocator.emplace(stream_exec_->platform(), + absl::Span({stream_exec_})); + allocator = &*se_allocator; + } + ScratchAllocator scratch_allocator(device_ordinal, allocator); + + TF_ASSIGN_OR_RETURN(CusolverContext context, + CusolverContext::Create(&stream)); + + bool changed = false; + for (HloInstruction* instruction : cusolver_calls) { + TF_ASSIGN_OR_RETURN( + bool result, + RunOnInstruction(&context, &scratch_allocator, instruction)); + changed |= result; + } + return changed; +} + +CusolverRewriter::CusolverRewriter(se::StreamExecutor* stream_exec, + DeviceMemoryAllocator* allocator) + : stream_exec_(stream_exec), allocator_(allocator) {} + +StatusOr CusolverRewriter::Run(HloModule* module) { + bool changed = false; + for (HloComputation* computation : module->MakeNonfusionComputations()) { + TF_ASSIGN_OR_RETURN(bool result, RunOnComputation(computation)); + changed |= result; + } + return changed; +} + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/cusolver_rewriter.h b/tensorflow/compiler/xla/service/gpu/cusolver_rewriter.h new file mode 100644 index 0000000000000000000000000000000000000000..c82233188f7de1e188876f13465f7face76a0a8b --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/cusolver_rewriter.h @@ -0,0 +1,48 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUSOLVER_REWRITER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUSOLVER_REWRITER_H_ + +#include "tensorflow/compiler/xla/service/device_memory_allocator.h" +#include "tensorflow/compiler/xla/service/gpu/cusolver_context.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace xla { +namespace gpu { + +// Rewrites Cholesky calls into CustomCall HLOs that call into cuSolver. +class CusolverRewriter : public HloModulePass { + public: + CusolverRewriter(se::StreamExecutor* stream_exec, + DeviceMemoryAllocator* allocator); + absl::string_view name() const override { return "cusolver-rewriter"; } + + StatusOr Run(HloModule* module) override; + + private: + StatusOr RunOnComputation(HloComputation* computation); + + se::StreamExecutor* stream_exec_; // never null + DeviceMemoryAllocator* allocator_; // may be null +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUSOLVER_REWRITER_H_ diff --git a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc index 80ddb3e5cde708f535f71b75587171ed975f939c..b024bbe4b5e170b265a507e60486a806cc8cb8c3 100644 --- a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "llvm/IR/DerivedTypes.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" // IWYU pragma: no_include "llvm/IR/Attributes.gen.inc" @@ -191,39 +192,6 @@ StatusOr GpuElementalIrEmitter::EmitPowerOp( PrimitiveType lhs_input_type = op->operand(0)->shape().element_type(); PrimitiveType rhs_input_type = op->operand(1)->shape().element_type(); PrimitiveType output_type = op->shape().element_type(); - llvm::Type* llvm_ty = lhs_value->getType(); - - auto make_sqrt = [&, this]() -> StatusOr { - // NVPTX has four relevant square root instructions: - // sqrt.approx{.ftz}.f32 - // sqrt.rn{.ftz}.f32 - // sqrt.rn.f64 - // rsqrt.approx.f64 - // We rely on LLVM's NVPTX backend to pick the right one based on our - // fast-math options. (If fast-math is enabled, llvm may compute the 64-bit - // sqrt from the rsqrt approximation.) - return EmitLlvmIntrinsicMathCall("llvm.sqrt", {lhs_value}, {lhs_input_type}, - output_type); - }; - - const HloInstruction* rhs = op->operand(1); - if (IsFPLiteralWithValue(rhs, .5)) { - VLOG(10) << "emitting pow(A, .5) as sqrt(A): " << op->ToString(); - return make_sqrt(); - } - - if (IsFPLiteralWithValue(rhs, -.5)) { - VLOG(10) << "emitting pow(A, -.5) as 1/sqrt(A): " << op->ToString(); - // LLVM's NVPTX backend knows how to transform 1/sqrt(A) into the NVPTX - // rsqrt.approx instruction. - // - // TODO(jlebar): Does this happen with fastmath disabled? If not, should - // we force-enable it? - TF_ASSIGN_OR_RETURN(auto* sqrt, make_sqrt()); - return FDiv(llvm::ConstantFP::get(llvm_ty, 1), sqrt); - } - - VLOG(10) << "emitting pow as regular call to pow(): " << op->ToString(); return EmitLibdeviceMathCall("__nv_pow", {lhs_value, rhs_value}, {lhs_input_type, rhs_input_type}, output_type); } @@ -270,6 +238,16 @@ StatusOr GpuElementalIrEmitter::EmitPow(PrimitiveType prim_type, prim_type); } +StatusOr GpuElementalIrEmitter::EmitSqrt(PrimitiveType prim_type, + llvm::Value* value) { + return EmitLibdeviceMathCall("__nv_sqrt", {value}, {prim_type}, prim_type); +} + +StatusOr GpuElementalIrEmitter::EmitRsqrt(PrimitiveType prim_type, + llvm::Value* value) { + return EmitLibdeviceMathCall("__nv_rsqrt", {value}, {prim_type}, prim_type); +} + StatusOr GpuElementalIrEmitter::EmitAtan2(PrimitiveType prim_type, llvm::Value* lhs, llvm::Value* rhs) { @@ -293,6 +271,16 @@ StatusOr GpuElementalIrEmitter::EmitTanh(PrimitiveType prim_type, return FPCast(fast_tanh, value->getType()); } +StatusOr GpuElementalIrEmitter::EmitRoundNearestAfz( + PrimitiveType prim_type, llvm::Value* value) { + // Use libdevice __nv_round instead of llvm.round. This is to workaround a + // bug in the PTX backend, which implements llvm.round with PTX cvt.rni. + // When the llvm.round is fixed, we may still want to use __nv_round here as + // expanding the non-trivial implementation early while inlining allows better + // optimizations. + return EmitLibdeviceMathCall("__nv_round", {value}, {prim_type}, prim_type); +} + llvm::Value* GpuElementalIrEmitter::EmitDeviceFunctionCall( const string& callee_name, absl::Span operands, absl::Span input_types, PrimitiveType output_type, diff --git a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h index e8b56a39ce58b6aab35c1c977553c7ff7e753273..e9d08177ad979871890a32374657d8479c0cf669 100644 --- a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h +++ b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h @@ -76,6 +76,12 @@ class GpuElementalIrEmitter : public ElementalIrEmitter { StatusOr EmitExpm1(PrimitiveType prim_type, llvm::Value* value) override; + StatusOr EmitSqrt(PrimitiveType prim_type, + llvm::Value* value) override; + + StatusOr EmitRsqrt(PrimitiveType prim_type, + llvm::Value* value) override; + StatusOr EmitPow(PrimitiveType prim_type, llvm::Value* lhs, llvm::Value* rhs) override; @@ -85,6 +91,9 @@ class GpuElementalIrEmitter : public ElementalIrEmitter { StatusOr EmitTanh(PrimitiveType prim_type, llvm::Value* value) override; + StatusOr EmitRoundNearestAfz(PrimitiveType prim_type, + llvm::Value* value) override; + llvm::Value* EmitThreadId() override; private: diff --git a/tensorflow/compiler/xla/service/gpu/fusion_merger.cc b/tensorflow/compiler/xla/service/gpu/fusion_merger.cc index 91930eccdff94bb2fc85636f3a4b2d661c618d87..a225c29ce922097b7d45dffff5a18b5c9932fb2b 100644 --- a/tensorflow/compiler/xla/service/gpu/fusion_merger.cc +++ b/tensorflow/compiler/xla/service/gpu/fusion_merger.cc @@ -62,7 +62,7 @@ double CalculateBytesReadByFusionParameter(HloInstruction* param) { // Iterate through 'instructions' accumulating byte sizes of each instruction // shape. For each 'instruction' in 'instructions', if all users of - // 'instruction' are Slice instructions, accumuates the byte sizes of each + // 'instruction' are Slice instructions, accumulates the byte sizes of each // Slice for a more accurate estimate of bytes read. double bytes = 0.0; for (auto& instruction : instructions) { @@ -95,27 +95,6 @@ double CalculateBytesReadByFusionInstruction(HloInstruction* fusion) { return bytes; } -// Returns the flops to bytes transferred ratio of instruction 'fusion'. -double CalculateFlopsToBytesRatio(HloInstruction* fusion) { - CHECK_EQ(HloOpcode::kFusion, fusion->opcode()); - // Calculate total bytes transferred in/out. - double bytes = CalculateBytesReadByFusionInstruction(fusion); - // Add bytes written to root instructions buffer. - if (fusion->IsMultiOutputFusion()) { - for (auto& operand : fusion->fused_expression_root()->operands()) { - bytes += ShapeUtil::ByteSizeOf(operand->shape()); - } - } else { - bytes += ShapeUtil::ByteSizeOf(fusion->fused_expression_root()->shape()); - } - // Calculate flops for all fused instructions. Use a null shape size function - // because we don't care about bytes accessed by the ops. - HloCostAnalysis analysis([](const Shape& shape) { return 0; }); - TF_CHECK_OK(fusion->fused_expression_root()->Accept(&analysis)); - // Return flops / bytes. - return bytes > 0.0 ? analysis.flop_count() / bytes : analysis.flop_count(); -} - // Returns bytes transferred by instruction 'fusion', including the bytes // that would be read by all users. double GetCurrentBytesTransferred(HloInstruction* fusion) { @@ -169,7 +148,6 @@ class FusionInstructionMerger { int num_fail_not_loop_fusion_ = 0; int num_fail_merge_all_users_ = 0; int num_fail_expensive_fused_instruction_ = 0; - int num_fail_flops_to_byte_ratio_ = 0; int num_fail_net_bytes_transferred_ratio_ = 0; TF_DISALLOW_COPY_AND_ASSIGN(FusionInstructionMerger); @@ -190,15 +168,12 @@ Status FusionInstructionMerger::Run() { << " not_loop_fusion: " << num_fail_not_loop_fusion_ << " merge_all_users: " << num_fail_merge_all_users_ << " expensive_instruction: " << num_fail_expensive_fused_instruction_ - << " flops_to_byte_ratio: " << num_fail_flops_to_byte_ratio_ << " net_bytes_transferred: " << num_fail_net_bytes_transferred_ratio_ << " }"; return Status::OK(); } Status FusionInstructionMerger::HandleFusion(HloInstruction* fusion) { - VLOG(3) << "FusionInstructionMerger ENTRY fusion: " << fusion->name() - << " flops_to_bytes_ratio: " << CalculateFlopsToBytesRatio(fusion); ++total_visited_; // Skip 'fusion' instruction if there are no users into which we can merge. if (fusion->users().empty()) { @@ -256,15 +231,6 @@ Status FusionInstructionMerger::HandleFusion(HloInstruction* fusion) { return Status::OK(); } - // Skip 'fusion' instruction if its flops to bytes transferred ratio - // exceeds the threshold value. - if (CalculateFlopsToBytesRatio(fusion) > - FusionMerger::GetThresholdFlopsToBytesRatio()) { - VLOG(3) << "Not merging " << fusion->name() - << ": flops-to-bytes ratio is not favorable."; - ++num_fail_flops_to_byte_ratio_; - return Status::OK(); - } // Skip 'fusion' instruction if merging it into all users would result in a // net increase in bytes transferred (currently allowing the net bytes // transferred to be exceeded up to ~10% in exhange for eliminating the @@ -288,7 +254,6 @@ Status FusionInstructionMerger::HandleFusion(HloInstruction* fusion) { } ++total_merged_; VLOG(2) << "Merged fusion instruction: " << fusion->name() - << " flops_to_bytes_ratio: " << CalculateFlopsToBytesRatio(fusion) << " merged_to_current_bytes_ratio: " << merged_to_current_bytes_ratio << " into users { " << absl::StrJoin(users, ", ", diff --git a/tensorflow/compiler/xla/service/gpu/fusion_merger.h b/tensorflow/compiler/xla/service/gpu/fusion_merger.h index f19996edfe3dd923aa686a19621ce28a4aed5a45..a49d68002f8de5bb5640731f3cd31572593ee837 100644 --- a/tensorflow/compiler/xla/service/gpu/fusion_merger.h +++ b/tensorflow/compiler/xla/service/gpu/fusion_merger.h @@ -37,8 +37,6 @@ class FusionMerger : public HloModulePass { absl::string_view name() const override { return "fusion merger"; } StatusOr Run(HloModule* module) override; - - static double GetThresholdFlopsToBytesRatio() { return 1.0; } }; } // namespace gpu diff --git a/tensorflow/compiler/xla/service/gpu/fusion_merger_test.cc b/tensorflow/compiler/xla/service/gpu/fusion_merger_test.cc index 7cc869ed9e89688d6ea06428a7bade3ebe55ea23..96ec060d533be66d35d7374a7eb5f93d8b732195 100644 --- a/tensorflow/compiler/xla/service/gpu/fusion_merger_test.cc +++ b/tensorflow/compiler/xla/service/gpu/fusion_merger_test.cc @@ -99,62 +99,6 @@ ENTRY MergeSharedFusionInstruction.Computation0 { EXPECT_EQ(7, operand2->fused_instruction_count()); } -// Tests that we do not merge a fusion instruction that above flops to bytes -// threshold. -// -// Fusion2 is not merged because it exceeds the threshold flops-to-bytes ratio. -TEST_F(FusionMergerTest, FlopsToBytesRatioThresholdExceeded) { - auto module = ParseHloString(R"( -HloModule FlopsToBytesRatioThresholdExceeded - -comp.2 { - state.param_1.1 = (f32[4]{0}, f32[4]{0}) parameter(0) - get-tuple-element.3 = f32[4]{0} get-tuple-element(state.param_1.1), index=0 - get-tuple-element.4 = f32[4]{0} get-tuple-element(state.param_1.1), index=2 - multiply.29 = f32[4]{0} multiply(get-tuple-element.3, get-tuple-element.4) - multiply.30 = f32[4]{0} multiply(get-tuple-element.3, multiply.29) - multiply.31 = f32[4]{0} multiply(get-tuple-element.3, multiply.30) - multiply.32 = f32[4]{0} multiply(get-tuple-element.3, multiply.31) - multiply.33 = f32[4]{0} multiply(get-tuple-element.3, multiply.32) - multiply.34 = f32[4]{0} multiply(get-tuple-element.3, multiply.33) - multiply.35 = f32[4]{0} multiply(get-tuple-element.3, multiply.34) - multiply.36 = f32[4]{0} multiply(get-tuple-element.3, multiply.35) - multiply.37 = f32[4]{0} multiply(get-tuple-element.3, multiply.36) - multiply.38 = f32[4]{0} multiply(get-tuple-element.3, multiply.37) - multiply.39 = f32[4]{0} multiply(get-tuple-element.3, multiply.38) - multiply.40 = f32[4]{0} multiply(get-tuple-element.3, multiply.39) - ROOT multiply.41 = f32[4]{0} multiply(get-tuple-element.3, multiply.40) -} - -comp.1 { - multiply.12.param_1.1 = f32[4]{0} parameter(1) - constant.param_1.3 = f32[4]{0} parameter(0) - add.3 = f32[4]{0} add(multiply.12.param_1.1, constant.param_1.3) - ROOT multiply.16 = f32[4]{0} multiply(add.3, constant.param_1.3) -} - -comp { - multiply.12.param_1 = f32[4]{0} parameter(1) - constant.param_1.1 = f32[4]{0} parameter(0) - multiply.15 = f32[4]{0} multiply(multiply.12.param_1, constant.param_1.1) - ROOT add.2 = f32[4]{0} add(multiply.15, constant.param_1.1) -} - -ENTRY FlopsToBytesRatioThresholdExceeded.Computation1 { - constant = f32[4]{0} constant({1, 1, 1, 1}) - state = (f32[4]{0}, f32[4]{0}) parameter(0) - fusion.2 = f32[4]{0} fusion(state), kind=kLoop, calls=comp.2 - fusion.3 = f32[4]{0} fusion(constant, fusion.2), kind=kLoop, calls=comp.1 - fusion.4 = f32[4]{0} fusion(constant, fusion.2), kind=kLoop, calls=comp - ROOT tuple = (f32[4]{0}, f32[4]{0}) tuple(fusion.3, fusion.4) -})") - .ValueOrDie(); - // Run fusion merger pass, which should detect that the flops/bytes of the - // shared fusion instruction exceeds the threshold ratio, and therefore - // cannot be merged with other fusion instructions. - EXPECT_FALSE(FusionMerger().Run(module.get()).ValueOrDie()); -} - // Tests that threshold for bytes transferred if merged is exceeded. // // Fusion2 is not merged because it exceeds the threshold bytes transferred. diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc index 3ed6553f9205803cfa17772b890c449cfb457c89..6b9cbdd94b334ab7a4f61a4e3e43250ed9648cd0 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc @@ -142,6 +142,16 @@ bool IsCustomCallToDnnConvolution(const HloInstruction& hlo) { target == kCudnnConvBiasActivationForwardCallTarget; } +const char* const kCusolverCholeskyCallTarget = "__cusolver$cholesky"; + +bool IsCustomCallToCusolver(const HloInstruction& hlo) { + if (hlo.opcode() != HloOpcode::kCustomCall) { + return false; + } + const auto& target = hlo.custom_call_target(); + return target == kCusolverCholeskyCallTarget; +} + bool ImplementedAsLibraryCall(const HloInstruction& hlo) { return ImplementedAsGemm(hlo) || IsCustomCallToDnnBatchNorm(hlo) || IsCustomCallToDnnConvolution(hlo); diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h index ebf4d926b7a280e10b09a2532caba7ad6ab3ceb2..f1a7aabb4db57b6818b29bdde73d87f0706f2827 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h +++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h @@ -131,6 +131,19 @@ extern const char* const kCudnnConvBiasActivationForwardCallTarget; // kConvolution opcode. bool IsCustomCallToDnnConvolution(const HloInstruction& hlo); +// Returns true if `hlo` will be implemented as a call to a cuSolver routine. +// +// This returns true if `hlo` is a CustomCall HLO with a call target equal to +// one of the kCusolver... constants, but returns *false* for HLOs with +// say, a kCholesky opcode. +bool IsCustomCallToCusolver(const HloInstruction& hlo); + +// Cholesky decomposition. Takes a (batched) matrix as input, and returns a +// tuple of (result, workspace, info), where result is the result of the +// Cholesky decomposition, workspace is scratch space for cuSolver, and info +// is a success/failure code per batch element. +extern const char* const kCusolverCholeskyCallTarget; + // Returns true if `hlo` will be implemented as a library call, e.g. cuBLAS gemm // or cuDNN convolution. bool ImplementedAsLibraryCall(const HloInstruction& hlo); diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc index 0cc65ebb52737aa9bb8866eb07278a2319aa797b..61aa981d77924a3c92b3df35fa5cbdcd665536be 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc @@ -39,6 +39,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor.h" #include "tensorflow/compiler/xla/service/gpu/buffer_allocations.h" +#include "tensorflow/compiler/xla/service/gpu/cholesky_thunk.h" #include "tensorflow/compiler/xla/service/gpu/conditional_thunk.h" #include "tensorflow/compiler/xla/service/gpu/convolution_thunk.h" #include "tensorflow/compiler/xla/service/gpu/copy_thunk.h" @@ -54,6 +55,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/ir_emitter_context.h" #include "tensorflow/compiler/xla/service/gpu/kernel_thunk.h" #include "tensorflow/compiler/xla/service/gpu/memset_thunk.h" +#include "tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.h" #include "tensorflow/compiler/xla/service/gpu/outfeed_thunk.h" #include "tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h" #include "tensorflow/compiler/xla/service/gpu/partition_assignment.h" @@ -74,6 +76,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/llvm_ir/sort_util.h" #include "tensorflow/compiler/xla/service/llvm_ir/tuple_ops.h" #include "tensorflow/compiler/xla/service/name_uniquer.h" +#include "tensorflow/compiler/xla/service/pattern_matcher.h" #include "tensorflow/compiler/xla/service/while_loop_analysis.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -102,6 +105,8 @@ using absl::StrCat; using llvm_ir::IrArray; using llvm_ir::IrName; +namespace m = match; + // If a dimensions is smaller than this, untiled transposition may be more // efficient. const int64 kMinDimensionToTransposeTiled = 16; @@ -476,6 +481,51 @@ Status IrEmitterUnnested::HandleCustomCall(HloInstruction* custom_call) { return Status::OK(); } + if (custom_call->custom_call_target() == kCusolverCholeskyCallTarget) { + TF_ASSIGN_OR_RETURN(CholeskyOptions options, + custom_call->backend_config()); + + const Shape& shape = custom_call->operand(0)->shape(); + int ndim = shape.dimensions_size(); + CHECK_GE(ndim, 2); + int64 n = shape.dimensions(ndim - 1); + + const auto& dims = shape.dimensions(); + int64 batch_size = std::accumulate(dims.begin(), dims.end() - 2, int64{1}, + [](int64 a, int64 b) { return a * b; }); + + auto operand_buffer = GetAllocationSlice(*custom_call->operand(0)); + + const auto& assn = ir_emitter_context_->buffer_assignment(); + auto a_buffer = assn.GetUniqueSlice(custom_call, {0}).ValueOrDie(); + auto workspace_buffer = assn.GetUniqueSlice(custom_call, {1}).ValueOrDie(); + auto info_buffer = assn.GetUniqueSlice(custom_call, {2}).ValueOrDie(); + + std::vector> thunks; + + if (operand_buffer != a_buffer) { + thunks.push_back(absl::make_unique( + /*source_address=*/operand_buffer, + /*destination_buffer=*/a_buffer, + /*mem_size=*/ShapeUtil::ByteSizeOf(shape), custom_call)); + } + + thunks.push_back(absl::make_unique( + options, a_buffer, workspace_buffer, info_buffer, + custom_call->operand(0)->shape().element_type(), batch_size, n, + custom_call)); + + // Elide the sequential thunk if there's no copy. + if (thunks.size() == 1) { + AddThunkToThunkSequence(std::move(thunks[0])); + } else { + AddThunkToThunkSequence( + absl::make_unique(std::move(thunks), custom_call)); + } + + return Status::OK(); + } + return IrEmitter::HandleCustomCall(custom_call); } @@ -926,13 +976,12 @@ Status IrEmitterUnnested::HandleWhile(HloInstruction* xla_while) { condition->root_instruction()->shape().element_type() == PRED) << "While condition computation must return bool"; // Build ForThunk for conformant while loops, otherwise build WhileThunk. - // TODO(b/112163966): Move trip count computation earlier in the pipeline. - if (auto loop_trip_count = ComputeWhileLoopTripCount(xla_while)) { - AddThunkToThunkSequence(BuildForThunk(xla_while, *loop_trip_count)); - VLOG(3) << "Built ForThunk for while: " << xla_while->name(); + auto config = xla_while->backend_config(); + if (config.ok() && config.ValueOrDie().has_known_trip_count()) { + AddThunkToThunkSequence( + BuildForThunk(xla_while, config.ValueOrDie().known_trip_count().n())); } else { AddThunkToThunkSequence(BuildWhileThunk(xla_while)); - VLOG(3) << "Built WhileThunk for while: " << xla_while->name(); } return Status::OK(); } @@ -1319,11 +1368,55 @@ Status IrEmitterUnnested::HandleTupleSelect(HloInstruction* tuple_select) { return IrEmitter::HandleTupleSelect(tuple_select); } +namespace { + +bool IsScalarAddComputation(HloComputation* computation) { + return Match(computation->root_instruction(), + m::AddAnyOrder(m::Parameter(0), m::Parameter(1)) + .WithShape(m::Shape().IsEffectiveScalar())); +} + +} // namespace + Status IrEmitterUnnested::HandleAllReduce(HloInstruction* crs) { + VLOG(2) << "AllReduce; replica count: " << hlo_module_config_.replica_count() + << "; operand count: " << crs->operand_count() + << "; NCCL is enabled: " << NcclAllReduceThunk::NcclIsEnabled(); + + // Note the replica_count == 1 case is handled via device-to-device copy + // below. + bool should_use_nccl_thunk = + hlo_module_config_.replica_count() > 1 && + crs->IsCrossReplicaAllReduce() && + crs->operand_count() == 1 && // One array to reduce. + crs->operand(0)->shape().element_type() == F32 && + // Check the computation is a summation. + IsScalarAddComputation(crs->to_apply()); + + if (should_use_nccl_thunk) { + CHECK(crs->operand(0)->shape().IsArray()) + << "Operands to all-reduce must be arrays: " << crs->ToString(); + AddThunkToThunkSequence(absl::make_unique( + /*replica_count=*/hlo_module_config_.replica_count(), + /*elements=*/ShapeUtil::ElementsIn(crs->operand(0)->shape()), + /*source_address=*/GetAllocationSlice(*crs->operand(0)), + /*destination_buffer=*/GetAllocationSlice(*crs), crs)); + return Status::OK(); + } + if (hlo_module_config_.replica_count() != 1) { - // TODO(b/33011107): Support nontrivial cross replica sum on GPU. - return Unimplemented( - "AllReduce with >1 replica is not implemented on GPU."); + // TODO(b/33011107): Support more AllReduce configurations on GPU. + string message = absl::StrFormat( + "Requested AllReduce not implemented on GPU; replica_count: %d; " + "operand_count: %d; IsCrossReplicaAllReduce: %d; NCCL support: %d", + hlo_module_config_.replica_count(), crs->operand_count(), + crs->IsCrossReplicaAllReduce(), NcclAllReduceThunk::NcclIsEnabled()); + if (crs->operand_count() > 0) { + absl::StrAppendFormat( + &message, "; first operand array element-type: %s", + PrimitiveType_Name(crs->operand(0)->shape().element_type())); + } + return Unimplemented("%s", message); } // CRS with one operand and one replica is simply the identity function. @@ -1977,41 +2070,32 @@ Status CheckWhileBuffersShareAllocation( // Checks that the buffers used in a conditional instruction are shared with the // operands and result as follows: // * The result buffer of the conditional should share the allocation with the -// result buffers of the true and false computations. -// * The buffer of operand 1 should share the allocation with the buffer of -// the parameter 0 instruction of the true computation. -// * The buffer of operand 2 should share the allocation with the buffer of -// the parameter 0 instruction of the false computation. +// result buffers of each branch computation. +// * The buffer of operand b+1 should share the allocation with the buffer of +// the parameter 0 instruction of the b'th computation. Status CheckConditionalBuffersShareAllocation( const HloInstruction* conditional, const BufferAssignment& buffer_assignment) { TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( conditional->shape(), [&](const Shape& /*subshape*/, const ShapeIndex& index) -> Status { - TF_RETURN_IF_ERROR(CheckHloBuffersShareAllocation( - conditional, conditional->true_computation()->root_instruction(), - index, buffer_assignment)); - TF_RETURN_IF_ERROR(CheckHloBuffersShareAllocation( - conditional, conditional->false_computation()->root_instruction(), - index, buffer_assignment)); + for (auto branch_computation : conditional->branch_computations()) { + TF_RETURN_IF_ERROR(CheckHloBuffersShareAllocation( + conditional, branch_computation->root_instruction(), index, + buffer_assignment)); + } return Status::OK(); })); - TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( - conditional->operand(1)->shape(), - [&](const Shape& /*subshape*/, const ShapeIndex& index) -> Status { - return CheckHloBuffersShareAllocation( - conditional->operand(1), - conditional->true_computation()->parameter_instruction(0), index, - buffer_assignment); - })); - TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( - conditional->operand(2)->shape(), - [&](const Shape& /*subshape*/, const ShapeIndex& index) -> Status { - return CheckHloBuffersShareAllocation( - conditional->operand(2), - conditional->false_computation()->parameter_instruction(0), index, - buffer_assignment); - })); + for (int j = 0; j < conditional->branch_count(); ++j) { + TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( + conditional->operand(j + 1)->shape(), + [&](const Shape& /*subshape*/, const ShapeIndex& index) -> Status { + return CheckHloBuffersShareAllocation( + conditional->operand(j + 1), + conditional->branch_computation(j)->parameter_instruction(0), + index, buffer_assignment); + })); + } return Status::OK(); } @@ -2064,22 +2148,20 @@ std::unique_ptr IrEmitterUnnested::BuildConditionalThunk( TF_CHECK_OK(CheckConditionalBuffersShareAllocation( hlo, ir_emitter_context_->buffer_assignment())); - HloComputation* true_computation = hlo->true_computation(); - IrEmitterUnnested ir_emitter_true(hlo_module_config_, true_computation, - ir_emitter_context_); - TF_CHECK_OK(true_computation->Accept(&ir_emitter_true)); - - HloComputation* false_computation = hlo->false_computation(); - IrEmitterUnnested ir_emitter_false(hlo_module_config_, false_computation, - ir_emitter_context_); - TF_CHECK_OK(false_computation->Accept(&ir_emitter_false)); + std::vector branch_operands; + std::vector branch_thunks; + for (int j = 0; j < hlo->branch_count(); ++j) { + branch_operands.emplace_back(GetAllocationSlice(*hlo->operand(j + 1))); + HloComputation* branch_computation = hlo->branch_computation(j); + IrEmitterUnnested ir_emitter(hlo_module_config_, branch_computation, + ir_emitter_context_); + TF_CHECK_OK(branch_computation->Accept(&ir_emitter)); + branch_thunks.push_back(std::move(*ir_emitter.ConsumeThunkSequence())); + } return absl::make_unique( - GetAllocationSlice(*hlo->operand(0)), - GetAllocationSlice(*hlo->operand(1)), - GetAllocationSlice(*hlo->operand(2)), - std::move(*ir_emitter_true.ConsumeThunkSequence()), - std::move(*ir_emitter_false.ConsumeThunkSequence()), hlo); + GetAllocationSlice(*hlo->operand(0)), branch_operands, + std::move(branch_thunks), hlo); } Status IrEmitterUnnested::EmitTargetElementLoopInThunk( diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h index f85e18bbf0798ef3d5b87e81d287d8aed691dfc4..9890ce122dfdc7444d769b6eb695a7c0932408c3 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h @@ -320,6 +320,9 @@ class IrEmitterUnnested : public IrEmitter { // Returns a FftThunk that calls cuFFT to implement `inst`. std::unique_ptr BuildFftThunk(const HloInstruction* inst); + // Returns a CholeskyThunk that calls cuSolver to implement `inst`. + std::unique_ptr BuildCholeskyThunk(const HloInstruction* inst); + // Returns a TriangularSolveThunk that calls cuBlas to implement `inst`. std::unique_ptr BuildTriangularSolveThunk(const HloInstruction* inst); @@ -356,9 +359,9 @@ class IrEmitterUnnested : public IrEmitter { std::unique_ptr BuildForThunk(const HloInstruction* hlo, const int64 loop_limit); - // Returns a ConditionalThunk that executes the thunk sequence for - // 'true_computation' or 'false_computation' depending on the value of the - // predicate in the given conditional instruction. + // Returns a ConditionalThunk which executes the thunk sequence for the + // 'branch_computation' corresponding to the predicate/branch_index of the + // given conditional instruction. std::unique_ptr BuildConditionalThunk(const HloInstruction* hlo); Status Postprocess(HloInstruction* hlo) override; diff --git a/tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.cc b/tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.cc new file mode 100644 index 0000000000000000000000000000000000000000..3051db3af4ae4380e4a38f50ad8ebc89642e645f --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.cc @@ -0,0 +1,356 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.h" + +#include "tensorflow/compiler/xla/util.h" + +#if GOOGLE_CUDA +#include "absl/synchronization/blocking_counter.h" +#include "third_party/nccl/nccl.h" +#include "tensorflow/core/lib/core/blocking_counter.h" +#include "tensorflow/core/platform/mutex.h" +#include "tensorflow/stream_executor/cuda/cuda_activation.h" +#endif + +namespace xla { +namespace gpu { + +/* static */ bool NcclAllReduceThunk::NcclIsEnabled() { +#if GOOGLE_CUDA + return true; +#else + return false; +#endif +} + +#if GOOGLE_CUDA +namespace { + +// GPU-replica-driving host threads (i.e. the threads that call +// GpuExecutable::Execute) build up this structure to describe their +// participating replica, and then call to +// GlobalRendezvousManager::SubmitParticipant. +struct ParticipantData { + // Number of replicas particiating in the AllReduce. + int64 replica_count; + + int64 element_count; + int64 device_ordinal; + int64 generation_counter; + + // TODO(b/125951860): We should vet that we're buffer allocating such that + // source_buffer == destination_buffer if that avoids a NCCL copy (will depend + // on how well the NCCL in-place implementation performs vs the out-of-place + // implementation). + se::DeviceMemoryBase source_data; + se::DeviceMemoryBase destination_data; + se::Stream* stream; + + NcclAllReduceThunk* originator; + + string ToString() const { + return absl::StrFormat( + "ParticipantData{replica_count=%d, element_count=%d, " + "device_ordinal=%d, generation_counter=%d, stream=%p, originator=%p}", + replica_count, element_count, device_ordinal, generation_counter, + stream, originator); + } +}; + +// Class that gets instantiated as a singleton in GetGlobalRendezvous() to +// coordinate participating threads in performing an AllReduce operation. +// +// This manager is responsible for establishing communication channels and +// ultimately enqueueing the NCCL library operation onto the participating +// streams. +class GlobalRendezvousManager { + public: + // The GpuExecutable-executing threads call this in order to a) establish the + // all-reduce rendezvous and b) enqueue the AllReduce operation on the caller + // thread's associated stream (given in "participant"). + // + // Implementation note: since the rendezvous we're creating here is global, we + // try to be paranoid about the fact that the *correct* one is happening. In + // an ideal world we'd have some StreamExecutor se::Platform level construct + // that we could use for cross-device networking primitives (e.g. via a + // NetworkSupport interface) that could be shared between TensorFlow and XLA, + // but this is a reasonable stopgap measure to get multi-GPU-replica up and + // running properly for single-host, single-concurrent-XLA-module usage. + Status SubmitParticipant(ParticipantData participant); + + // Returns the current generation number of AllReduce operations. + // (Currently one AllReduce operation occurs per generation.) + int64 GetCurrentGeneration() { + tensorflow::mutex_lock lock(mutex_); + return current_generation_; + } + + private: + // Called by the primary thread to set up the communication links. + // + // TODO(b/125951860): This performs lots of (presumably) unnecessary host-side + // synchronization so that we can be paranoid about semantics in the earliest + // implementation. In the limit we should only need to synchronize host + // replica threads when the "number of replicas" or "participating device + // ordinals" change, to set up a new NCCL "communication" context, at which + // point we can enqueue onto device streams without host synchronization in + // our code -- this will likely be helpful for "lots of little AllReduce" + // cases. + Status InitializeCommunicationChannels() EXCLUSIVE_LOCKS_REQUIRED(mutex_); + + // Called when all necessary participants are present, the functionality + // that's implemented by all executing threads lives in here. + Status DoAllReduce(ParticipantData data, ncclComm_t comm); + + // Puts all state back into a "reset" state for the next generation of + // AllReduce requests. + void DeinitializeGeneration() EXCLUSIVE_LOCKS_REQUIRED(mutex_) { + for (ncclComm_t& comm : comms_) { + ncclCommDestroy(comm); + } + comms_.clear(); + participants_.clear(); + current_generation_++; + initialized_ = false; + done_ = absl::nullopt; + } + + tensorflow::mutex mutex_; + tensorflow::condition_variable all_participants_present_; + tensorflow::condition_variable deinitialized_; + + // Communication handles that correspond to the participants below. + std::vector comms_ GUARDED_BY(mutex_); + + Status initialize_status_ GUARDED_BY(mutex_); + std::vector participants_ GUARDED_BY(mutex_); + int64 current_generation_ GUARDED_BY(mutex_) = 0; + bool initialized_ GUARDED_BY(mutex_) = false; + + // The participating threads wait for this to count down in order to know we + // can begin the teardown process. + absl::optional done_; +}; + +Status GlobalRendezvousManager::SubmitParticipant(ParticipantData participant) { + auto all_participants_present = [this, &participant]() + EXCLUSIVE_LOCKS_REQUIRED(mutex_) -> bool { + return participants_.size() >= participant.replica_count; + }; + + // We remember the participant index at which we are inserted and use that + // same index for referring to auxiliary metadata (e.g. the ncclComm_t handle + // index) below. + int64 index; + + { + tensorflow::mutex_lock lock(mutex_); + + // Spot check for consistent replica counts among submitting threads. + if (!participants_.empty() && + (participants_.back().replica_count != participant.replica_count || + participants_.back().originator != participant.originator)) { + return InvalidArgument( + "Running two XLA modules with AllReduces in parallel is not " + "supported. It is possible this is due to a bug where were try to " + "run two different AllReduces from the same module at once. " + "(Attempted a rendezvous with a different replica count from other " + "participants; existing: %s; submitted: %s)", + participants_.back().ToString(), participant.ToString()); + } + index = participants_.size(); + participants_.push_back(participant); + + if (all_participants_present()) { + all_participants_present_.notify_all(); + } + } + + // We pull into our thread a) the communication handle and b) whether we're + // the "primary" thread for this rendezvous -- the "primary" thread has some + // additional responsibilities for setup/teardown. + ncclComm_t comm; + bool primary; + + { + tensorflow::mutex_lock lock(mutex_); + while (!all_participants_present()) { + // Once all the participants have arrived, all participating threads will + // cross this barrier, though only (the first) one will be the "primary". + all_participants_present_.wait(lock); + } + + // Somebody will be the first -- that thread has some additional + // responsibilities. + primary = !initialized_; + + CHECK_EQ(participant.generation_counter, current_generation_); + + // Bump the generation counter so the other threads know we've completed the + // global rendezvous and have set up the AllReduce. + if (primary) { + VLOG(3) << "Primary initializing accounting data."; + initialized_ = true; + done_.emplace(participant.replica_count); + initialize_status_ = InitializeCommunicationChannels(); + VLOG(3) << "Done initializing communication channels; status: " + << initialize_status_; + if (!initialize_status_.ok()) { + DeinitializeGeneration(); + } + } + + if (!initialize_status_.ok()) { + // TODO(b/125951860): If this fails once, it will fail forever. + return initialize_status_; + } + + comm = comms_[index]; + + // Drop the lock at the end of scope so other participants may enter. + } + + VLOG(3) << "Performing all reduce from device ordinal: " + << participant.device_ordinal; + + Status all_reduce_status = DoAllReduce(participant, comm); + + VLOG(3) << "Waiting for all participants to complete enqueue."; + + done_->DecrementCount(); + + if (primary) { + // Primary thread clears out the AllReduce state when everybody is done to + // make it clean-slate for any subsequent AllReduce request (e.g. number of + // replicas may change in the next request). + // + // Note surrounding TODOs for only reinitializing this when the replica + // count / participants actually change -- lots of "playing it safe" + // happening in this first cut. + done_->Wait(); + VLOG(3) << "All participants completed enqueue."; + VLOG(3) << "Primary thread clearing."; + tensorflow::mutex_lock lock(mutex_); + DeinitializeGeneration(); + VLOG(3) << "Generation is now: " << current_generation_; + deinitialized_.notify_all(); + } else { + VLOG(3) << "Waiting to deinitialize."; + tensorflow::mutex_lock lock(mutex_); + while (initialized_) { + deinitialized_.wait(lock); + } + } + + VLOG(3) << "Returning status: " << all_reduce_status; + return all_reduce_status; +} + +Status GlobalRendezvousManager::InitializeCommunicationChannels() { + std::vector ordinals; + for (ParticipantData& data : participants_) { + ordinals.push_back(data.device_ordinal); + } + comms_.resize(ordinals.size()); + VLOG(3) << "Participants: " << participants_.size() + << "; initializing comms."; + ncclResult_t result = ncclCommInitAll(comms_.data(), comms_.size(), + /*devlist=*/ordinals.data()); + if (result != ncclSuccess) { + comms_.clear(); + return InternalError( + "Failed to initialize NCCL communication channels for %d participants: " + "%s", + participants_.size(), ncclGetErrorString(result)); + } + return Status::OK(); +} + +Status GlobalRendezvousManager::DoAllReduce(ParticipantData participant, + ncclComm_t comm) { + se::StreamExecutor* executor = participant.stream->parent(); + se::cuda::ScopedActivateExecutorContext scoped_context(executor); + cudaStream_t* cu_stream = reinterpret_cast( + participant.stream->implementation()->GpuStreamMemberHack()); + VLOG(3) << "Using stream pointer: " << cu_stream + << " on device: " << participant.device_ordinal; + void* send_buffer = participant.source_data.opaque(); + void* recv_buffer = participant.destination_data.opaque(); + ncclResult_t result = ncclAllReduce(send_buffer, recv_buffer, + /*count=*/participant.element_count, + /*datatype=*/ncclFloat, + /*op=*/ncclSum, + /*comm=*/comm, + /*stream=*/*cu_stream); + TF_RET_CHECK(ncclSuccess == result) + << "Failed to perform all-reduce: " << ncclGetErrorString(result); + + VLOG(3) << "Done performing all reduce for ordinal: " + << participant.device_ordinal; + + return Status::OK(); +} + +static GlobalRendezvousManager* GetGlobalRendezvous() { + static auto* manager = new GlobalRendezvousManager; + return manager; +} + +} // namespace + +Status NcclAllReduceThunk::ExecuteOnStream( + const BufferAllocations& buffer_allocations, se::Stream* stream, + HloExecutionProfiler* profiler) { + auto* global_rendezvous = GetGlobalRendezvous(); + + ParticipantData participant; + participant.replica_count = replica_count_; + participant.element_count = element_count_; + participant.device_ordinal = stream->parent()->device_ordinal(); + participant.generation_counter = global_rendezvous->GetCurrentGeneration(); + participant.source_data = buffer_allocations.GetDeviceAddress(source_buffer_); + participant.destination_data = + buffer_allocations.GetDeviceAddress(destination_buffer_); + participant.stream = stream; + participant.originator = this; + + return GetGlobalRendezvous()->SubmitParticipant(std::move(participant)); +} +#else + +Status NcclAllReduceThunk::ExecuteOnStream( + const BufferAllocations& buffer_allocations, se::Stream* stream, + HloExecutionProfiler* profiler) { + return Unimplemented( + "NCCL support is not available: this binary was not built with a CUDA " + "compiler, which is necessary to build the NCCL source library."); +} + +#endif // GOOGLE_CUDA + +NcclAllReduceThunk::NcclAllReduceThunk( + int64 replica_count, int64 element_count, + const BufferAllocation::Slice& source_buffer, + const BufferAllocation::Slice& destination_buffer, + const HloInstruction* all_reduce) + : Thunk(Thunk::kNcclAllReduce, all_reduce), + replica_count_(replica_count), + element_count_(element_count), + source_buffer_(source_buffer), + destination_buffer_(destination_buffer) {} + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.h b/tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.h new file mode 100644 index 0000000000000000000000000000000000000000..1a8d1356c0023e2c7f49c3731693e10beba54a6d --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/nccl_all_reduce_thunk.h @@ -0,0 +1,62 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_NCCL_ALL_REDUCE_THUNK_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_NCCL_ALL_REDUCE_THUNK_H_ + +#include "tensorflow/compiler/xla/service/buffer_assignment.h" +#include "tensorflow/compiler/xla/service/gpu/buffer_allocations.h" +#include "tensorflow/compiler/xla/service/gpu/hlo_execution_profiler.h" +#include "tensorflow/compiler/xla/service/gpu/thunk.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { +namespace gpu { + +// Thunk that performs a NCCL-based All-Reduce among CUDA GPU-based replicas. +class NcclAllReduceThunk : public Thunk { + public: + // Returns whether NCCL operations appear possible to perform; e.g. if we + // haven't done a build with the CUDA compiler enabled, we can't compile the + // NCCL header, and thus this will be false. + // + // When this is false, the ExecuteOnStream() call will simply return a status + // error. + static bool NcclIsEnabled(); + + // TODO(b/125951860): Plumb more datatypes / reduction operators. Initial + // implementation is simply F32 summation. + NcclAllReduceThunk(int64 replica_count, int64 element_count, + const BufferAllocation::Slice& source_buffer, + const BufferAllocation::Slice& destination_buffer, + const HloInstruction* all_reduce); + + Status ExecuteOnStream(const BufferAllocations& buffer_allocations, + se::Stream* stream, + HloExecutionProfiler* profiler) override; + + private: + const int64 replica_count_; + const int64 element_count_; + const BufferAllocation::Slice source_buffer_; + const BufferAllocation::Slice destination_buffer_; +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_NCCL_ALL_REDUCE_THUNK_H_ diff --git a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc index 6e00e4b4ff8c493f00fae3355215fb13fb5f4f10..ad75e2dd4341698d7bf0682382fce54c5154d753 100644 --- a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc +++ b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc @@ -46,6 +46,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/cudnn_conv_padding_legalization.h" #include "tensorflow/compiler/xla/service/gpu/cudnn_conv_rewriter.h" #include "tensorflow/compiler/xla/service/gpu/cudnn_fused_conv_rewriter.h" +#include "tensorflow/compiler/xla/service/gpu/cusolver_rewriter.h" #include "tensorflow/compiler/xla/service/gpu/fusion_merger.h" #include "tensorflow/compiler/xla/service/gpu/gpu_constants.h" #include "tensorflow/compiler/xla/service/gpu/gpu_copy_insertion.h" @@ -87,6 +88,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/tuple_simplifier.h" #include "tensorflow/compiler/xla/service/while_loop_constant_sinking.h" #include "tensorflow/compiler/xla/service/while_loop_simplifier.h" +#include "tensorflow/compiler/xla/service/while_loop_trip_count_annotator.h" #include "tensorflow/compiler/xla/service/zero_sized_hlo_elimination.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" @@ -247,15 +249,27 @@ Status OptimizeHloModule(HloModule* hlo_module, se::StreamExecutor* stream_exec, TransposeFolding::NeverFoldTranspose); pipeline.AddPass(/*is_layout_sensitive=*/false); pipeline.AddPass(); + + // Run WhileLoopTripCountAnnotator at the end of the simplification + // pipeline, before layout assignment and fusion. This pass does some + // pattern-matching on while bodies/conditions, and this is where the HLO is + // "nicest". + // + // It's important that we don't make semantic changes (e.g. unrolling) to + // any `while` loops after this point, because otherwise the trip-count + // annotations added by this pass may not be correct after the + // modifications. + pipeline.AddPass(); TF_RETURN_IF_ERROR(pipeline.Run(hlo_module).status()); } { // Convert convolutions into CustomCalls to cudnn, then canonicalize them - // (CudnnConvPaddingLegalization). + // (CudnnConvPaddingLegalization). Also expand cuSolver calls. HloPassPipeline pipeline("conv_canonicalization"); pipeline.AddInvariantChecker(/*layout_sensitive=*/false, /*allow_mixed_precision=*/false); + pipeline.AddPass(stream_exec, device_allocator); pipeline.AddPass(); pipeline.AddPass(); pipeline.AddPass(); @@ -328,6 +342,7 @@ Status OptimizeHloModule(HloModule* hlo_module, se::StreamExecutor* stream_exec, // wouldn't be able to simplify away the new_tuple bits. pipeline.AddPass(stream_exec, device_allocator, compiler); + // Clean up new_tuple described above. pipeline.AddPass(); diff --git a/tensorflow/compiler/xla/service/gpu/partition_assignment.cc b/tensorflow/compiler/xla/service/gpu/partition_assignment.cc index bfed4f5230dfe37bca48560ce83a2dd82c8950a4..10bc82488ff56135f4585e62c2f71c11a359e542 100644 --- a/tensorflow/compiler/xla/service/gpu/partition_assignment.cc +++ b/tensorflow/compiler/xla/service/gpu/partition_assignment.cc @@ -41,7 +41,7 @@ std::ostream& operator<<(std::ostream& out, int64 ThreadsPerBlockLimit(const se::DeviceDescription& device_desc) { int64 threads_per_block = device_desc.threads_per_block_limit(); - if (threads_per_block == 0) { + if (threads_per_block <= 0) { static std::atomic log_count{0}; if (log_count.fetch_add(1) < 8) { LOG(WARNING) << "Attempting to calculate launch dimensions for GPU " @@ -71,18 +71,17 @@ LaunchDimensions CalculateLaunchDimensions( num_elements = num_elements / unroll_factor; // Since we don't do any inter-warp communication, we're free to choose any - // block size we want, subject to hardware constraints. We choose the - // smallest block size that allows the GPU to reach full occupancy (assuming - // the kernel uses sufficiently few registers). This gives us max performance - // when the kernel uses few registers, and lets us scale down gracefully as - // the kernel uses more registers. + // block size we want, subject to hardware constraints. We choose the largest + // block size allowed, as empirically, this is a performance win on almost + // (but not all) benchmarks. // - // Specifically, we choose the number of threads per block such that + // My guess is that using a larger block size encourages ptxas to decrease + // per-thread register usage, thus allowing for higher occupancy, but I + // haven't verified this. // - // * = - + // TODO(jlebar): Investigate this further, and tune this heuristic so we can + // run faster on the few benchmarks where smaller block size helps. int64 threads_per_block = ThreadsPerBlockLimit(device_desc); - if (num_elements < threads_per_block) { threads_per_block = num_elements; VLOG(2) << "Update # of threads per block to the element count (" diff --git a/tensorflow/compiler/xla/service/gpu/scratch_allocator.cc b/tensorflow/compiler/xla/service/gpu/scratch_allocator.cc new file mode 100644 index 0000000000000000000000000000000000000000..197367e81687eeddea8778267075e66ef1819341 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/scratch_allocator.cc @@ -0,0 +1,43 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/gpu/scratch_allocator.h" + +namespace xla { +namespace gpu { + +StatusOr> ScratchAllocator::AllocateBytes( + se::Stream* stream, int64 byte_size) { + CHECK_GE(byte_size, 0) << "byte_size must be positive."; + if (byte_size > GetMemoryLimitInBytes(stream)) { + return se::port::Status( + se::port::error::RESOURCE_EXHAUSTED, + absl::StrFormat( + "Allocating %d bytes exceeds the memory limit of %d bytes.", + byte_size, GetMemoryLimitInBytes(stream))); + } + + TF_ASSIGN_OR_RETURN(OwningDeviceMemory allocated_buffer, + memory_allocator_->Allocate(device_ordinal_, byte_size, + /*retry_on_failure=*/false)); + total_allocated_bytes_ += byte_size; + + se::DeviceMemoryBase buffer_addr = allocated_buffer.AsDeviceMemoryBase(); + allocated_buffers_.push_back(std::move(allocated_buffer)); + return se::DeviceMemory(buffer_addr); +} + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/scratch_allocator.h b/tensorflow/compiler/xla/service/gpu/scratch_allocator.h new file mode 100644 index 0000000000000000000000000000000000000000..620c7e78912eb7d9730bae02aab8f85b5fd2c096 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/scratch_allocator.h @@ -0,0 +1,61 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_SCRATCH_ALLOCATOR_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_SCRATCH_ALLOCATOR_H_ + +#include + +#include "tensorflow/compiler/xla/service/device_memory_allocator.h" +#include "tensorflow/compiler/xla/service/owning_device_memory.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace xla { +namespace gpu { + +class ScratchAllocator : public se::ScratchAllocator { + public: + ScratchAllocator(int device_ordinal, DeviceMemoryAllocator* memory_allocator) + : device_ordinal_(device_ordinal), memory_allocator_(memory_allocator) {} + + int64 GetMemoryLimitInBytes(se::Stream* stream) override { + return 1LL << 32; // 4GB. TODO(jlebar): Tune this? + } + int64 TotalAllocatedBytes() { return total_allocated_bytes_; } + + StatusOr> AllocateBytes(se::Stream* stream, + int64 byte_size) override; + + template + StatusOr> Allocate(se::Stream* stream, + int64 num_elements) { + TF_ASSIGN_OR_RETURN(se::DeviceMemory bytes, + AllocateBytes(stream, num_elements * sizeof(T))); + return se::DeviceMemory(bytes); + } + + private: + const int device_ordinal_; + DeviceMemoryAllocator* memory_allocator_; + std::vector allocated_buffers_; + int64 total_allocated_bytes_ = 0; +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_SCRATCH_ALLOCATOR_H_ diff --git a/tensorflow/compiler/xla/service/gpu/thunk.cc b/tensorflow/compiler/xla/service/gpu/thunk.cc index a677617727c04811584cbaa295d164ed27273bb2..6b98cbb6570057e6d78107cc887fe3a138babcef 100644 --- a/tensorflow/compiler/xla/service/gpu/thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/thunk.cc @@ -20,6 +20,8 @@ namespace gpu { std::ostream& operator<<(std::ostream& os, Thunk::Kind kind) { switch (kind) { + case Thunk::kCholesky: + return os << "kCholesky"; case Thunk::kConditional: return os << "kConditional"; case Thunk::kConvolution: @@ -32,6 +34,8 @@ std::ostream& operator<<(std::ostream& os, Thunk::Kind kind) { return os << "kCudnnBatchNormForwardInference"; case Thunk::kCudnnBatchNormForwardTraining: return os << "kCudnnBatchNormForwardTraining"; + case Thunk::kNcclAllReduce: + return os << "kNcclAllReduce"; case Thunk::kFft: return os << "kFft"; case Thunk::kGemm: diff --git a/tensorflow/compiler/xla/service/gpu/thunk.h b/tensorflow/compiler/xla/service/gpu/thunk.h index bc69af897a01775d2d33d46067464b10e049f3e1..442506f002c61545bc7da52c9ac14a5e7f6fc4f8 100644 --- a/tensorflow/compiler/xla/service/gpu/thunk.h +++ b/tensorflow/compiler/xla/service/gpu/thunk.h @@ -42,12 +42,14 @@ class GpuExecutable; class Thunk { public: enum Kind { + kCholesky, kConditional, kConvolution, kCopy, kCudnnBatchNormBackward, kCudnnBatchNormForwardInference, kCudnnBatchNormForwardTraining, + kNcclAllReduce, kFft, kGemm, kInfeed, diff --git a/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc b/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc index 2dce7749bbd8da2673ae607eee3d731d9917e8fe..77e49f0e46bbefb57f611bec26727189c92c6fa6 100644 --- a/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc +++ b/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc @@ -106,24 +106,6 @@ class WhileTransformerTest : public HloTestBase { return while_hlo; } - void RunFusionPasses() { - // Run standard fusion passes. - TF_ASSERT_OK(gpu::GpuInstructionFusion(/*may_duplicate=*/false) - .Run(module_.get()) - .status()); - TF_ASSERT_OK(gpu::GpuInstructionFusion(/*may_duplicate=*/true) - .Run(module_.get()) - .status()); - } - - void RunCopyInsertionPass() { - HloVerifier verifier(/*layout_sensitive=*/false, - /*allow_mixed_precision=*/false); - TF_ASSERT_OK(verifier.Run(module_.get()).status()); - CopyInsertion copy_insertion; - TF_ASSERT_OK(copy_insertion.Run(module_.get()).status()); - } - Shape GetLoopStateShape(const int64 ind_var_tuple_index) { if (ind_var_tuple_index == 0) { return ShapeUtil::MakeTupleShape( @@ -146,10 +128,6 @@ TEST_F(WhileTransformerTest, InductionVariableAtTupleElement0) { module_->AddEmbeddedComputation(BuildConditionComputation(0, 10)); auto body = module_->AddEmbeddedComputation(BuildBodyComputation(0, 1, 1)); auto while_hlo = BuildWhileInstruction(condition, body, 0, 0); - // Run HLO Optimization passes. - RunFusionPasses(); - RunCopyInsertionPass(); - auto result = ComputeWhileLoopTripCount(while_hlo); ASSERT_TRUE(result); EXPECT_EQ(10, *result); @@ -161,10 +139,6 @@ TEST_F(WhileTransformerTest, InductionVariableAtTupleElement1) { module_->AddEmbeddedComputation(BuildConditionComputation(1, 10)); auto body = module_->AddEmbeddedComputation(BuildBodyComputation(1, 0, 1)); auto while_hlo = BuildWhileInstruction(condition, body, 1, 0); - // Run HLO Optimization passes. - RunFusionPasses(); - RunCopyInsertionPass(); - auto result = ComputeWhileLoopTripCount(while_hlo); ASSERT_TRUE(result); EXPECT_EQ(10, *result); @@ -176,10 +150,6 @@ TEST_F(WhileTransformerTest, ImpossibleLoopLimit) { module_->AddEmbeddedComputation(BuildConditionComputation(0, 5)); auto body = module_->AddEmbeddedComputation(BuildBodyComputation(0, 1, 1)); auto while_hlo = BuildWhileInstruction(condition, body, 0, 10); - // Run HLO Optimization passes. - RunFusionPasses(); - RunCopyInsertionPass(); - auto result = ComputeWhileLoopTripCount(while_hlo); ASSERT_TRUE(result); EXPECT_EQ(0, *result); @@ -191,10 +161,6 @@ TEST_F(WhileTransformerTest, InvalidLoopIncrement) { module_->AddEmbeddedComputation(BuildConditionComputation(0, 10)); auto body = module_->AddEmbeddedComputation(BuildBodyComputation(0, 1, -1)); auto while_hlo = BuildWhileInstruction(condition, body, 0, 0); - // Run HLO Optimization passes. - RunFusionPasses(); - RunCopyInsertionPass(); - auto result = ComputeWhileLoopTripCount(while_hlo); ASSERT_FALSE(result); } diff --git a/tensorflow/compiler/xla/service/hlo.proto b/tensorflow/compiler/xla/service/hlo.proto index ae9e3169fd9b7a4655ab91ffb1589b845402ba8d..1413ce3062d135ca14adb6af43fceee0ad79b789 100644 --- a/tensorflow/compiler/xla/service/hlo.proto +++ b/tensorflow/compiler/xla/service/hlo.proto @@ -29,12 +29,13 @@ limitations under the License. syntax = "proto3"; package xla; + import "tensorflow/compiler/xla/xla_data.proto"; option cc_enable_arenas = true; // Serialization of HloInstruction. -// Next ID: 62 +// Next ID: 63 message HloInstructionProto { reserved 10; reserved "parameter_name"; @@ -200,6 +201,9 @@ message HloInstructionProto { // Options for TriangularSolve xla.TriangularSolveOptions triangular_solve_options = 59; + // Options for Cholesky + xla.CholeskyOptions cholesky_options = 62; + // Describes how parameters behave with regards to replicas. xla.ParameterReplication parameter_replication = 61; } diff --git a/tensorflow/compiler/xla/service/hlo_alias_analysis.cc b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc index e511f1951c5dd07ebb64fa38fd5b7f6a0e87b429..7d02f4b3d756df9d1fcbddfa85df2a41a62d9169 100644 --- a/tensorflow/compiler/xla/service/hlo_alias_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc @@ -293,7 +293,7 @@ class BufferValueMap { VLOG(3) << " value @ " << position << " is root of " << callsite.instruction()->name() - << "; true/false branch roots must share buffer among them : " + << "; branch computation roots must share buffer among them : " << cond_value.ToShortString(); aliased_buffers->push_back(GetBufferForValue(cond_value)); } diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc index 817e15f9ff10a9b7e1a502265c85f70fdd681dd9..48a51d302bbf054d904c54ab933d87fc910d0714 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.cc +++ b/tensorflow/compiler/xla/service/hlo_computation.cc @@ -124,6 +124,24 @@ HloInstruction* HloComputation::AddParameter( return instructions_.back().get(); } +HloInstruction* HloComputation::AddEntryComputationParameter( + std::unique_ptr instruction) { + CHECK_EQ(instruction->opcode(), HloOpcode::kParameter); + CHECK_EQ(instruction->parameter_number(), num_parameters()); + CHECK(parent()->entry_computation() == this); + + HloModuleConfig config = parent()->config(); + config.mutable_entry_computation_layout()->add_parameter_layout( + ShapeLayout(instruction->shape())); + parent()->set_config(config); + + instruction->set_parent(this); + param_instructions_.push_back(instruction.get()); + AddInstructionInternal(std::move(instruction)); + + return instructions_.back().get(); +} + Status HloComputation::RemoveParameter(int64 param_no) { CHECK_GE(param_no, 0); CHECK_LT(param_no, param_instructions_.size()); diff --git a/tensorflow/compiler/xla/service/hlo_computation.h b/tensorflow/compiler/xla/service/hlo_computation.h index 212dfa15a13185f1050103739fad8b560270d401..a48cfa1f1b22ffd748fe9fe3ddb7f36d8d0dee4d 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.h +++ b/tensorflow/compiler/xla/service/hlo_computation.h @@ -117,11 +117,20 @@ class HloComputation { // instruction. Status RemoveUnusedParameters(); - // Add new parameter instruction to the computation. + // Adds a new parameter instruction to a fusion computation. + // // This should be a new parameter. Instruction will be appended to parameters // and inserted to the instruction list. HloInstruction* AddParameter(std::unique_ptr instruction); + // Adds a new parameter instruction to the entry computation and update + // the parent module config to reflect the change. + // + // This should be a new parameter. Instruction will be appended to parameters + // and inserted to the instruction list. + HloInstruction* AddEntryComputationParameter( + std::unique_ptr instruction); + // Remove an instruction from the computation. The instruction must have no // users. Instruction is deallocated with this call. Status RemoveInstruction(HloInstruction* instruction); diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc index 29ac263c5f39b5ec1f02a232704adcd3e3f21f60..13b1c82709523fc98b02551d14bc9a9cdacc5fc1 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc @@ -91,9 +91,10 @@ Status HloCostAnalysis::HandleElementwiseOp( auto opcode = hlo_instruction->opcode(); // We treat transcendental operations separately since one transcendental // operation can correspond to several floating point ops. - if (opcode == HloOpcode::kExp || opcode == HloOpcode::kPower || - opcode == HloOpcode::kTanh || opcode == HloOpcode::kSin || - opcode == HloOpcode::kCos) { + if (opcode == HloOpcode::kExp || opcode == HloOpcode::kLog || + opcode == HloOpcode::kPower || opcode == HloOpcode::kSqrt || + opcode == HloOpcode::kRsqrt || opcode == HloOpcode::kTanh || + opcode == HloOpcode::kSin || opcode == HloOpcode::kCos) { current_properties_[kTranscendentalsKey] = computation_count; } else { // Note: transcendental operations are considered a separate category from @@ -556,11 +557,22 @@ Status HloCostAnalysis::HandleTriangularSolve(const HloInstruction* hlo) { // Estimate as batch * mn^2 / 2 flops. int64 elems = a_shape.dimensions(a_shape.dimensions_size() - 1); elems *= ShapeUtil::ElementsIn(b_shape); - // Each output elment requires reduction_widht FMA operations. current_properties_[kFlopsKey] = kFmaFlops * elems; return Status::OK(); } +Status HloCostAnalysis::HandleCholesky(const HloInstruction* hlo) { + float bytes_accessed = GetShapeSize(hlo->operand(0)->shape()) / 2.0f; + current_properties_[kBytesAccessedKey] = bytes_accessed; + + const Shape& a_shape = hlo->operand(0)->shape(); + // Estimate as batch * n^3 / 3 flops. + int64 elems = a_shape.dimensions(a_shape.dimensions_size() - 1); + elems *= ShapeUtil::ElementsIn(a_shape); + current_properties_[kFlopsKey] = elems / 3; + return Status::OK(); +} + Status HloCostAnalysis::HandleAllReduce(const HloInstruction* crs) { // We assume 2 replicas, so that each output element is the sum of two input // elements. @@ -672,19 +684,22 @@ Status HloCostAnalysis::HandleWhile(const HloInstruction* xla_while) { } Status HloCostAnalysis::HandleConditional(const HloInstruction* conditional) { - // Compute the cost of the true and false computations and take the maximum - // from those for each property. + // Compute the cost of the branch computations and take the maximum from those + // for each property. TF_ASSIGN_OR_RETURN( - const Properties true_computation_properties, - ProcessUnnestedSubcomputation(conditional->true_computation())); - TF_ASSIGN_OR_RETURN( - const Properties false_computation_properties, - ProcessUnnestedSubcomputation(conditional->false_computation())); - current_properties_ = true_computation_properties; - for (const auto& property : false_computation_properties) { - if (!tensorflow::gtl::InsertIfNotPresent(¤t_properties_, property)) { - current_properties_[property.first] = - std::max(current_properties_[property.first], property.second); + const Properties branch0_computation_properties, + ProcessUnnestedSubcomputation(conditional->branch_computation(0))); + current_properties_ = branch0_computation_properties; + for (int j = 1; j < conditional->branch_count(); ++j) { + TF_ASSIGN_OR_RETURN( + const Properties branch_computation_properties, + ProcessUnnestedSubcomputation(conditional->branch_computation(j))); + for (const auto& property : branch_computation_properties) { + if (!tensorflow::gtl::InsertIfNotPresent(¤t_properties_, + property)) { + auto& current_property = current_properties_[property.first]; + current_property = std::max(current_property, property.second); + } } } current_should_compute_bottleneck_time_ = false; diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.h b/tensorflow/compiler/xla/service/hlo_cost_analysis.h index 96357dec68e390251c43c2c3fc6f5a5612063fbd..4480554de507f20b5d44b87a19e58236252bad1d 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.h +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.h @@ -72,6 +72,7 @@ class HloCostAnalysis : public ConstDfsHloVisitor { Status HandleConvolution(const HloInstruction* convolution) override; Status HandleFft(const HloInstruction* fft) override; Status HandleTriangularSolve(const HloInstruction* hlo) override; + Status HandleCholesky(const HloInstruction* hlo) override; Status HandleAllReduce(const HloInstruction* crs) override; Status HandleAllToAll(const HloInstruction* hlo) override; Status HandleCollectivePermute(const HloInstruction* hlo) override; diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc index 3144a84805454488f417391f40ed6b9e9facc752..ce55db1b66d1a5ea85d3580c146c647fc2999fb9 100644 --- a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc @@ -414,11 +414,11 @@ bool HloDataflowAnalysis::UpdateCallValueSet(HloInstruction* call) { bool HloDataflowAnalysis::UpdateConditionalValueSet( HloInstruction* conditional) { CHECK_EQ(conditional->opcode(), HloOpcode::kConditional); - const InstructionValueSet* const inputs[] = { - &GetInstructionValueSet( - conditional->true_computation()->root_instruction()), - &GetInstructionValueSet( - conditional->false_computation()->root_instruction())}; + std::vector inputs(conditional->branch_count()); + for (int j = 0; j < conditional->branch_count(); ++j) { + inputs[j] = &GetInstructionValueSet( + conditional->branch_computation(j)->root_instruction()); + } if (ssa_form_) { return Phi(conditional, inputs); } else { @@ -546,20 +546,23 @@ bool HloDataflowAnalysis::UpdateParameterValueSet(HloInstruction* parameter) { } else if (callsite.instruction()->opcode() == HloOpcode::kConditional) { CHECK_EQ(parameter->parameter_number(), 0); auto conditional = callsite.instruction(); - // Conditional has 3 operands. Operand 0 is the predicate, operand 1 is - // the argument to the true computation and operand 2 is the argument to - // the false computation. + // Conditional has branch_count+1 operands. Operand 0 is the branch_index, + // operands 1 and onward are the arguments to the branch computations. // - // If the parameter belongs to conditional's true computation, then + // If the parameter belongs to conditional's branch 0 computation, then // operand 1 is forwarded to this parameter instruction. If the parameter - // belongs to conditional's false computation, then operand 2 is forwarded - // to this parameter instruction. - if (parameter->parent() == conditional->true_computation()) { - inputs.push_back(&GetInstructionValueSet(conditional->operand(1))); - } else { - CHECK_EQ(parameter->parent(), conditional->false_computation()); - inputs.push_back(&GetInstructionValueSet(conditional->operand(2))); + // belongs to conditional's branch 5 computation, then operand 6 is + // forwarded to this parameter instruction. + bool found_parent = false; + for (int j = 0; j < conditional->branch_count(); ++j) { + if (parameter->parent() == conditional->branch_computation(j)) { + inputs.push_back( + &GetInstructionValueSet(conditional->operand(j + 1))); + found_parent = true; + break; + } } + CHECK(found_parent); need_phi = true; } else { LOG(FATAL) << "CallContext::kSequential computations should only be " @@ -710,19 +713,17 @@ void HloDataflowAnalysis::Propagate() { // parameter(s) of the computation need to be updated. if (user->opcode() == HloOpcode::kConditional) { // If operand 0 is the use of instruction, then no parameters need to be - // updated, since that is the predicate of the conditional. - // If operand 1 is the use of instruction, then the true_computation's - // parameter need to be updated. - // If operand 2 is the use of instruction, then the false_computation's - // parameter need to be updated. + // updated, since that is the branch_index of the conditional. + // If operand n+1 is the use of instruction, then the branch_computation + // n's parameter need to be updated. // - // Note that the same instruction can be used in both operand 1 and - // operand 2. - if (user->operand(1) == instruction) { - add_to_worklist(user->true_computation()->parameter_instruction(0)); - } - if (user->operand(2) == instruction) { - add_to_worklist(user->false_computation()->parameter_instruction(0)); + // Note that the same instruction can be used in multiple branches' + // operands. + for (int j = 0; j < user->branch_count(); ++j) { + if (user->operand(j + 1) == instruction) { + add_to_worklist( + user->branch_computation(j)->parameter_instruction(0)); + } } } else { for (HloComputation* called_computation : user->called_computations()) { @@ -744,8 +745,8 @@ void HloDataflowAnalysis::Propagate() { const CallGraphNode& call_graph_node = call_graph_->GetNode(instruction->parent()); for (const CallSite& callsite : call_graph_node.caller_callsites()) { - if ((callsite.instruction()->opcode() == HloOpcode::kCall) || - (callsite.instruction()->opcode() == HloOpcode::kConditional)) { + if (callsite.instruction()->opcode() == HloOpcode::kCall || + callsite.instruction()->opcode() == HloOpcode::kConditional) { add_to_worklist(callsite.instruction()); } else if (callsite.instruction()->opcode() == HloOpcode::kWhile) { // Add the while itself, and the body and condition parameters. diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.cc b/tensorflow/compiler/xla/service/hlo_evaluator.cc index 4d6487700b24cfd3b89aece58e5ad6d7bb43a800..63bdbc52e82a123e8b8b8beeaa0a76546e076be4 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator.cc +++ b/tensorflow/compiler/xla/service/hlo_evaluator.cc @@ -1270,28 +1270,27 @@ Status HloEvaluator::HandleFusion(HloInstruction* fusion) { } Status HloEvaluator::HandleConditional(HloInstruction* conditional) { - const auto& pred = GetEvaluatedLiteralFor(conditional->operand(0)); - const auto& true_computation_arg = - GetEvaluatedLiteralFor(conditional->operand(1)); - const auto& false_computation_arg = - GetEvaluatedLiteralFor(conditional->operand(2)); - - auto* true_computation = conditional->true_computation(); - auto* false_computation = conditional->false_computation(); + const auto& branch_index_literal = + GetEvaluatedLiteralFor(conditional->operand(0)); + int branch_index; + if (conditional->operand(0)->shape().element_type() == PRED) { + branch_index = branch_index_literal.Get({}) ? 0 : 1; + } else { + branch_index = branch_index_literal.Get({}); + if (branch_index < 0 || branch_index >= conditional->branch_count()) { + branch_index = conditional->branch_count() - 1; + } + } + const auto& branch_computation_arg = + GetEvaluatedLiteralFor(conditional->operand(1 + branch_index)); HloEvaluator embedded_evaluator; embedded_evaluator.set_dynamic_dimension_inference( dynamic_dimension_inference_); - Literal result; - if (pred.Get({})) { - result = - embedded_evaluator.Evaluate(*true_computation, {&true_computation_arg}) - .ConsumeValueOrDie(); - } else { - result = embedded_evaluator - .Evaluate(*false_computation, {&false_computation_arg}) - .ConsumeValueOrDie(); - } + Literal result = embedded_evaluator + .Evaluate(*conditional->branch_computation(branch_index), + {&branch_computation_arg}) + .ConsumeValueOrDie(); evaluated_[conditional] = std::move(result); return Status::OK(); diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h index d516a6258c80bda168ef4c6fd976e60946eb8b5b..2d8a578985e8f603d4056bee8619725095ebc7bb 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h +++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h @@ -329,10 +329,10 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { template < typename NativeT, typename std::enable_if::value>::type* = nullptr> - Status HandleLog1p(HloInstruction* expm1) { + Status HandleLog1p(HloInstruction* log1p) { TF_ASSIGN_OR_RETURN( - parent_->evaluated_[expm1], - ElementWiseUnaryOp(expm1, [](ElementwiseT elem_operand) { + parent_->evaluated_[log1p], + ElementWiseUnaryOp(log1p, [](ElementwiseT elem_operand) { return std::log1p(elem_operand); })); return Status::OK(); @@ -664,6 +664,23 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { return Status::OK(); } + Status HandleSqrt(HloInstruction* sqrt) override { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[sqrt], + ElementWiseUnaryOp(sqrt, [](ElementwiseT elem_operand) { + return std::sqrt(elem_operand); + })); + return Status::OK(); + } + + Status HandleRsqrt(HloInstruction* rsqrt) override { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[rsqrt], + ElementWiseUnaryOp(rsqrt, [](ElementwiseT elem_operand) { + return static_cast(1) / std::sqrt(elem_operand); + })); + return Status::OK(); + } + template ::value>::type* = nullptr> Status HandleRemainder(HloInstruction* remainder) { diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc index 49300b3ffe2f755d103af7877ab3fee5298eeb3e..9623edcf5eba706e8e193484c856b80918dd5eb3 100644 --- a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc +++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc @@ -829,8 +829,7 @@ string HloDotDumper::GetInstructionNodeInlinedOperands( // collected from profiling tools. Those constants may not have a valid // literal. if (elem_count.has_value() && *elem_count <= 8 && constant->HasLiteral()) { - return StrFormat("%s (%s)", constant->literal().ToString(), - ShapeUtil::HumanString(constant->shape())); + return constant->literal().ToString(); } // Otherwise, print e.g. "%constant.42 (s32[100])". @@ -953,6 +952,7 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { case HloOpcode::kRemainder: case HloOpcode::kRng: case HloOpcode::kRoundNearestAfz: + case HloOpcode::kRsqrt: case HloOpcode::kSelect: case HloOpcode::kShiftLeft: case HloOpcode::kShiftRightArithmetic: @@ -961,6 +961,7 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { case HloOpcode::kSin: case HloOpcode::kSlice: case HloOpcode::kSort: + case HloOpcode::kSqrt: case HloOpcode::kSubtract: case HloOpcode::kTanh: // De-emphasize scalar-shaped elementwise ops -- they're generally @@ -1016,6 +1017,7 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { case HloOpcode::kDot: case HloOpcode::kFft: case HloOpcode::kTriangularSolve: + case HloOpcode::kCholesky: return kDarkBlue; case HloOpcode::kReducePrecision: return kRed; diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc index 33c2270eb0a847d088776a2d9d67e341a69dbae2..bb45eb4fa0d55aabc7ad449f07d62c68a804f8f8 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.cc +++ b/tensorflow/compiler/xla/service/hlo_instruction.cc @@ -82,6 +82,15 @@ StatusOr> HloInstruction::CreateFromProto( const auto computations = [&computation_map, &proto](int index) { return computation_map.at(proto.called_computation_ids(index)); }; + const auto all_computations = [&computation_map, &proto]() { + std::vector result(proto.called_computation_ids_size()); + std::transform(proto.called_computation_ids().begin(), + proto.called_computation_ids().end(), result.begin(), + [&computation_map](int64 computation_id) { + return computation_map.at(computation_id); + }); + return result; + }; TF_RET_CHECK( absl::c_all_of(proto.operand_ids(), @@ -96,36 +105,31 @@ StatusOr> HloInstruction::CreateFromProto( Shape shape(proto.shape()); TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(shape)); + absl::optional arity = HloOpcodeArity(opcode); + if (arity) { + TF_RET_CHECK(proto.operand_ids_size() == *arity) + << proto.opcode() << " instruction should have " << *arity + << " operands but sees " << proto.operand_ids_size(); + } + switch (opcode) { // Ops migrated to subclasses. case HloOpcode::kBatchNormTraining: - TF_RET_CHECK(proto.operand_ids_size() == 3) - << "BatchNormTraining instruction should have 3 operands but sees " - << proto.operand_ids_size(); instruction = CreateBatchNormTraining(shape, operands(0), operands(1), operands(2), proto.epsilon(), proto.feature_index()); break; case HloOpcode::kBatchNormInference: - TF_RET_CHECK(proto.operand_ids_size() == 5) - << "BatchNormInference instruction should have 5 operands but sees " - << proto.operand_ids_size(); instruction = CreateBatchNormInference( shape, operands(0), operands(1), operands(2), operands(3), operands(4), proto.epsilon(), proto.feature_index()); break; case HloOpcode::kBatchNormGrad: - TF_RET_CHECK(proto.operand_ids_size() == 5) - << "BatchNormGrad instruction should have 5 operands but sees " - << proto.operand_ids_size(); instruction = CreateBatchNormGrad(shape, operands(0), operands(1), operands(2), operands(3), operands(4), proto.epsilon(), proto.feature_index()); break; case HloOpcode::kFft: { - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "Fft instruction should have 1 operand but sees " - << proto.operand_ids_size(); std::vector fft_length(proto.fft_length().begin(), proto.fft_length().end()); instruction = CreateFft(shape, operands(0), proto.fft_type(), @@ -133,43 +137,30 @@ StatusOr> HloInstruction::CreateFromProto( break; } case HloOpcode::kTriangularSolve: { - TF_RET_CHECK(proto.operand_ids_size() == 2) - << "Triangular solve instruction should have 2 operands but sees " - << proto.operand_ids_size(); instruction = CreateTriangularSolve(shape, operands(0), operands(1), proto.triangular_solve_options()); break; } + case HloOpcode::kCholesky: { + instruction = + CreateCholesky(shape, operands(0), proto.cholesky_options()); + break; + } case HloOpcode::kSend: - TF_RET_CHECK(proto.operand_ids_size() == 2) - << "Send instruction should have 2 operand but sees " - << proto.operand_ids_size(); instruction = CreateSend(operands(0), operands(1), proto.channel_id(), proto.is_host_transfer()); break; case HloOpcode::kSendDone: - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "SendDone instruction should have 1 operand but sees " - << proto.operand_ids_size(); instruction = CreateSendDone(operands(0), proto.is_host_transfer()); break; case HloOpcode::kRecv: - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "Recv instruction should have 1 operand but sees " - << proto.operand_ids_size(); instruction = CreateRecv(shape.tuple_shapes(0), operands(0), proto.channel_id(), proto.is_host_transfer()); break; case HloOpcode::kRecvDone: - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "RecvDone instruction should have 1 operand but sees " - << proto.operand_ids_size(); instruction = CreateRecvDone(operands(0), proto.is_host_transfer()); break; case HloOpcode::kReverse: - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "Reverse instruction should have 1 operand but sees " - << proto.operand_ids_size(); instruction = CreateReverse(shape, operands(0), std::vector(proto.dimensions().begin(), proto.dimensions().end())); @@ -181,6 +172,26 @@ StatusOr> HloInstruction::CreateFromProto( instruction = CreateConcatenate(shape, all_operands(), proto.dimensions(0)); break; + case HloOpcode::kConditional: { + TF_RET_CHECK(proto.called_computation_ids_size() > 0) + << "conditional should have at least 1 called computation"; + if (operands(0)->shape().element_type() == PRED) { + TF_RET_CHECK(proto.called_computation_ids_size() == 2) + << "conditional should have exactly 2 called computations but got " + << proto.called_computation_ids_size(); + } + TF_RET_CHECK(proto.operand_ids_size() == + proto.called_computation_ids_size() + 1) + << "conditional should have one branch_index operand plus one " + "operand per called computation but got " + << proto.operand_ids_size() << " operands for " + << proto.called_computation_ids_size() << " branch computations"; + auto cond_operands = all_operands(); + instruction = + CreateConditional(shape, cond_operands[0], all_computations(), + absl::MakeSpan(cond_operands).subspan(1)); + break; + } case HloOpcode::kReduce: TF_RET_CHECK(proto.operand_ids_size() % 2 == 0) << "Reduce instruction should have an even number of operands but " @@ -218,18 +229,12 @@ StatusOr> HloInstruction::CreateFromProto( break; } case HloOpcode::kTranspose: - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "Transpose instruction should have 1 operand but sees " - << proto.operand_ids_size(); instruction = CreateTranspose(shape, operands(0), std::vector(proto.dimensions().begin(), proto.dimensions().end())); break; case HloOpcode::kBroadcast: - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "Broadcast instruction should have 1 operand but sees " - << proto.operand_ids_size(); instruction = CreateBroadcast(shape, operands(0), std::vector(proto.dimensions().begin(), @@ -242,9 +247,6 @@ StatusOr> HloInstruction::CreateFromProto( instruction = CreateMap(shape, all_operands(), computations(0)); break; case HloOpcode::kSlice: { - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "Slice instruction should have 1 operand but sees " - << proto.operand_ids_size(); std::vector slice_starts, slice_limits, slice_strides; for (const HloInstructionProto::SliceDimensions& slice_dimensions : proto.slice_dimensions()) { @@ -268,9 +270,6 @@ StatusOr> HloInstruction::CreateFromProto( break; } case HloOpcode::kTrace: { - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "Trace instruction should have 1 operand but sees " - << proto.operand_ids_size(); TF_RET_CHECK(proto.has_literal()); TF_ASSIGN_OR_RETURN(auto literal, Literal::CreateFromProto(proto.literal())); @@ -310,16 +309,10 @@ StatusOr> HloInstruction::CreateFromProto( } break; case HloOpcode::kGetTupleElement: - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "GetTupleElement instruction should have 1 operand but sees " - << proto.operand_ids_size(); instruction = CreateGetTupleElement(shape, operands(0), proto.tuple_index()); break; case HloOpcode::kReducePrecision: - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "ReducePrecision instruction should have 1 operand but sees " - << proto.operand_ids_size(); instruction = CreateReducePrecision( shape, operands(0), proto.exponent_bits(), proto.mantissa_bits()); break; @@ -329,16 +322,10 @@ StatusOr> HloInstruction::CreateFromProto( << "Infeed should have a tuple shape with 2 operands, but has: " << shape; const Shape& data_shape = ShapeUtil::GetTupleElementShape(shape, 0); - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "Infeed instruction should have 1 operand but sees " - << proto.operand_ids_size(); instruction = CreateInfeed(data_shape, operands(0), proto.infeed_config()); } break; case HloOpcode::kOutfeed: { - TF_RET_CHECK(proto.operand_ids_size() == 2) - << "Outfeed instruction should have 2 operands but sees " - << proto.operand_ids_size(); Shape outfeed_shape(proto.outfeed_shape()); TF_RETURN_IF_ERROR( ShapeUtil::ValidateShapeWithOptionalLayout(outfeed_shape)); @@ -372,9 +359,6 @@ StatusOr> HloInstruction::CreateFromProto( break; } case HloOpcode::kCollectivePermute: { - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "CollectivePermute instruction should have 1 operand but sees " - << proto.operand_ids_size(); std::vector> source_target_pairs( proto.source_target_pairs_size()); for (int i = 0; i < source_target_pairs.size(); i++) { @@ -386,16 +370,10 @@ StatusOr> HloInstruction::CreateFromProto( break; } case HloOpcode::kReplicaId: { - TF_RET_CHECK(proto.operand_ids_size() == 0) - << "ReplicaId instruction should have 0 operand but sees " - << proto.operand_ids_size(); instruction = CreateReplicaId(); break; } case HloOpcode::kConvolution: { - TF_RET_CHECK(proto.operand_ids_size() == 2) - << "Convolution instruction should have 2 operands but sees " - << proto.operand_ids_size(); TF_RET_CHECK(proto.has_window()); TF_RET_CHECK(proto.has_convolution_dimension_numbers()); PrecisionConfig precision_config = proto.precision_config(); @@ -409,9 +387,6 @@ StatusOr> HloInstruction::CreateFromProto( break; } case HloOpcode::kReduceWindow: - TF_RET_CHECK(proto.operand_ids_size() == 2) - << "ReduceWindow instruction should have 2 operands but sees " - << proto.operand_ids_size(); TF_RET_CHECK(proto.called_computation_ids_size() == 1) << "ReduceWindow should have 1 called computation but sees " << proto.called_computation_ids_size(); @@ -419,9 +394,6 @@ StatusOr> HloInstruction::CreateFromProto( proto.window(), computations(0)); break; case HloOpcode::kSelectAndScatter: - TF_RET_CHECK(proto.operand_ids_size() == 3) - << "SelectAndScatter instruction should have 3 operands but sees " - << proto.operand_ids_size(); TF_RET_CHECK(proto.called_computation_ids_size() == 2) << "SelectAndScatter should have 2 called computations but sees " << proto.called_computation_ids_size(); @@ -464,9 +436,6 @@ StatusOr> HloInstruction::CreateFromProto( std::max(static_cast(proto.batch_group_count()), 1LL)); break; case HloOpcode::kPad: - TF_RET_CHECK(proto.operand_ids_size() == 2) - << "Pad instruction should have 2 operands but sees " - << proto.operand_ids_size(); TF_RET_CHECK(proto.has_padding_config()); instruction = CreatePad(shape, operands(0), operands(1), proto.padding_config()); @@ -512,9 +481,6 @@ StatusOr> HloInstruction::CreateFromProto( break; } case HloOpcode::kGather: { - TF_RET_CHECK(proto.operand_ids_size() == 2) - << "Gather instruction should have 2 operands but sees " - << proto.operand_ids_size(); TF_RET_CHECK(proto.has_gather_dimension_numbers()) << "Gather instruction should have GatherDimensionNumbers set."; std::unique_ptr gather_dimension_numbers = @@ -529,9 +495,6 @@ StatusOr> HloInstruction::CreateFromProto( break; } case HloOpcode::kScatter: { - TF_RET_CHECK(proto.operand_ids_size() == 3) - << "Scatter instruction should have 3 operands but sees " - << proto.operand_ids_size(); TF_RET_CHECK(proto.has_scatter_dimension_numbers()) << "Scatter instruction should have ScatterDimensionNumbers set."; TF_RET_CHECK(proto.called_computation_ids_size() == 1) @@ -553,9 +516,6 @@ StatusOr> HloInstruction::CreateFromProto( case HloOpcode::kDot: { TF_RET_CHECK(proto.has_dot_dimension_numbers()) << "Dot instruction should have dot_dimension_numbers."; - TF_RET_CHECK(proto.operand_ids_size() == 2) - << "Dot instruction should have 2 operands but sees " - << proto.operand_ids_size(); PrecisionConfig precision_config = proto.precision_config(); precision_config.mutable_operand_precision()->Resize( proto.operand_ids_size(), PrecisionConfig::DEFAULT); @@ -565,9 +525,6 @@ StatusOr> HloInstruction::CreateFromProto( break; } case HloOpcode::kDomain: { - TF_RET_CHECK(proto.operand_ids_size() == 1) - << "Domain instruction should have 1 operands but sees " - << proto.operand_ids_size(); std::shared_ptr entry_hlo_sharding; std::shared_ptr exit_hlo_sharding; if (proto.has_domain_entry_sharding()) { @@ -589,7 +546,6 @@ StatusOr> HloInstruction::CreateFromProto( break; } case HloOpcode::kGetDimensionSize: - TF_RET_CHECK(proto.operand_ids_size() == 1); TF_RET_CHECK(proto.dimensions_size() == 1); instruction = CreateGetDimensionSize(shape, operands(0), proto.dimensions(0)); @@ -709,8 +665,10 @@ HloInstruction::CreateGetTupleElement(const Shape& shape, case HloOpcode::kNot: case HloOpcode::kNegate: case HloOpcode::kReal: + case HloOpcode::kRsqrt: case HloOpcode::kSign: case HloOpcode::kSin: + case HloOpcode::kSqrt: case HloOpcode::kTanh: break; default: @@ -810,6 +768,11 @@ HloInstruction::CreateTriangularSolve(const Shape& shape, HloInstruction* a, return absl::make_unique(shape, a, b, options); } +/* static */ std::unique_ptr HloInstruction::CreateCholesky( + const Shape& shape, HloInstruction* a, const CholeskyOptions& options) { + return absl::make_unique(shape, a, options); +} + /* static */ std::unique_ptr HloInstruction::CreateDot( const Shape& shape, HloInstruction* lhs, HloInstruction* rhs, const DotDimensionNumbers& dimension_numbers, @@ -964,6 +927,21 @@ HloInstruction::CreateAddDependency(HloInstruction* data_operand, return instruction; } +/* static */ std::unique_ptr HloInstruction::CreateConditional( + const Shape& shape, HloInstruction* branch_index, + absl::Span branch_computations, + absl::Span branch_computation_args) { + auto instruction = + absl::WrapUnique(new HloInstruction(HloOpcode::kConditional, shape)); + instruction->AppendOperand(branch_index); + CHECK_EQ(branch_computations.size(), branch_computation_args.size()); + for (int i = 0; i < branch_computations.size(); ++i) { + instruction->called_computations_.push_back(branch_computations[i]); + instruction->AppendOperand(branch_computation_args[i]); + } + return instruction; +} + /* static */ std::unique_ptr HloInstruction::CreateSlice( const Shape& shape, HloInstruction* operand, absl::Span start_indices, @@ -1370,6 +1348,7 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( case HloOpcode::kDomain: case HloOpcode::kGetDimensionSize: case HloOpcode::kTriangularSolve: + case HloOpcode::kCholesky: clone = CloneWithNewOperandsImpl(shape, new_operands, context); break; // Unary ops. @@ -1390,8 +1369,10 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( case HloOpcode::kNot: case HloOpcode::kNegate: case HloOpcode::kReal: + case HloOpcode::kRsqrt: case HloOpcode::kSign: case HloOpcode::kSin: + case HloOpcode::kSqrt: case HloOpcode::kTanh: CHECK_EQ(new_operands.size(), 1); clone = CreateUnary(shape, opcode_, new_operands[0]); @@ -1460,10 +1441,10 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( CreateWhile(shape, while_condition(), while_body(), new_operands[0]); break; case HloOpcode::kConditional: - CHECK_EQ(new_operands.size(), 3); - clone = CreateConditional(shape, new_operands[0], new_operands[1], - true_computation(), new_operands[2], - false_computation()); + CHECK_EQ(new_operands.size(), branch_count() + 1); + clone = CreateConditional(shape, new_operands[0], + absl::MakeSpan(branch_computations()), + new_operands.subspan(1)); break; case HloOpcode::kAfterAll: if (new_operands.empty()) { @@ -1751,12 +1732,14 @@ bool HloInstruction::IdenticalSlowPath( case HloOpcode::kReshape: case HloOpcode::kReplicaId: case HloOpcode::kRoundNearestAfz: + case HloOpcode::kRsqrt: case HloOpcode::kSelect: case HloOpcode::kShiftLeft: case HloOpcode::kShiftRightArithmetic: case HloOpcode::kShiftRightLogical: case HloOpcode::kSign: case HloOpcode::kSin: + case HloOpcode::kSqrt: case HloOpcode::kSubtract: case HloOpcode::kTanh: case HloOpcode::kTuple: @@ -1772,16 +1755,16 @@ bool HloInstruction::IdenticalSlowPath( case HloOpcode::kCall: return eq_computations(to_apply(), other.to_apply()); case HloOpcode::kConditional: - return eq_computations(true_computation(), other.true_computation()) && - eq_computations(false_computation(), other.false_computation()); - - case HloOpcode::kWhile: { - if (eq_computations(while_body(), other.while_body()) && - eq_computations(while_condition(), other.while_condition())) { - return true; + for (int j = 0; j < branch_count(); ++j) { + if (!eq_computations(branch_computation(j), + other.branch_computation(j))) { + return false; + } } - return false; - } + return true; + case HloOpcode::kWhile: + return (eq_computations(while_body(), other.while_body()) && + eq_computations(while_condition(), other.while_condition())); // Ops migrated to subclasses should never come to this line. // TODO(b/80131774): Remove this switch when migration is complete. @@ -1826,6 +1809,7 @@ bool HloInstruction::IdenticalSlowPath( case HloOpcode::kDomain: case HloOpcode::kGetDimensionSize: case HloOpcode::kTriangularSolve: + case HloOpcode::kCholesky: LOG(FATAL) << "Base class impl called for opcode with subclass: " << opcode(); } @@ -1880,7 +1864,11 @@ Status HloInstruction::ReplaceUseWith(HloInstruction* user, << "this shape: " << ShapeUtil::HumanString(shape()) << ", replacement shape: " << ShapeUtil::HumanString(new_producer->shape()); + return ReplaceUseWithDifferentShape(user, new_producer); +} +Status HloInstruction::ReplaceUseWithDifferentShape( + HloInstruction* user, HloInstruction* new_producer) { VLOG(3) << "Replacing uses of " << name() << " in " << user->name() << " with " << new_producer->name(); @@ -2039,28 +2027,41 @@ HloInstruction* HloInstruction::while_init() const { HloComputation* HloInstruction::true_computation() const { CHECK_EQ(HloOpcode::kConditional, opcode_); + CHECK_EQ(PRED, operand(0)->shape().element_type()); return called_computations_[kTrueComputationIndex]; } HloComputation* HloInstruction::false_computation() const { CHECK_EQ(HloOpcode::kConditional, opcode_); + CHECK_EQ(PRED, operand(0)->shape().element_type()); return called_computations_[kFalseComputationIndex]; } -void HloInstruction::set_true_computation(HloComputation* true_computation) { - // Don't allow changing the computation for fused instructions so we don't - // have to recompute called_instructions for the entire fusion instruction. - CHECK(!IsFused()); - CHECK_EQ(HloOpcode::kConditional, opcode_); - called_computations_[kTrueComputationIndex] = true_computation; +const std::vector& HloInstruction::branch_computations() + const { + CHECK(HloOpcode::kConditional == opcode_); + return called_computations_; +} + +int HloInstruction::branch_count() const { + CHECK(HloOpcode::kConditional == opcode_); + return called_computations_.size(); +} + +HloComputation* HloInstruction::branch_computation(int b) const { + CHECK(HloOpcode::kConditional == opcode_); + CHECK_GE(b, 0); + CHECK_LT(b, called_computations_.size()); + return called_computations_[b]; } -void HloInstruction::set_false_computation(HloComputation* false_computation) { +void HloInstruction::set_branch_computation(int b, + HloComputation* computation) { // Don't allow changing the computation for fused instructions so we don't // have to recompute called_instructions for the entire fusion instruction. CHECK(!IsFused()); CHECK_EQ(HloOpcode::kConditional, opcode_); - called_computations_[kFalseComputationIndex] = false_computation; + called_computations_[b] = computation; } string HloInstruction::SignatureString() const { @@ -2107,8 +2108,10 @@ bool HloInstruction::IsElementwiseImpl( case HloOpcode::kNegate: case HloOpcode::kReal: case HloOpcode::kReducePrecision: + case HloOpcode::kRsqrt: case HloOpcode::kSign: case HloOpcode::kSin: + case HloOpcode::kSqrt: case HloOpcode::kTanh: CHECK_EQ(1, operand_count()); return true; @@ -2261,10 +2264,21 @@ std::vector HloInstruction::ExtraAttributesToString( extra.push_back( StrCat("scatter=", PrintName(scatter()->name(), options))); } else if (opcode() == HloOpcode::kConditional) { - extra.push_back(StrCat("true_computation=", - PrintName(true_computation()->name(), options))); - extra.push_back(StrCat("false_computation=", - PrintName(false_computation()->name(), options))); + if (operand(0)->shape().element_type() == PRED) { + extra.push_back(StrCat("true_computation=", + PrintName(true_computation()->name(), options))); + extra.push_back( + StrCat("false_computation=", + PrintName(false_computation()->name(), options))); + } else { + extra.push_back(StrCat( + "branch_computations={", + StrJoin(branch_computations(), ", ", + [&](string* out, const HloComputation* computation) { + StrAppend(out, PrintName(computation->name(), options)); + }), + "}")); + } } else if (opcode() == HloOpcode::kCall || opcode() == HloOpcode::kMap || opcode() == HloOpcode::kReduceWindow || opcode() == HloOpcode::kReduce || @@ -2296,10 +2310,20 @@ std::vector HloInstruction::ExtraAttributesToString( extra.push_back(StrCat("scatter=\n", scatter()->ToString(new_options))); break; case HloOpcode::kConditional: - extra.push_back(StrCat("true_computation=\n", - true_computation()->ToString(new_options))); - extra.push_back(StrCat("false_computation=\n", - false_computation()->ToString(new_options))); + if (operand(0)->shape().element_type() == PRED) { + extra.push_back(StrCat("true_computation=\n", + true_computation()->ToString(new_options))); + extra.push_back(StrCat("false_computation=\n", + false_computation()->ToString(new_options))); + } else { + extra.push_back(StrCat( + "branch_computations={\n", + StrJoin(branch_computations(), ",\n", + [&](string* out, const HloComputation* computation) { + StrAppend(out, computation->ToString(new_options)); + }), + "\n}")); + } break; case HloOpcode::kCall: case HloOpcode::kMap: @@ -2545,6 +2569,10 @@ Status HloInstruction::Visit(DfsHloVisitorBase* visitor) { return visitor->HandleCos(this); case HloOpcode::kSin: return visitor->HandleSin(this); + case HloOpcode::kSqrt: + return visitor->HandleSqrt(this); + case HloOpcode::kRsqrt: + return visitor->HandleRsqrt(this); case HloOpcode::kReal: return visitor->HandleReal(this); case HloOpcode::kImag: @@ -2615,6 +2643,8 @@ Status HloInstruction::Visit(DfsHloVisitorBase* visitor) { return visitor->HandleGetDimensionSize(this); case HloOpcode::kTriangularSolve: return visitor->HandleTriangularSolve(this); + case HloOpcode::kCholesky: + return visitor->HandleCholesky(this); // These opcodes are not handled here. case HloOpcode::kTrace: @@ -3499,4 +3529,9 @@ const DomainMetadata& HloInstruction::user_side_metadata() const { const TriangularSolveOptions& HloInstruction::triangular_solve_options() const { return Cast(this)->triangular_solve_options(); } + +const CholeskyOptions& HloInstruction::cholesky_options() const { + return Cast(this)->cholesky_options(); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h index 8ac1636d7159c7cb478856737d93387be49f1ba1..e731137a7fab0ef7e8b94fce825cb72ca33c543e 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.h +++ b/tensorflow/compiler/xla/service/hlo_instruction.h @@ -448,6 +448,9 @@ class HloInstruction { const Shape& shape, HloInstruction* a, HloInstruction* b, const TriangularSolveOptions& options); + static std::unique_ptr CreateCholesky( + const Shape& shape, HloInstruction* a, const CholeskyOptions& options); + // Creates a dot op with operands 'lhs' and 'rhs' with contracting and batch // dimensions specified in 'dimension_numbers'. static std::unique_ptr CreateDot( @@ -708,6 +711,11 @@ class HloInstruction { HloInstruction* true_computation_arg, HloComputation* true_computation, HloInstruction* false_computation_arg, HloComputation* false_computation); + static std::unique_ptr CreateConditional( + const Shape& shape, HloInstruction* branch_index, + absl::Span branch_computations, + absl::Span branch_computation_args); + static std::unique_ptr CreateGather( const Shape& shape, HloInstruction* operand, HloInstruction* start_indices, @@ -949,6 +957,10 @@ class HloInstruction { // operands of it which could be created due to this replacement. Status ReplaceUseWith(HloInstruction* user, HloInstruction* new_producer); + // Same as ReplaceUseWith(), but new_producer can have a different shape. + Status ReplaceUseWithDifferentShape(HloInstruction* user, + HloInstruction* new_producer); + // Replaces the specified operand with new_operand. The old and new operands // must have compatible shapes ignoring floating-point precision. // @@ -1050,14 +1062,23 @@ class HloInstruction { HloInstruction* while_init() const; - // Gets/sets the true and false HloComputation for Conditional. The setters - // should only be called by HloModule or HloComputation methods. + // Gets/sets the true and false HloComputation for Conditional. // - // Precondition: The instruction is a Conditional instruction. + // Precondition: The instruction is a predicated Conditional instruction. HloComputation* true_computation() const; HloComputation* false_computation() const; - void set_true_computation(HloComputation* true_computation); - void set_false_computation(HloComputation* false_computation); + + // Gets the branch HloComputations for Conditional. + // + // Precondition: The instruction is a Conditional instruction. + const std::vector& branch_computations() const; + int branch_count() const; + HloComputation* branch_computation(int b) const; + // Sets a branch HloComputation for Conditional. + // The setter should only be called by HloModule or HloComputation methods. + // + // Precondition: The instruction is a Conditional instruction. + void set_branch_computation(int b, HloComputation* computation); // Returns a string for the signature of this instruction if considered as a // function, e.g. the signature of an F32 add is (F32, F32) -> F32. @@ -1590,6 +1611,9 @@ class HloInstruction { // Delegates to HloTriangularSolveInstruction::triangular_solve_options(). const TriangularSolveOptions& triangular_solve_options() const; + // Delegates to HloCholeskyInstruction::cholesky_options(). + const CholeskyOptions& cholesky_options() const; + // Old methods kept for smooth subclassing transition END. protected: diff --git a/tensorflow/compiler/xla/service/hlo_instructions.cc b/tensorflow/compiler/xla/service/hlo_instructions.cc index 905a6fe08b4430ad862edf0886a57c9f7e9f7977..7d18b35c2bb2d6b0a6435ecb397512838fe4b72e 100644 --- a/tensorflow/compiler/xla/service/hlo_instructions.cc +++ b/tensorflow/compiler/xla/service/hlo_instructions.cc @@ -30,6 +30,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_sharding_metadata.h" #include "tensorflow/compiler/xla/window_util.h" +#include "tensorflow/core/platform/protobuf.h" namespace xla { namespace { @@ -201,6 +202,46 @@ std::unique_ptr HloFftInstruction::CloneWithNewOperandsImpl( fft_length_); } +namespace { + +// Converts a protocol buffer message (e.g., TriangularSolveOptions) to a vector +// of "key=value" attribute strings generically, using protocol buffer +// reflection. +// +// Currently implements a small subset of cases; feel free to add more as +// needed. +std::vector AttributeProtoToStringVector( + const tensorflow::protobuf::Message& message) { + const tensorflow::protobuf::Reflection* reflection = message.GetReflection(); + std::vector fields; + reflection->ListFields(message, &fields); + + std::vector output; + for (const tensorflow::protobuf::FieldDescriptor* field : fields) { + string s = absl::StrCat(field->name(), "="); + CHECK(!field->is_repeated()) << "Repeated fields aren't implemented"; + switch (field->type()) { + case tensorflow::protobuf::FieldDescriptor::TYPE_BOOL: { + bool val = reflection->GetBool(message, field); + absl::StrAppend(&s, val ? "true" : "false"); + break; + } + case tensorflow::protobuf::FieldDescriptor::TYPE_ENUM: { + const tensorflow::protobuf::EnumValueDescriptor* evd = + reflection->GetEnum(message, field); + absl::StrAppend(&s, evd->name()); + break; + } + default: + LOG(FATAL) << "Unimplemented field type: " << field->DebugString(); + } + output.push_back(std::move(s)); + } + return output; +} + +} // namespace + HloTriangularSolveInstruction::HloTriangularSolveInstruction( const Shape& shape, HloInstruction* a, HloInstruction* b, const TriangularSolveOptions& options) @@ -218,14 +259,7 @@ HloInstructionProto HloTriangularSolveInstruction::ToProto() const { std::vector HloTriangularSolveInstruction::ExtraAttributesToStringImpl( const HloPrintOptions& options) const { - return { - StrCat("left_side=", - triangular_solve_options_.left_side() ? "true" : "false"), - StrCat("lower=", triangular_solve_options_.lower() ? "true" : "false"), - StrCat("unit_diagonal=", - triangular_solve_options_.unit_diagonal() ? "true" : "false"), - StrCat("transpose_a=", TriangularSolveOptions_Transpose_Name( - triangular_solve_options_.transpose_a()))}; + return AttributeProtoToStringVector(triangular_solve_options_); } bool HloTriangularSolveInstruction::IdenticalSlowPath( @@ -252,6 +286,44 @@ HloTriangularSolveInstruction::CloneWithNewOperandsImpl( shape, new_operands[0], new_operands[1], triangular_solve_options()); } +HloCholeskyInstruction::HloCholeskyInstruction(const Shape& shape, + HloInstruction* a, + const CholeskyOptions& options) + : HloInstruction(HloOpcode::kCholesky, shape), cholesky_options_(options) { + AppendOperand(a); +} + +HloInstructionProto HloCholeskyInstruction::ToProto() const { + HloInstructionProto proto = HloInstruction::ToProto(); + *proto.mutable_cholesky_options() = cholesky_options_; + return proto; +} + +std::vector HloCholeskyInstruction::ExtraAttributesToStringImpl( + const HloPrintOptions& options) const { + return AttributeProtoToStringVector(cholesky_options_); +} + +bool HloCholeskyInstruction::IdenticalSlowPath( + const HloInstruction& other, + const std::function& + eq_computations) const { + const auto& casted_other = static_cast(other); + const auto& options = cholesky_options(); + const auto& other_options = casted_other.cholesky_options(); + + return options.lower() == other_options.lower(); +} + +std::unique_ptr +HloCholeskyInstruction::CloneWithNewOperandsImpl( + const Shape& shape, absl::Span new_operands, + HloCloneContext* context) const { + CHECK_EQ(new_operands.size(), 1); + return absl::make_unique(shape, new_operands[0], + cholesky_options()); +} + HloSendRecvInstruction::HloSendRecvInstruction(HloOpcode opcode, const Shape& shape, int64 channel_id, diff --git a/tensorflow/compiler/xla/service/hlo_instructions.h b/tensorflow/compiler/xla/service/hlo_instructions.h index 4d23cb671f24623f56faa9b69015cef21752a799..43aa12c10f2f03bd8a4e59ec32cd2bbdec11f3a2 100644 --- a/tensorflow/compiler/xla/service/hlo_instructions.h +++ b/tensorflow/compiler/xla/service/hlo_instructions.h @@ -159,6 +159,31 @@ class HloTriangularSolveInstruction : public HloInstruction { TriangularSolveOptions triangular_solve_options_; }; +class HloCholeskyInstruction : public HloInstruction { + public: + explicit HloCholeskyInstruction(const Shape& shape, HloInstruction* a, + const CholeskyOptions& options); + const CholeskyOptions& cholesky_options() const { return cholesky_options_; } + + // Returns a serialized representation of this instruction. + HloInstructionProto ToProto() const override; + + private: + std::vector ExtraAttributesToStringImpl( + const HloPrintOptions& options) const override; + bool IdenticalSlowPath( + const HloInstruction& other, + const std::function& + eq_computations) const override; + + // Implementation for non-common logic of CloneWithNewOperands. + std::unique_ptr CloneWithNewOperandsImpl( + const Shape& shape, absl::Span new_operands, + HloCloneContext* context) const override; + + CholeskyOptions cholesky_options_; +}; + class HloSendRecvInstruction : public HloInstruction { public: // Returns the channel id associated with the instruction. The id is diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc index 8322870cfd6a89fc6f863da8fd4a3576e8845cd7..edcda8f9a7b974b95a12348577c335a3e6d8fcce 100644 --- a/tensorflow/compiler/xla/service/hlo_module.cc +++ b/tensorflow/compiler/xla/service/hlo_module.cc @@ -158,17 +158,12 @@ void HloModule::ReplaceComputations( break; } case HloOpcode::kConditional: { - HloComputation* new_true_computation = - tensorflow::gtl::FindWithDefault( - replacements, instruction->true_computation(), nullptr); - if (new_true_computation != nullptr) { - instruction->set_true_computation(new_true_computation); - } - HloComputation* new_false_computation = - tensorflow::gtl::FindWithDefault( - replacements, instruction->false_computation(), nullptr); - if (new_false_computation != nullptr) { - instruction->set_false_computation(new_false_computation); + for (int b = 0; b < instruction->branch_count(); ++b) { + HloComputation* new_computation = tensorflow::gtl::FindWithDefault( + replacements, instruction->branch_computation(b), nullptr); + if (new_computation != nullptr) { + instruction->set_branch_computation(b, new_computation); + } } break; } diff --git a/tensorflow/compiler/xla/service/hlo_module.h b/tensorflow/compiler/xla/service/hlo_module.h index b6fe6a5cdbd0934014f1152acd48c7a5973bead3..2c63247eea8292f52e95b6171100221336450c13 100644 --- a/tensorflow/compiler/xla/service/hlo_module.h +++ b/tensorflow/compiler/xla/service/hlo_module.h @@ -167,6 +167,12 @@ class HloModule { // Gets the number of computations in this module. int64 computation_count() const { return computations_.size(); } + // Returns the mutable computation for the given index. + HloComputation* mutable_computation(int64 idx) { + CHECK(idx >= 0 && idx < computations_.size()); + return computations_[idx].get(); + } + // Gets the number of instructions in this module. int64 instruction_count() const; diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc index b877081be5775bf6c75a69ffeba28d0f2cc17f90..bc258a77000d17cdb6b1d1005b6dac70e300e398 100644 --- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc +++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc @@ -45,11 +45,8 @@ string HloModuleGroupMetadata::TrackedInstruction::ToString() const { case ComputationKind::kWhileBody: repr += ":WHILE_BODY"; break; - case ComputationKind::kConditionalTrue: - repr += ":CONDITIONAL_TRUE"; - break; - case ComputationKind::kConditionalFalse: - repr += ":CONDITIONAL_FALSE"; + case ComputationKind::kConditionalBranch: + repr += absl::StrCat(":CONDITIONAL_BRANCH_", index_); break; case ComputationKind::kCallFunction: repr += ":CALL"; @@ -307,10 +304,10 @@ Status HloModuleGroupMetadata::RecordInstructions() { tracked_instructions_[hlo->while_body()] = TrackedInstruction(hlo, ComputationKind::kWhileBody); } else if (hlo->opcode() == HloOpcode::kConditional) { - tracked_instructions_[hlo->true_computation()] = - TrackedInstruction(hlo, ComputationKind::kConditionalTrue); - tracked_instructions_[hlo->false_computation()] = - TrackedInstruction(hlo, ComputationKind::kConditionalFalse); + for (int b = 0; b < hlo->branch_count(); ++b) { + tracked_instructions_[hlo->branch_computation(b)] = + TrackedInstruction(hlo, ComputationKind::kConditionalBranch, b); + } } else if (hlo->opcode() == HloOpcode::kCall) { tracked_instructions_[hlo->to_apply()] = TrackedInstruction(hlo, ComputationKind::kCallFunction); diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h index 84f7f2f31339ae9e98ea2301b6e6d94fcf4dedbb..07becfc3638a550b661e2ee0d4f10ac5e836e481 100644 --- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h +++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h @@ -67,8 +67,7 @@ class HloModuleGroupMetadata { kInvalid, kWhileCondition, kWhileBody, - kConditionalTrue, - kConditionalFalse, + kConditionalBranch, kCallFunction, }; @@ -80,12 +79,13 @@ class HloModuleGroupMetadata { class TrackedInstruction { public: TrackedInstruction() = default; - TrackedInstruction(HloInstruction* instruction, ComputationKind kind) - : instruction_(instruction), kind_(kind) {} + TrackedInstruction(HloInstruction* instruction, ComputationKind kind, + int index = -1) + : instruction_(instruction), kind_(kind), index_(index) {} bool operator==(const TrackedInstruction& rhs) const { return instruction_->opcode() == rhs.instruction_->opcode() && - kind_ == rhs.kind_; + kind_ == rhs.kind_ && index_ == rhs.index_; } bool operator!=(const TrackedInstruction& rhs) const { return !operator==(rhs); @@ -98,6 +98,7 @@ class HloModuleGroupMetadata { private: HloInstruction* instruction_ = nullptr; ComputationKind kind_ = ComputationKind::kInvalid; + int index_ = -1; }; // Represents a channel and the instructions that form the channel. diff --git a/tensorflow/compiler/xla/service/hlo_opcode.cc b/tensorflow/compiler/xla/service/hlo_opcode.cc index 4551a1c2e259b06818f913cb6a9e782436b7e594..548fbb873aa646e061fb990454bb555d098607d8 100644 --- a/tensorflow/compiler/xla/service/hlo_opcode.cc +++ b/tensorflow/compiler/xla/service/hlo_opcode.cc @@ -53,8 +53,8 @@ StatusOr StringToHloOpcode(const string& opcode_name) { bool HloOpcodeIsComparison(HloOpcode opcode) { switch (opcode) { -#define CASE_IS_COMPARISON(enum_name, ...) \ - case HloOpcode::enum_name: \ +#define CASE_IS_COMPARISON(enum_name, opcode_name, ...) \ + case HloOpcode::enum_name: \ return HAS_PROPERTY(kHloOpcodeIsComparison, __VA_ARGS__); HLO_OPCODE_LIST(CASE_IS_COMPARISON) #undef CASE_IS_COMPARISON @@ -63,14 +63,25 @@ bool HloOpcodeIsComparison(HloOpcode opcode) { bool HloOpcodeIsVariadic(HloOpcode opcode) { switch (opcode) { -#define CASE_IS_VARIADIC(enum_name, ...) \ - case HloOpcode::enum_name: \ - return HAS_PROPERTY(kHloOpcodeIsVariadic, __VA_ARGS__); +#define CASE_IS_VARIADIC(enum_name, opcode_name, arity, ...) \ + case HloOpcode::enum_name: \ + return arity == kHloOpcodeIsVariadic; HLO_OPCODE_LIST(CASE_IS_VARIADIC) #undef CASE_IS_VARIADIC } } +absl::optional HloOpcodeArity(HloOpcode opcode) { + switch (opcode) { +#define CASE_ARITY(enum_name, opcode_name, arity, ...) \ + case HloOpcode::enum_name: \ + return arity == kHloOpcodeIsVariadic ? absl::nullopt \ + : absl::make_optional(arity); + HLO_OPCODE_LIST(CASE_ARITY) +#undef CASE_ARITY + } +} + #undef HAS_PROPERTY #undef RESOLVE #undef CHECK_DEFAULT diff --git a/tensorflow/compiler/xla/service/hlo_opcode.h b/tensorflow/compiler/xla/service/hlo_opcode.h index 35626ba37541fc7a984ad05d12ebc22e9a08a550..3e144c4472f09c50b68866fa50ddee409793293f 100644 --- a/tensorflow/compiler/xla/service/hlo_opcode.h +++ b/tensorflow/compiler/xla/service/hlo_opcode.h @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "absl/types/optional.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" @@ -30,9 +31,9 @@ namespace xla { // See the XLA documentation for the semantics of each opcode. // // Each entry has the format: -// (enum_name, opcode_name) +// (enum_name, opcode_name, arity) // or -// (enum_name, opcode_name, p1 | p2 | ...) +// (enum_name, opcode_name, arity, p1 | p2 | ...) // // with p1, p2, ... are members of HloOpcodeProperty. They are combined // using bitwise-or. @@ -44,104 +45,107 @@ namespace xla { // - In fully qualified names (HloInstruction::FullyQualifiedName()), to // separate the qualifiers (name of the computation and potentially the // fusion instruction) from the name -#define HLO_OPCODE_LIST(V) \ - V(kAbs, "abs") \ - V(kAdd, "add") \ - V(kAddDependency, "add-dependency") \ - V(kAfterAll, "after-all", kHloOpcodeIsVariadic) \ - V(kAllReduce, "all-reduce") \ - V(kAllToAll, "all-to-all") \ - V(kAtan2, "atan2") \ - V(kBatchNormGrad, "batch-norm-grad") \ - V(kBatchNormInference, "batch-norm-inference") \ - V(kBatchNormTraining, "batch-norm-training") \ - V(kBitcast, "bitcast") \ - V(kBitcastConvert, "bitcast-convert") \ - V(kBroadcast, "broadcast") \ - V(kCall, "call", kHloOpcodeIsVariadic) \ - V(kCeil, "ceil") \ - V(kClamp, "clamp") \ - V(kCollectivePermute, "collective-permute") \ - V(kClz, "count-leading-zeros") \ - V(kComplex, "complex") \ - V(kConcatenate, "concatenate", kHloOpcodeIsVariadic) \ - V(kConditional, "conditional") \ - V(kConstant, "constant") \ - V(kConvert, "convert") \ - V(kConvolution, "convolution") \ - V(kCopy, "copy") \ - V(kCos, "cosine") \ - V(kCustomCall, "custom-call") \ - V(kDivide, "divide") \ - V(kDomain, "domain") \ - V(kDot, "dot") \ - V(kDynamicSlice, "dynamic-slice") \ - V(kDynamicUpdateSlice, "dynamic-update-slice") \ - V(kEq, "equal-to", kHloOpcodeIsComparison) \ - V(kExp, "exponential") \ - V(kExpm1, "exponential-minus-one") \ - V(kFft, "fft") \ - V(kFloor, "floor") \ - V(kFusion, "fusion", kHloOpcodeIsVariadic) \ - V(kGather, "gather") \ - V(kGe, "greater-than-or-equal-to", kHloOpcodeIsComparison) \ - V(kGetDimensionSize, "get-dimension-size") \ - V(kGetTupleElement, "get-tuple-element") \ - V(kGt, "greater-than", kHloOpcodeIsComparison) \ - V(kImag, "imag") \ - V(kInfeed, "infeed") \ - V(kIota, "iota") \ - V(kIsFinite, "is-finite") \ - V(kLe, "less-than-or-equal-to", kHloOpcodeIsComparison) \ - V(kLog, "log") \ - V(kLog1p, "log-plus-one") \ - V(kAnd, "and") \ - V(kNot, "not") \ - V(kOr, "or") \ - V(kXor, "xor") \ - V(kLt, "less-than", kHloOpcodeIsComparison) \ - V(kMap, "map", kHloOpcodeIsVariadic) \ - V(kMaximum, "maximum") \ - V(kMinimum, "minimum") \ - V(kMultiply, "multiply") \ - V(kNe, "not-equal-to", kHloOpcodeIsComparison) \ - V(kNegate, "negate") \ - V(kOutfeed, "outfeed") \ - V(kPad, "pad") \ - V(kParameter, "parameter") \ - V(kPower, "power") \ - V(kReal, "real") \ - V(kRecv, "recv") \ - V(kRecvDone, "recv-done") \ - V(kReduce, "reduce") \ - V(kReducePrecision, "reduce-precision") \ - V(kReduceWindow, "reduce-window") \ - V(kRemainder, "remainder") \ - V(kReshape, "reshape") \ - V(kReverse, "reverse") \ - V(kRng, "rng") \ - V(kRoundNearestAfz, "round-nearest-afz") \ - V(kScatter, "scatter") \ - V(kSelect, "select") \ - V(kSelectAndScatter, "select-and-scatter") \ - V(kSend, "send") \ - V(kSendDone, "send-done") \ - V(kShiftLeft, "shift-left") \ - V(kShiftRightArithmetic, "shift-right-arithmetic") \ - V(kShiftRightLogical, "shift-right-logical") \ - V(kSign, "sign") \ - V(kSin, "sine") \ - V(kSlice, "slice") \ - V(kSort, "sort") \ - V(kSubtract, "subtract") \ - V(kTanh, "tanh") \ - V(kTrace, "trace") \ - V(kTranspose, "transpose") \ - V(kTriangularSolve, "triangular-solve") \ - V(kTuple, "tuple", kHloOpcodeIsVariadic) \ - V(kTupleSelect, "tuple-select") \ - V(kWhile, "while") \ - V(kReplicaId, "replica-id") +#define HLO_OPCODE_LIST(V) \ + V(kAbs, "abs", 1) \ + V(kAdd, "add", 2) \ + V(kAddDependency, "add-dependency", 2) \ + V(kAfterAll, "after-all", kHloOpcodeIsVariadic) \ + V(kAllReduce, "all-reduce", kHloOpcodeIsVariadic) \ + V(kAllToAll, "all-to-all", kHloOpcodeIsVariadic) \ + V(kAtan2, "atan2", 2) \ + V(kBatchNormGrad, "batch-norm-grad", 5) \ + V(kBatchNormInference, "batch-norm-inference", 5) \ + V(kBatchNormTraining, "batch-norm-training", 3) \ + V(kBitcast, "bitcast", 1) \ + V(kBitcastConvert, "bitcast-convert", 1) \ + V(kBroadcast, "broadcast", 1) \ + V(kCall, "call", kHloOpcodeIsVariadic) \ + V(kCeil, "ceil", 1) \ + V(kCholesky, "cholesky", 1) \ + V(kClamp, "clamp", 3) \ + V(kCollectivePermute, "collective-permute", 1) \ + V(kClz, "count-leading-zeros", 1) \ + V(kComplex, "complex", 2) \ + V(kConcatenate, "concatenate", kHloOpcodeIsVariadic) \ + V(kConditional, "conditional", kHloOpcodeIsVariadic) \ + V(kConstant, "constant", 0) \ + V(kConvert, "convert", 1) \ + V(kConvolution, "convolution", 2) \ + V(kCopy, "copy", 1) \ + V(kCos, "cosine", 1) \ + V(kCustomCall, "custom-call", kHloOpcodeIsVariadic) \ + V(kDivide, "divide", 2) \ + V(kDomain, "domain", 1) \ + V(kDot, "dot", 2) \ + V(kDynamicSlice, "dynamic-slice", kHloOpcodeIsVariadic) \ + V(kDynamicUpdateSlice, "dynamic-update-slice", kHloOpcodeIsVariadic) \ + V(kEq, "equal-to", 2, kHloOpcodeIsComparison) \ + V(kExp, "exponential", 1) \ + V(kExpm1, "exponential-minus-one", 1) \ + V(kFft, "fft", 1) \ + V(kFloor, "floor", 1) \ + V(kFusion, "fusion", kHloOpcodeIsVariadic) \ + V(kGather, "gather", 2) \ + V(kGe, "greater-than-or-equal-to", 2, kHloOpcodeIsComparison) \ + V(kGetDimensionSize, "get-dimension-size", 1) \ + V(kGetTupleElement, "get-tuple-element", 1) \ + V(kGt, "greater-than", 2, kHloOpcodeIsComparison) \ + V(kImag, "imag", 1) \ + V(kInfeed, "infeed", 1) \ + V(kIota, "iota", 0) \ + V(kIsFinite, "is-finite", 1) \ + V(kLe, "less-than-or-equal-to", 2, kHloOpcodeIsComparison) \ + V(kLog, "log", 1) \ + V(kLog1p, "log-plus-one", 1) \ + V(kAnd, "and", 2) \ + V(kNot, "not", 1) \ + V(kOr, "or", 2) \ + V(kXor, "xor", 2) \ + V(kLt, "less-than", 2, kHloOpcodeIsComparison) \ + V(kMap, "map", kHloOpcodeIsVariadic) \ + V(kMaximum, "maximum", 2) \ + V(kMinimum, "minimum", 2) \ + V(kMultiply, "multiply", 2) \ + V(kNe, "not-equal-to", 2, kHloOpcodeIsComparison) \ + V(kNegate, "negate", 1) \ + V(kOutfeed, "outfeed", 2) \ + V(kPad, "pad", 2) \ + V(kParameter, "parameter", 0) \ + V(kPower, "power", 2) \ + V(kReal, "real", 1) \ + V(kRecv, "recv", 1) \ + V(kRecvDone, "recv-done", 1) \ + V(kReduce, "reduce", kHloOpcodeIsVariadic) \ + V(kReducePrecision, "reduce-precision", 1) \ + V(kReduceWindow, "reduce-window", 2) \ + V(kRemainder, "remainder", 2) \ + V(kReplicaId, "replica-id", 0) \ + V(kReshape, "reshape", 1) \ + V(kReverse, "reverse", 1) \ + V(kRng, "rng", kHloOpcodeIsVariadic) \ + V(kRoundNearestAfz, "round-nearest-afz", 1) \ + V(kRsqrt, "rsqrt", 1) \ + V(kScatter, "scatter", 3) \ + V(kSelect, "select", 3) \ + V(kSelectAndScatter, "select-and-scatter", 3) \ + V(kSend, "send", 2) \ + V(kSendDone, "send-done", 1) \ + V(kShiftLeft, "shift-left", 2) \ + V(kShiftRightArithmetic, "shift-right-arithmetic", 2) \ + V(kShiftRightLogical, "shift-right-logical", 2) \ + V(kSign, "sign", 1) \ + V(kSin, "sine", 1) \ + V(kSlice, "slice", 1) \ + V(kSort, "sort", kHloOpcodeIsVariadic) \ + V(kSqrt, "sqrt", 1) \ + V(kSubtract, "subtract", 2) \ + V(kTanh, "tanh", 1) \ + V(kTrace, "trace", 1) \ + V(kTranspose, "transpose", 1) \ + V(kTriangularSolve, "triangular-solve", 2) \ + V(kTuple, "tuple", kHloOpcodeIsVariadic) \ + V(kTupleSelect, "tuple-select", 3) \ + V(kWhile, "while", 1) enum class HloOpcode { #define DECLARE_ENUM(enum_name, opcode_name, ...) enum_name, @@ -149,12 +153,16 @@ enum class HloOpcode { #undef DECLARE_ENUM }; +// Arity value that denotes that an operator is variadic. +enum { + kHloOpcodeIsVariadic = -1, +}; + // List of properties associated with opcodes. // Properties are defined as increasing powers of two, so that we can use // bitwise-or to combine properties, and bitwise-and to test for them. enum HloOpcodeProperty { kHloOpcodeIsComparison = 1 << 0, - kHloOpcodeIsVariadic = 1 << 1, }; // Returns a string representation of the opcode. @@ -173,6 +181,10 @@ bool HloOpcodeIsComparison(HloOpcode opcode); // Returns true iff the given opcode has variadic operands. bool HloOpcodeIsVariadic(HloOpcode opcode); +// Returns the arity of opcode. If the opcode is variadic, +// returns nullopt. +absl::optional HloOpcodeArity(HloOpcode opcode); + // Returns the number of HloOpcode values. inline const uint32_t HloOpcodeCount() { #define HLO_COUNT_ONE(...) +1 diff --git a/tensorflow/compiler/xla/service/hlo_opcode_test.cc b/tensorflow/compiler/xla/service/hlo_opcode_test.cc index 6f3f83f63a05fafaa3f3ddcff8a7cac7cb7b06d5..910cc25a5917183e24599df0da1b4127e639f536 100644 --- a/tensorflow/compiler/xla/service/hlo_opcode_test.cc +++ b/tensorflow/compiler/xla/service/hlo_opcode_test.cc @@ -54,11 +54,20 @@ TEST(HloOpcodeTest, OpcodeProperties) { EXPECT_FALSE(HloOpcodeIsComparison(opcode)); } switch (opcode) { + case HloOpcode::kAfterAll: + case HloOpcode::kAllReduce: + case HloOpcode::kAllToAll: case HloOpcode::kCall: case HloOpcode::kConcatenate: + case HloOpcode::kConditional: + case HloOpcode::kCustomCall: + case HloOpcode::kDynamicSlice: + case HloOpcode::kDynamicUpdateSlice: case HloOpcode::kFusion: case HloOpcode::kMap: - case HloOpcode::kAfterAll: + case HloOpcode::kReduce: + case HloOpcode::kRng: + case HloOpcode::kSort: case HloOpcode::kTuple: EXPECT_TRUE(HloOpcodeIsVariadic(opcode)); break; diff --git a/tensorflow/compiler/xla/service/hlo_ordering.cc b/tensorflow/compiler/xla/service/hlo_ordering.cc index 0cec61c257bb84e467290fb52ec9063a32ed558d..831771fe63b8dd4c276ad3ec05ea90b4d475e7e0 100644 --- a/tensorflow/compiler/xla/service/hlo_ordering.cc +++ b/tensorflow/compiler/xla/service/hlo_ordering.cc @@ -66,24 +66,31 @@ bool HloOrdering::ExecutesBefore(const HloInstruction* a, } } - // If the common ancestor is a conditional instruction, even though the true - // and false computations are not really ordered per-se, we define the true - // computation to be ordered before the false one. - // This ensures that buffers can still be shared among the two computations + // If the common ancestor is a conditional instruction, even though the branch + // computations are not really ordered per-se, we define the 0th branch + // computation to be ordered before the 1st one, before the 2nd and so forth. + // This ensures that buffers can still be shared among branch computations // as they will forcibly have disjoint liveness. if (a_ancestor == b_ancestor && - a_ancestor->opcode() == HloOpcode::kConditional) { - const HloComputation* true_computation = a_ancestor->true_computation(); - const HloComputation* false_computation = a_ancestor->false_computation(); - if (call_graph_->InstructionIsNestedIn(a, true_computation) && - call_graph_->InstructionIsNestedIn(b, false_computation)) { + (a_ancestor->opcode() == HloOpcode::kConditional)) { + int a_branch = -1; + int b_branch = -1; + for (int j = 0; j < a_ancestor->branch_count(); ++j) { + if (call_graph_->InstructionIsNestedIn( + a, a_ancestor->branch_computation(j))) { + a_branch = j; + } + if (call_graph_->InstructionIsNestedIn( + b, a_ancestor->branch_computation(j))) { + b_branch = j; + } + } + if (a_branch != -1 && a_branch < b_branch) { return true; } - // If 'b' is the conditional ancestor, and 'a' is within the true or false - // computations, 'a' executes before 'b'. - if (b == a_ancestor && - (call_graph_->InstructionIsNestedIn(a, true_computation) || - call_graph_->InstructionIsNestedIn(a, false_computation))) { + // If 'b' is the conditional ancestor, and 'a' is within a branch + // computation, 'a' executes before 'b'. + if (b == a_ancestor && a_branch != -1) { return true; } } @@ -144,17 +151,17 @@ bool HloOrdering::IsDefinedBefore(const HloValue& a, const HloValue& b) const { b.defining_instruction()->while_condition()))) { return true; } - // If 'b' is a conditional phi and 'a' is in the true or false computation, - // then 'a' executes before 'b'. + // If 'b' is a conditional phi and 'a' is in some branch computation, then 'a' + // executes before 'b'. if (b.is_phi() && - b.defining_instruction()->opcode() == HloOpcode::kConditional && - (call_graph_->InstructionIsNestedIn( - a.defining_instruction(), - b.defining_instruction()->true_computation()) || - call_graph_->InstructionIsNestedIn( - a.defining_instruction(), - b.defining_instruction()->false_computation()))) { - return true; + b.defining_instruction()->opcode() == HloOpcode::kConditional) { + for (int j = 0; j < b.defining_instruction()->branch_count(); ++j) { + if (call_graph_->InstructionIsNestedIn( + a.defining_instruction(), + b.defining_instruction()->branch_computation(j))) { + return true; + } + } } return ExecutesBefore(a.defining_instruction(), b.defining_instruction()); } @@ -225,17 +232,14 @@ bool HloOrdering::UseIsBeforeValueDefinition( if (use.instruction->opcode() == HloOpcode::kConditional) { const HloInstruction* conditional = use.instruction; - if (call_graph_->InstructionIsNestedIn(value.defining_instruction(), - conditional->true_computation())) { - VLOG(4) << " use is conditional " << use.instruction->name() - << " and def is in TRUE computation"; - return true; - } - if (call_graph_->InstructionIsNestedIn(value.defining_instruction(), - conditional->false_computation())) { - VLOG(4) << " use is conditional " << use.instruction->name() - << " and def is in FALSE computation"; - return true; + for (int j = 0; j < conditional->branch_count(); ++j) { + if (call_graph_->InstructionIsNestedIn( + value.defining_instruction(), + conditional->branch_computation(j))) { + VLOG(4) << " use is conditional " << use.instruction->name() + << " and def is in " << j << "th branch computation"; + return true; + } } if (value.defining_instruction() == use.instruction) { VLOG(4) << " use is conditional " << use << " and def is " diff --git a/tensorflow/compiler/xla/service/hlo_parser.cc b/tensorflow/compiler/xla/service/hlo_parser.cc index f448571082e52e4b81db1c68d1e1470935386139..fd55d92c04a7a2d0b89e2bd2e1e745d22a3a48e1 100644 --- a/tensorflow/compiler/xla/service/hlo_parser.cc +++ b/tensorflow/compiler/xla/service/hlo_parser.cc @@ -22,6 +22,7 @@ limitations under the License. #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" #include "absl/strings/str_split.h" +#include "absl/types/span.h" #include "absl/types/variant.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/literal_util.h" @@ -35,6 +36,7 @@ limitations under the License. #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/lib/gtl/map_util.h" +#include "tensorflow/core/platform/protobuf.h" namespace xla { @@ -179,8 +181,8 @@ class HloParser { kBracedInt64List, kBracedInt64ListList, kHloComputation, + kBracedHloComputationList, kFftType, - kTriangularSolveTranspose, kWindow, kConvolutionDimensionNumbers, kSharding, @@ -237,6 +239,21 @@ class HloParser { bool ParseAttributeHelper(const std::unordered_map& attrs, std::unordered_set* seen_attrs); + // Parses an attribute string into a protocol buffer `message`. + // Since proto3 has no notion of mandatory fields, `required_attrs` gives the + // set of mandatory attributes. + bool ParseAttributesAsProtoMessage( + const std::unordered_set& required_attrs, + tensorflow::protobuf::Message* message); + + // Parses one attribute. If it has already been seen, return error. Returns + // true and adds to seen_attrs on success. + // + // Do not call this except in ParseAttributesAsProtoMessage. + bool ParseAttributeAsProtoMessageHelper( + tensorflow::protobuf::Message* message, + std::unordered_set* seen_attrs); + // Parses a name and finds the corresponding hlo computation. bool ParseComputationName(HloComputation** value); // Parses a list of names and finds the corresponding hlo instructions. @@ -261,6 +278,8 @@ class HloParser { bool ParseSliceRanges(SliceRanges* result); bool ParsePrecisionList(std::vector* result); + bool ParseHloComputation(HloComputation** result); + bool ParseHloComputationList(std::vector* result); bool ParseShapeList(std::vector* result); bool ParseInt64List(const TokKind start, const TokKind end, const TokKind delim, std::vector* result); @@ -281,7 +300,6 @@ class HloParser { bool ParseTiles(std::vector* tiles); bool ParseOpcode(HloOpcode* result); bool ParseFftType(FftType* result); - bool ParseTriangularSolveTranspose(TriangularSolveOptions::Transpose* result); bool ParseFusionKind(HloInstruction::FusionKind* result); bool ParseRandomDistribution(RandomDistribution* result); bool ParsePrecision(PrecisionConfig::Precision* result); @@ -725,8 +743,10 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, case HloOpcode::kNot: case HloOpcode::kNegate: case HloOpcode::kReal: + case HloOpcode::kRsqrt: case HloOpcode::kSign: case HloOpcode::kSin: + case HloOpcode::kSqrt: case HloOpcode::kTanh: { if (!ParseOperands(&operands, /*expected_size=*/1) || !ParseAttributes(attrs)) { @@ -1102,37 +1122,28 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kTriangularSolve: { - optional left_side; - optional lower; - optional unit_diagonal; - optional transpose_a; - attrs["left_side"] = {/*required=*/false, AttrTy::kBool, &left_side}; - attrs["lower"] = {/*required=*/false, AttrTy::kBool, &lower}; - attrs["unit_diagonal"] = {/*required=*/false, AttrTy::kBool, - &unit_diagonal}; - attrs["transpose_a"] = {/*required=*/false, - AttrTy::kTriangularSolveTranspose, &transpose_a}; + TriangularSolveOptions options; if (!ParseOperands(&operands, /*expected_size=*/2) || - !ParseAttributes(attrs)) { + !ParseAttributesAsProtoMessage( + /*required_attrs=*/std::unordered_set(), &options)) { return false; } - TriangularSolveOptions options; - if (left_side) { - options.set_left_side(*left_side); - } - if (lower) { - options.set_lower(*lower); - } - if (unit_diagonal) { - options.set_unit_diagonal(*unit_diagonal); - } - options.set_transpose_a( - transpose_a ? *transpose_a : TriangularSolveOptions::NO_TRANSPOSE); instruction = builder->AddInstruction(HloInstruction::CreateTriangularSolve( shape, operands[0], operands[1], options)); break; } + case HloOpcode::kCholesky: { + CholeskyOptions options; + if (!ParseOperands(&operands, /*expected_size=*/1) || + !ParseAttributesAsProtoMessage( + /*required_attrs=*/std::unordered_set(), &options)) { + return false; + } + instruction = builder->AddInstruction( + HloInstruction::CreateCholesky(shape, operands[0], options)); + break; + } case HloOpcode::kBroadcast: { optional> broadcast_dimensions; attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, @@ -1429,18 +1440,36 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, case HloOpcode::kConditional: { optional true_computation; optional false_computation; - attrs["true_computation"] = {/*required=*/true, AttrTy::kHloComputation, - &true_computation}; - attrs["false_computation"] = {/*required=*/true, AttrTy::kHloComputation, - &false_computation}; - if (!ParseOperands(&operands, /*expected_size=*/3) || - !ParseAttributes(attrs)) { + optional> branch_computations; + if (!ParseOperands(&operands)) { + return false; + } + const bool branch_index_is_bool = + operands[0]->shape().element_type() == PRED; + if (branch_index_is_bool) { + attrs["true_computation"] = {/*required=*/true, AttrTy::kHloComputation, + &true_computation}; + attrs["false_computation"] = { + /*required=*/true, AttrTy::kHloComputation, &false_computation}; + } else { + attrs["branch_computations"] = {/*required=*/true, + AttrTy::kBracedHloComputationList, + &branch_computations}; + } + if (!ParseAttributes(attrs)) { + return false; + } + if (branch_index_is_bool) { + branch_computations.emplace({*true_computation, *false_computation}); + } + if (branch_computations->empty() || + operands.size() != branch_computations->size() + 1) { return false; } instruction = builder->AddInstruction(HloInstruction::CreateConditional( - shape, /*pred=*/operands[0], - /*true_computation_arg=*/operands[1], *true_computation, - /*false_computation_arg=*/operands[2], *false_computation)); + shape, /*branch_index=*/operands[0], + absl::MakeSpan(*branch_computations), + absl::MakeSpan(operands).subspan(1))); break; } case HloOpcode::kCustomCall: { @@ -2676,35 +2705,27 @@ bool HloParser::ParseAttributeHelper( } case AttrTy::kHloComputation: { HloComputation* result = nullptr; - if (lexer_.GetKind() == TokKind::kLbrace) { - // This means it is a nested computation. - if (!ParseInstructionList(&result, /*computation_name=*/"_")) { - return false; - } - } else { - // This means it is a computation name. - if (!ParseComputationName(&result)) { - return false; - } + if (!ParseHloComputation(&result)) { + return false; } static_cast*>(attr_out_ptr)->emplace(result); return true; } - case AttrTy::kFftType: { - FftType result; - if (!ParseFftType(&result)) { + case AttrTy::kBracedHloComputationList: { + std::vector result; + if (!ParseHloComputationList(&result)) { return false; } - static_cast*>(attr_out_ptr)->emplace(result); + static_cast>*>(attr_out_ptr) + ->emplace(result); return true; } - case AttrTy::kTriangularSolveTranspose: { - TriangularSolveOptions::Transpose result; - if (!ParseTriangularSolveTranspose(&result)) { + case AttrTy::kFftType: { + FftType result; + if (!ParseFftType(&result)) { return false; } - static_cast*>(attr_out_ptr) - ->emplace(result); + static_cast*>(attr_out_ptr)->emplace(result); return true; } case AttrTy::kWindow: { @@ -2859,6 +2880,95 @@ bool HloParser::ParseAttributeHelper( return true; } +// attributes ::= (',' attribute)* +bool HloParser::ParseAttributesAsProtoMessage( + const std::unordered_set& required_attrs, + tensorflow::protobuf::Message* message) { + LocTy loc = lexer_.GetLoc(); + std::unordered_set seen_attrs; + while (EatIfPresent(TokKind::kComma)) { + if (!ParseAttributeAsProtoMessageHelper(message, &seen_attrs)) { + return false; + } + } + // Check that all required attrs were seen. + for (const string& attr : required_attrs) { + if (seen_attrs.find(attr) == seen_attrs.end()) { + return Error(loc, + StrFormat("attribute %s is expected but not seen", attr)); + } + } + return true; +} + +bool HloParser::ParseAttributeAsProtoMessageHelper( + tensorflow::protobuf::Message* message, + std::unordered_set* seen_attrs) { + LocTy loc = lexer_.GetLoc(); + string name; + if (!ParseAttributeName(&name)) { + return Error(loc, "error parsing attributes"); + } + VLOG(1) << "Parsing attribute " << name; + if (!seen_attrs->insert(name).second) { + return Error(loc, StrFormat("attribute %s already exists", name)); + } + const tensorflow::protobuf::Descriptor* descriptor = message->GetDescriptor(); + const tensorflow::protobuf::FieldDescriptor* fd = + descriptor->FindFieldByName(name); + if (!fd) { + string allowed_attrs = "Allowed attributes: "; + + for (int i = 0; i < descriptor->field_count(); ++i) { + if (i == 0) { + absl::StrAppend(&allowed_attrs, descriptor->field(i)->name()); + } else { + absl::StrAppend(&allowed_attrs, ", ", descriptor->field(i)->name()); + } + } + return Error(loc, StrFormat("unexpected attribute \"%s\". %s", name, + allowed_attrs)); + } + const tensorflow::protobuf::Reflection* reflection = message->GetReflection(); + CHECK(!fd->is_repeated()); // Repeated fields not implemented. + bool success = [&] { + switch (fd->type()) { + case tensorflow::protobuf::FieldDescriptor::TYPE_BOOL: { + bool result; + if (!ParseBool(&result)) { + return false; + } + reflection->SetBool(message, fd, result); + return true; + } + case tensorflow::protobuf::FieldDescriptor::TYPE_ENUM: { + if (lexer_.GetKind() != TokKind::kIdent) { + return TokenError( + StrFormat("expects %s type", fd->enum_type()->name())); + } + string val = lexer_.GetStrVal(); + const tensorflow::protobuf::EnumValueDescriptor* evd = + fd->enum_type()->FindValueByName(val); + if (evd == nullptr) { + return TokenError(StrFormat("expects %s type but sees: %s", + fd->enum_type()->name(), val)); + } + reflection->SetEnum(message, fd, evd); + lexer_.Lex(); + return true; + } + default: + LOG(ERROR) << "Unimplemented protocol buffer type " + << fd->DebugString(); + return false; + } + }(); + if (!success) { + return Error(loc, StrFormat("error parsing attribute %s", name)); + } + return true; +} + bool HloParser::ParseComputationName(HloComputation** value) { string name; LocTy loc = lexer_.GetLoc(); @@ -3143,6 +3253,29 @@ bool HloParser::ParsePrecisionList( parse_and_add_item); } +bool HloParser::ParseHloComputation(HloComputation** result) { + if (lexer_.GetKind() == TokKind::kLbrace) { + // This means it is a nested computation. + return ParseInstructionList(result, /*computation_name=*/"_"); + } + // This means it is a computation name. + return ParseComputationName(result); +} + +bool HloParser::ParseHloComputationList(std::vector* result) { + auto parse_and_add_item = [&]() { + HloComputation* computation; + if (!ParseHloComputation(&computation)) { + return false; + } + LOG(INFO) << "parsed computation " << computation->name(); + result->push_back(computation); + return true; + }; + return ParseList(TokKind::kLbrace, TokKind::kRbrace, TokKind::kComma, + parse_and_add_item); +} + // shapelist ::= '{' shapes '}' // precision_elements // ::= /*empty*/ @@ -3623,22 +3756,6 @@ bool HloParser::ParseFftType(FftType* result) { return true; } -bool HloParser::ParseTriangularSolveTranspose( - TriangularSolveOptions::Transpose* result) { - VLOG(1) << "ParseTriangularSolveTranspose"; - if (lexer_.GetKind() != TokKind::kIdent) { - return TokenError("expects triangular solve transpose type"); - } - string val = lexer_.GetStrVal(); - if (!TriangularSolveOptions_Transpose_Parse(val, result) || - !TriangularSolveOptions_Transpose_IsValid(*result)) { - return TokenError( - StrFormat("expects triangular solve transpose type but sees: %s", val)); - } - lexer_.Lex(); - return true; -} - bool HloParser::ParseFusionKind(HloInstruction::FusionKind* result) { VLOG(1) << "ParseFusionKind"; if (lexer_.GetKind() != TokKind::kIdent) { diff --git a/tensorflow/compiler/xla/service/hlo_parser_test.cc b/tensorflow/compiler/xla/service/hlo_parser_test.cc index 8e3f1e44b9562334130aa565ed447a78899fad53..1ba2d718ecc532d5244c077a92f68921f3957f4f 100644 --- a/tensorflow/compiler/xla/service/hlo_parser_test.cc +++ b/tensorflow/compiler/xla/service/hlo_parser_test.cc @@ -575,6 +575,19 @@ ENTRY %Transpose.v3 (input: c128[1,2,3]) -> c128[1,2,3] { ROOT %transpose = c128[1,2,3]{2,1,0} transpose(c128[1,2,3]{2,1,0} %input), dimensions={0,1,2} } +)" +}, +// Triangular solve +{ +"TriangularSolve", +R"(HloModule TriangularSolve_module + +ENTRY %SimpleRightLowerNotranspose.4 (a.1: f32[4,4], b.2: f32[3,4]) -> f32[3,4] { + %a.1 = f32[4,4]{1,0} parameter(0) + %b.2 = f32[3,4]{1,0} parameter(1) + ROOT %triangular-solve.3 = f32[3,4]{1,0} triangular-solve(f32[4,4]{1,0} %a.1, f32[3,4]{1,0} %b.2), lower=true, transpose_a=NO_TRANSPOSE +} + )" }, // Dynamic slice @@ -939,6 +952,36 @@ ENTRY %ParseC128Literal () -> c128[2] { ROOT %c = c128[2]{0} constant({(1, 2), (-inf, nan)}) } +)" +}, +// Indexed Conditional +{ +"IndexedConditional", +R"(HloModule indexed_conditional + +%Negate (x: f32[]) -> f32[] { + %x = f32[] parameter(0) + ROOT %negate = f32[] negate(f32[] %x) +} + +%Identity (y: f32[]) -> f32[] { + %y = f32[] parameter(0) + ROOT %copy = f32[] copy(f32[] %y) +} + +%Floor (z: f32[]) -> f32[] { + %z = f32[] parameter(0) + ROOT %floor = f32[] floor(f32[] %z) +} + +ENTRY %Parameters1.v4 () -> f32[] { + %constant = s32[] constant(1) + %constant.1 = f32[] constant(56) + %constant.2 = f32[] constant(12) + %constant.3 = f32[] constant(13) + ROOT %conditional = f32[] conditional(s32[] %constant, f32[] %constant.1, f32[] %constant.2, f32[] %constant.3), branch_computations={%Negate, %Identity, %Floor} +} + )" }, }); @@ -1178,10 +1221,40 @@ ENTRY Sort { )" }, -// Conditional +// Indexed Conditional +{ +"IndexedConditional", +R"(HloModule indexed_conditional + +Negate { + x = f32[] parameter(0) + ROOT negate = f32[] negate(x) +} + +Identity { + y = f32[] parameter(0) + ROOT copy = f32[] copy(y) +} + +Floor { + z = f32[] parameter(0) + ROOT floor = f32[] floor(z) +} + +ENTRY Parameters1.v4 { + constant = s32[] constant(1) + constant.1 = f32[] constant(56) + constant.2 = f32[] constant(12) + constant.3 = f32[] constant(13) + ROOT conditional = f32[] conditional(constant, constant.1, constant.2, constant.3), branch_computations={Negate, Identity, Floor} +} + +)" +}, +// Predicated Conditional { -"Conditional", -R"(HloModule conditional +"PredicatedConditional", +R"(HloModule pred_conditional Negate { x = f32[] parameter(0) @@ -2304,6 +2377,31 @@ TEST(HloParserSingleOpTest, CanonicalOpWithNested) { text); } +TEST(HloParserSingleOpTest, CanonicalOpIndexedConditionalInlinedBranches) { + const string text = + R"(f32[5,10]{1,0} conditional(s32[], f32[5,10]{1,0}, f32[5,10]{1,0}, f32[5,10]{1,0}), branch_computations={ +{ + tmp_0 = f32[5,10]{1,0} parameter(0) + ROOT tmp_1 = f32[5,10]{1,0} ceil(f32[5,10]{1,0} tmp_0) +}, +{ + tmp_0 = f32[5,10]{1,0} parameter(0) + ROOT tmp_1 = f32[5,10]{1,0} floor(f32[5,10]{1,0} tmp_0) +}, +{ + tmp_0 = f32[5,10]{1,0} parameter(0) + ROOT tmp_1 = f32[5,10]{1,0} copy(f32[5,10]{1,0} tmp_0) +} +})"; + + TF_ASSERT_OK_AND_ASSIGN(auto module, ParseHloString(text)); + const HloComputation* computation = module->entry_computation(); + ASSERT_NE(computation, nullptr); + EXPECT_EQ( + computation->root_instruction()->ToString(HloPrintOptions::Canonical()), + text); +} + TEST(HloParserSingleOpTest, SingleOpWithNested) { const string text = R"(%fusion = f32[3,2,1,1]{3,2,1,0} fusion(f32[3,2,1,1]{3,2,1,0} %p0, f32[2]{0} %p1), kind=kLoop, calls= diff --git a/tensorflow/compiler/xla/service/hlo_runner.cc b/tensorflow/compiler/xla/service/hlo_runner.cc index 5a5401e351384867016a3a9addfd43d57091848c..8f44e1b37ee2e4df53679c226682788279b30aa3 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.cc +++ b/tensorflow/compiler/xla/service/hlo_runner.cc @@ -27,6 +27,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/transfer_manager.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/core/common_runtime/eigen_thread_pool.h" +#include "tensorflow/core/lib/core/blocking_counter.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" @@ -269,8 +270,8 @@ StatusOr HloRunner::ExecuteWithDeviceBuffers( } StatusOr> HloRunner::ExecuteReplicated( - std::unique_ptr module, - const ReplicatedExecuteOptions& options) { + std::unique_ptr module, const ReplicatedExecuteOptions& options, + bool use_threads) { TF_ASSIGN_OR_RETURN( std::unique_ptr executable, CreateExecutable(std::move(module), options.run_hlo_passes)); @@ -369,9 +370,39 @@ StatusOr> HloRunner::ExecuteReplicated( } LOG(INFO) << "Replicated execution started"; - TF_ASSIGN_OR_RETURN(std::vector results, - executable->ExecuteOnStreams(service_run_options, - argument_buffer_slices)); + std::vector results; + if (!use_threads) { + TF_ASSIGN_OR_RETURN(results, + executable->ExecuteOnStreams(service_run_options, + argument_buffer_slices)); + } else { + tensorflow::mutex mutex; + std::vector> thread_results( + options.num_replicas); + { + LOG(INFO) << "Creating thread pool for " << options.num_replicas + << " replicas"; + tensorflow::thread::ThreadPool pool(tensorflow::Env::Default(), + "replicas", options.num_replicas); + for (int64 i = 0; i < options.num_replicas; ++i) { + pool.Schedule([&, i] { + auto result = executable->ExecuteOnStream( + &service_run_options[i], argument_buffer_slices[i], nullptr); + tensorflow::mutex_lock lock(mutex); + thread_results[i] = std::move(result); + }); + } + + // Note: the thread pool destructor guarantees it completes all work + // before we leave this scope. + } + for (auto& thread_result : thread_results) { + if (!thread_result.ok()) { + return thread_result.status(); + } + results.push_back(std::move(thread_result).ValueOrDie()); + } + } LOG(INFO) << "Replicated execution terminated"; std::vector exec_results; diff --git a/tensorflow/compiler/xla/service/hlo_runner.h b/tensorflow/compiler/xla/service/hlo_runner.h index 098989cd4c78fb5ad57cd6700fbf99c50064f225..88a137e6452f487d250c879bbd9ac7a15d07fc59 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.h +++ b/tensorflow/compiler/xla/service/hlo_runner.h @@ -165,9 +165,13 @@ class HloRunner { // Executes a given HLO module into a set of replicas, and returns a map // with the replica number as key, and the corresponding returned literal as // value. + // + // use_threads indicates whether this replicated computation will be executed + // with a thread-per-replica, vs using an implicitly async call such as + // Executable::ExecuteOnStreams. StatusOr> ExecuteReplicated( std::unique_ptr module, - const ReplicatedExecuteOptions& options); + const ReplicatedExecuteOptions& options, bool use_threads = false); // If backend is not created in the constructor, creates and returns the // default backend. If creation fails, crashes the program. diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc index 97e6ea9dad04238c2e6f1a49fba6d880ef3169c2..375ae2c477d7a0aea8445d9c237991eee3353a04 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.cc +++ b/tensorflow/compiler/xla/service/hlo_verifier.cc @@ -58,15 +58,6 @@ bool IsCallerInstruction(HloInstruction* hlo) { } } -Status ShapeVerifier::Preprocess(HloInstruction* hlo) { - if (!hlo->called_computations().empty() && !IsCallerInstruction(hlo)) { - return InternalError( - "Called computations specified for non-caller instruction %s", - hlo->ToString()); - } - return VerifyNotSparse(hlo->shape()); -} - namespace { Status CheckOperandCount(const HloInstruction* hlo, int expected) { @@ -91,6 +82,21 @@ Status CheckParameterCount(const HloInstruction* calling_instruction, } // namespace +Status ShapeVerifier::Preprocess(HloInstruction* hlo) { + if (!hlo->called_computations().empty() && !IsCallerInstruction(hlo)) { + return InternalError( + "Called computations specified for non-caller instruction %s", + hlo->ToString()); + } + TF_RETURN_IF_ERROR(VerifyNotSparse(hlo->shape())); + + absl::optional arity = HloOpcodeArity(hlo->opcode()); + if (arity) { + TF_RETURN_IF_ERROR(CheckOperandCount(hlo, *arity)); + } + return Status::OK(); +} + Status ShapeVerifier::HandleElementwiseUnary(HloInstruction* hlo) { return CheckUnaryShape(hlo); } @@ -122,14 +128,12 @@ Status ShapeVerifier::HandleConcatenate(HloInstruction* concatenate) { } Status ShapeVerifier::HandleConvert(HloInstruction* convert) { - TF_RETURN_IF_ERROR(CheckOperandCount(convert, 1)); return CheckShape(convert, ShapeInference::InferConvertShape( convert->operand(0)->shape(), convert->shape().element_type())); } Status ShapeVerifier::HandleBitcastConvert(HloInstruction* convert) { - TF_RETURN_IF_ERROR(CheckOperandCount(convert, 1)); return CheckShape(convert, ShapeInference::InferBitcastConvertShape( convert->operand(0)->shape(), convert->shape().element_type())); @@ -140,7 +144,6 @@ Status ShapeVerifier::HandleCopy(HloInstruction* copy) { } Status ShapeVerifier::HandleDot(HloInstruction* dot) { - TF_RETURN_IF_ERROR(CheckOperandCount(dot, 2)); TF_ASSIGN_OR_RETURN(const Shape expected, ShapeInference::InferDotOpShape( dot->operand(0)->shape(), dot->operand(1)->shape(), @@ -149,7 +152,6 @@ Status ShapeVerifier::HandleDot(HloInstruction* dot) { } Status ShapeVerifier::HandleConvolution(HloInstruction* convolution) { - TF_RETURN_IF_ERROR(CheckOperandCount(convolution, 2)); TF_ASSIGN_OR_RETURN( const Shape expected, ShapeInference::InferConvolveShape( @@ -160,7 +162,6 @@ Status ShapeVerifier::HandleConvolution(HloInstruction* convolution) { } Status ShapeVerifier::HandleFft(HloInstruction* fft) { - TF_RETURN_IF_ERROR(CheckOperandCount(fft, 1)); TF_ASSIGN_OR_RETURN( const Shape expected, ShapeInference::InferFftShape(fft->operand(0)->shape(), fft->fft_type(), @@ -169,7 +170,6 @@ Status ShapeVerifier::HandleFft(HloInstruction* fft) { } Status ShapeVerifier::HandleTriangularSolve(HloInstruction* hlo) { - TF_RETURN_IF_ERROR(CheckOperandCount(hlo, 2)); TF_ASSIGN_OR_RETURN(const Shape expected, ShapeInference::InferTriangularSolveShape( hlo->operand(0)->shape(), hlo->operand(1)->shape(), @@ -177,6 +177,13 @@ Status ShapeVerifier::HandleTriangularSolve(HloInstruction* hlo) { return CheckShape(hlo, expected); } +Status ShapeVerifier::HandleCholesky(HloInstruction* hlo) { + TF_RETURN_IF_ERROR(CheckOperandCount(hlo, 1)); + TF_ASSIGN_OR_RETURN(const Shape expected, ShapeInference::InferCholeskyShape( + hlo->operand(0)->shape())); + return CheckShape(hlo, expected); +} + Status ShapeVerifier::HandleAllReduce(HloInstruction* crs) { std::vector operand_shapes; for (const HloInstruction* operand : crs->operands()) { @@ -199,13 +206,11 @@ Status ShapeVerifier::HandleReplicaId(HloInstruction* hlo) { } Status ShapeVerifier::HandleCollectivePermute(HloInstruction* hlo) { - TF_RETURN_IF_ERROR(CheckOperandCount(hlo, 1)); return CheckShape(hlo, ShapeInference::InferCollectivePermuteShape( hlo->operand(0)->shape())); } Status ShapeVerifier::HandleReducePrecision(HloInstruction* reduce_precision) { - TF_RETURN_IF_ERROR(CheckOperandCount(reduce_precision, 1)); return CheckShape(reduce_precision, ShapeInference::InferReducePrecisionShape( reduce_precision->operand(0)->shape(), reduce_precision->exponent_bits(), @@ -239,7 +244,6 @@ Status ShapeVerifier::CheckOperandAndParameter( } Status ShapeVerifier::HandleInfeed(HloInstruction* instruction) { - TF_RETURN_IF_ERROR(CheckOperandCount(instruction, 1)); HloInfeedInstruction* infeed = Cast(instruction); TF_RETURN_IF_ERROR(CheckIsTokenOperand(instruction, 0)); @@ -250,7 +254,6 @@ Status ShapeVerifier::HandleInfeed(HloInstruction* instruction) { } Status ShapeVerifier::HandleOutfeed(HloInstruction* instruction) { - TF_RETURN_IF_ERROR(CheckOperandCount(instruction, 2)); HloOutfeedInstruction* outfeed = Cast(instruction); TF_RETURN_IF_ERROR(CheckIsTokenOperand(instruction, 1)); @@ -326,7 +329,6 @@ Status ShapeVerifier::HandleRng(HloInstruction* instruction) { } Status ShapeVerifier::HandleReverse(HloInstruction* reverse) { - TF_RETURN_IF_ERROR(CheckOperandCount(reverse, 1)); return CheckShape( reverse, ShapeInference::InferReverseShape(reverse->operand(0)->shape(), reverse->dimensions())); @@ -387,7 +389,6 @@ Status ShapeVerifier::HandleSort(HloInstruction* sort) { } Status ShapeVerifier::HandleConstant(HloInstruction* constant) { - TF_RETURN_IF_ERROR(CheckOperandCount(constant, 0)); if (!Cast(constant)->HasLiteral()) { return InternalError("Constant is required to have a valid literal: %s", constant->ToString()); @@ -396,7 +397,6 @@ Status ShapeVerifier::HandleConstant(HloInstruction* constant) { } Status ShapeVerifier::HandleIota(HloInstruction* instruction) { - TF_RETURN_IF_ERROR(CheckOperandCount(instruction, 0)); auto* iota = Cast(instruction); if (!iota->shape().IsArray()) { return InternalError("Iota does not support non-array result."); @@ -414,7 +414,6 @@ Status ShapeVerifier::HandleIota(HloInstruction* instruction) { } Status ShapeVerifier::HandleGetTupleElement(HloInstruction* get_tuple_element) { - TF_RETURN_IF_ERROR(CheckOperandCount(get_tuple_element, 1)); return CheckShape(get_tuple_element, ShapeInference::InferGetTupleElementShape( get_tuple_element->operand(0)->shape(), @@ -462,7 +461,6 @@ Status ShapeVerifier::HandleReduce(HloInstruction* reduce) { } Status ShapeVerifier::HandleBitcast(HloInstruction* bitcast) { - TF_RETURN_IF_ERROR(CheckOperandCount(bitcast, 1)); // Bitcasts are not allowed to change the element type. if (bitcast->operand(0)->shape().element_type() != bitcast->shape().element_type()) { @@ -475,7 +473,6 @@ Status ShapeVerifier::HandleBitcast(HloInstruction* bitcast) { } Status ShapeVerifier::HandleBroadcast(HloInstruction* broadcast) { - TF_RETURN_IF_ERROR(CheckOperandCount(broadcast, 1)); // HLO broadcast has no exact analog at the proto level so there is no // ShapeInference method. Check the output shape explicitly. const Shape& operand_shape = broadcast->operand(0)->shape(); @@ -495,7 +492,6 @@ Status ShapeVerifier::HandleBroadcast(HloInstruction* broadcast) { } Status ShapeVerifier::HandleReshape(HloInstruction* reshape) { - TF_RETURN_IF_ERROR(CheckOperandCount(reshape, 1)); // Check for mixed precision. const Shape& operand_shape = reshape->operand(0)->shape(); TF_RET_CHECK(SameElementType(reshape->shape(), operand_shape)); @@ -505,14 +501,12 @@ Status ShapeVerifier::HandleReshape(HloInstruction* reshape) { } Status ShapeVerifier::HandleTranspose(HloInstruction* transpose) { - TF_RETURN_IF_ERROR(CheckOperandCount(transpose, 1)); return CheckShape( transpose, ShapeInference::InferTransposeShape( transpose->operand(0)->shape(), transpose->dimensions())); } Status ShapeVerifier::HandleParameter(HloInstruction* hlo) { - TF_RETURN_IF_ERROR(CheckOperandCount(hlo, 0)); return Status::OK(); } @@ -572,7 +566,6 @@ Status ShapeVerifier::HandleCustomCall(HloInstruction* instruction) { } Status ShapeVerifier::HandleSlice(HloInstruction* slice) { - TF_RETURN_IF_ERROR(CheckOperandCount(slice, 1)); return CheckShape(slice, ShapeInference::InferSliceShape( slice->operand(0)->shape(), slice->slice_starts(), @@ -627,7 +620,6 @@ Status ShapeVerifier::HandleMap(HloInstruction* map) { } Status ShapeVerifier::HandleReduceWindow(HloInstruction* reduce_window) { - TF_RETURN_IF_ERROR(CheckOperandCount(reduce_window, 2)); TF_RETURN_IF_ERROR(CheckShape( reduce_window, ShapeInference::InferReduceWindowShape( @@ -642,7 +634,6 @@ Status ShapeVerifier::HandleReduceWindow(HloInstruction* reduce_window) { } Status ShapeVerifier::HandleSelectAndScatter(HloInstruction* instruction) { - TF_RETURN_IF_ERROR(CheckOperandCount(instruction, 3)); return CheckShape( instruction, ShapeInference::InferSelectAndScatterShape( @@ -653,7 +644,6 @@ Status ShapeVerifier::HandleSelectAndScatter(HloInstruction* instruction) { } Status ShapeVerifier::HandleWhile(HloInstruction* xla_while) { - TF_RETURN_IF_ERROR(CheckOperandCount(xla_while, 1)); TF_RETURN_IF_ERROR( CheckParameterCount(xla_while, xla_while->while_body(), 1)); TF_RETURN_IF_ERROR( @@ -677,33 +667,32 @@ Status ShapeVerifier::HandleWhile(HloInstruction* xla_while) { } Status ShapeVerifier::HandleConditional(HloInstruction* conditional) { - TF_RETURN_IF_ERROR(CheckOperandCount(conditional, 3)); - TF_RETURN_IF_ERROR( - CheckParameterCount(conditional, conditional->true_computation(), 1)); - TF_RETURN_IF_ERROR( - CheckParameterCount(conditional, conditional->false_computation(), 1)); - TF_RETURN_IF_ERROR(CheckOperandAndParameter( - conditional, 1, conditional->true_computation(), 0)); - TF_RETURN_IF_ERROR(CheckOperandAndParameter( - conditional, 2, conditional->false_computation(), 0)); - TF_RETURN_IF_ERROR( - CheckShape(conditional, - conditional->true_computation()->root_instruction()->shape())); - TF_RETURN_IF_ERROR(CheckShape( - conditional, - conditional->false_computation()->root_instruction()->shape())); + const int num_branches = conditional->branch_count(); + if (conditional->operand(0)->shape().element_type() == PRED) { + TF_RET_CHECK(num_branches == 2); + } else { + TF_RET_CHECK(num_branches >= 1); + } + TF_RETURN_IF_ERROR(CheckOperandCount(conditional, num_branches + 1)); + for (int j = 0; j < num_branches; ++j) { + TF_RETURN_IF_ERROR(CheckParameterCount( + conditional, conditional->branch_computation(j), 1)); + TF_RETURN_IF_ERROR(CheckOperandAndParameter( + conditional, j + 1, conditional->branch_computation(j), 0)); + TF_RETURN_IF_ERROR(CheckShape( + conditional, + conditional->branch_computation(j)->root_instruction()->shape())); + } return Status::OK(); } Status ShapeVerifier::HandlePad(HloInstruction* pad) { - TF_RETURN_IF_ERROR(CheckOperandCount(pad, 2)); return CheckShape(pad, ShapeInference::InferPadShape(pad->operand(0)->shape(), pad->operand(1)->shape(), pad->padding_config())); } Status ShapeVerifier::HandleSend(HloInstruction* send) { - TF_RETURN_IF_ERROR(CheckOperandCount(send, 2)); return CheckShape(send, ShapeUtil::MakeTupleShape({send->operand(0)->shape(), ShapeUtil::MakeShape(U32, {}), @@ -711,12 +700,10 @@ Status ShapeVerifier::HandleSend(HloInstruction* send) { } Status ShapeVerifier::HandleSendDone(HloInstruction* send_done) { - TF_RETURN_IF_ERROR(CheckOperandCount(send_done, 1)); return CheckShape(send_done, ShapeUtil::MakeTokenShape()); } Status ShapeVerifier::HandleRecv(HloInstruction* recv) { - TF_RETURN_IF_ERROR(CheckOperandCount(recv, 1)); return CheckShape( recv, ShapeUtil::MakeTupleShape( {ShapeUtil::GetTupleElementShape(recv->shape(), 0), @@ -724,7 +711,6 @@ Status ShapeVerifier::HandleRecv(HloInstruction* recv) { } Status ShapeVerifier::HandleRecvDone(HloInstruction* recv_done) { - TF_RETURN_IF_ERROR(CheckOperandCount(recv_done, 1)); return CheckShape( recv_done, ShapeUtil::MakeTupleShape( @@ -734,7 +720,6 @@ Status ShapeVerifier::HandleRecvDone(HloInstruction* recv_done) { Status ShapeVerifier::HandleBatchNormTraining( HloInstruction* batch_norm_training) { - TF_RETURN_IF_ERROR(CheckOperandCount(batch_norm_training, 3)); return CheckShape(batch_norm_training, ShapeInference::InferBatchNormTrainingShape( batch_norm_training->operand(0)->shape(), @@ -745,7 +730,6 @@ Status ShapeVerifier::HandleBatchNormTraining( Status ShapeVerifier::HandleBatchNormInference( HloInstruction* batch_norm_inference) { - TF_RETURN_IF_ERROR(CheckOperandCount(batch_norm_inference, 5)); return CheckShape(batch_norm_inference, ShapeInference::InferBatchNormInferenceShape( batch_norm_inference->operand(0)->shape(), @@ -757,7 +741,6 @@ Status ShapeVerifier::HandleBatchNormInference( } Status ShapeVerifier::HandleBatchNormGrad(HloInstruction* batch_norm_grad) { - TF_RETURN_IF_ERROR(CheckOperandCount(batch_norm_grad, 5)); return CheckShape(batch_norm_grad, ShapeInference::InferBatchNormGradShape( batch_norm_grad->operand(0)->shape(), batch_norm_grad->operand(1)->shape(), @@ -825,7 +808,6 @@ Status CheckMixedPrecisionOperands(const HloInstruction* instruction) { } // namespace Status ShapeVerifier::HandleGather(HloInstruction* gather) { - TF_RETURN_IF_ERROR(CheckOperandCount(gather, 2)); return CheckShape( gather, ShapeInference::InferGatherShape( @@ -834,7 +816,6 @@ Status ShapeVerifier::HandleGather(HloInstruction* gather) { } Status ShapeVerifier::HandleScatter(HloInstruction* scatter) { - TF_RETURN_IF_ERROR(CheckOperandCount(scatter, 3)); return CheckShape( scatter, ShapeInference::InferScatterShape( scatter->operand(0)->shape(), scatter->operand(1)->shape(), @@ -852,7 +833,6 @@ Status ShapeVerifier::HandleAfterAll(HloInstruction* token) { } Status ShapeVerifier::HandleAddDependency(HloInstruction* add_dependency) { - TF_RETURN_IF_ERROR(CheckOperandCount(add_dependency, 2)); TF_RETURN_IF_ERROR(CheckIsTokenOperand(add_dependency, 1)); return CheckShape(add_dependency, add_dependency->operand(0)->shape()); } @@ -934,14 +914,12 @@ Status ShapeVerifier::CheckShape(const HloInstruction* instruction, } Status ShapeVerifier::CheckUnaryShape(const HloInstruction* instruction) { - TF_RETURN_IF_ERROR(CheckOperandCount(instruction, 1)); return CheckShape(instruction, ShapeInference::InferUnaryOpShape(instruction->opcode(), instruction->operand(0))); } Status ShapeVerifier::CheckBinaryShape(const HloInstruction* instruction) { - TF_RETURN_IF_ERROR(CheckOperandCount(instruction, 2)); return CheckShape( instruction, ShapeInference::InferBinaryOpShape(instruction->opcode(), instruction->operand(0), @@ -949,7 +927,6 @@ Status ShapeVerifier::CheckBinaryShape(const HloInstruction* instruction) { } Status ShapeVerifier::CheckTernaryShape(const HloInstruction* instruction) { - TF_RETURN_IF_ERROR(CheckOperandCount(instruction, 3)); return CheckShape(instruction, ShapeInference::InferTernaryOpShape( instruction->opcode(), instruction->operand(0), @@ -1395,17 +1372,13 @@ class InstructionVerifier : public DfsHloVisitorWithDefault { } Status HandleConditional(HloInstruction* conditional) override { - if (conditional->true_computation()->num_parameters() != 1) { - return FailedPrecondition( - "True computation %s of %s must have 1 parameter insted of %d", - conditional->true_computation()->name(), conditional->ToString(), - conditional->true_computation()->num_parameters()); - } - if (conditional->false_computation()->num_parameters() != 1) { - return FailedPrecondition( - "False computation %s of %s must have 1 parameter insted of %d", - conditional->false_computation()->name(), conditional->ToString(), - conditional->false_computation()->num_parameters()); + for (int b = 0; b < conditional->branch_count(); ++b) { + if (conditional->branch_computation(b)->num_parameters() != 1) { + return FailedPrecondition( + "Branch computation %s of %s must have 1 parameter insted of %d", + conditional->branch_computation(b)->name(), conditional->ToString(), + conditional->branch_computation(b)->num_parameters()); + } } return Status::OK(); } diff --git a/tensorflow/compiler/xla/service/hlo_verifier.h b/tensorflow/compiler/xla/service/hlo_verifier.h index a9b5e9a3e6eec19e125188a192694fcaadfe2322..d427a1586c3cd1d1abbd6606f33067e36cabad98 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.h +++ b/tensorflow/compiler/xla/service/hlo_verifier.h @@ -52,6 +52,7 @@ class ShapeVerifier : public DfsHloVisitor { Status HandleDot(HloInstruction* dot) override; Status HandleConvolution(HloInstruction* convolution) override; Status HandleFft(HloInstruction* fft) override; + Status HandleCholesky(HloInstruction* hlo) override; Status HandleTriangularSolve(HloInstruction* hlo) override; Status HandleAllReduce(HloInstruction* crs) override; Status HandleAllToAll(HloInstruction* hlo) override; diff --git a/tensorflow/compiler/xla/service/instruction_fusion.cc b/tensorflow/compiler/xla/service/instruction_fusion.cc index a5767774f2dcebf8fd309823d48dcc3269b3f594..e4a78af7c72fa53b1739f6ce4ea02a86ced0fc68 100644 --- a/tensorflow/compiler/xla/service/instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/instruction_fusion.cc @@ -126,6 +126,7 @@ bool IsAlwaysDuplicable(const HloInstruction& instruction) { case HloOpcode::kBatchNormInference: case HloOpcode::kBatchNormTraining: case HloOpcode::kCall: + case HloOpcode::kCholesky: case HloOpcode::kConditional: case HloOpcode::kConvolution: case HloOpcode::kAllReduce: @@ -151,11 +152,13 @@ bool IsAlwaysDuplicable(const HloInstruction& instruction) { case HloOpcode::kReduceWindow: case HloOpcode::kRemainder: case HloOpcode::kRng: + case HloOpcode::kRsqrt: case HloOpcode::kScatter: case HloOpcode::kSelectAndScatter: case HloOpcode::kSend: case HloOpcode::kSendDone: case HloOpcode::kSort: + case HloOpcode::kSqrt: case HloOpcode::kTanh: case HloOpcode::kTrace: case HloOpcode::kTriangularSolve: diff --git a/tensorflow/compiler/xla/service/interpreter/BUILD b/tensorflow/compiler/xla/service/interpreter/BUILD index 8cd936268994c2a25c2c0debe0a003d1d05cbd0b..599489b3785be50ba7a145f298a13d6bb995a1cf 100644 --- a/tensorflow/compiler/xla/service/interpreter/BUILD +++ b/tensorflow/compiler/xla/service/interpreter/BUILD @@ -32,6 +32,7 @@ cc_library( "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla/service:algebraic_simplifier", + "//tensorflow/compiler/xla/service:cholesky_expander", "//tensorflow/compiler/xla/service:compiler", "//tensorflow/compiler/xla/service:computation_placer", "//tensorflow/compiler/xla/service:dynamic_index_splitter", diff --git a/tensorflow/compiler/xla/service/interpreter/compiler.cc b/tensorflow/compiler/xla/service/interpreter/compiler.cc index 792773c676984aa280c1b20cb7fd0fc7c9425f6c..a8f8ab4f725d904a529dbd50c1c199972a1c0895 100644 --- a/tensorflow/compiler/xla/service/interpreter/compiler.cc +++ b/tensorflow/compiler/xla/service/interpreter/compiler.cc @@ -20,6 +20,7 @@ limitations under the License. #include "absl/memory/memory.h" #include "tensorflow/compiler/xla/service/algebraic_simplifier.h" +#include "tensorflow/compiler/xla/service/cholesky_expander.h" #include "tensorflow/compiler/xla/service/computation_placer.h" #include "tensorflow/compiler/xla/service/cpu/custom_call_target_registry.h" #include "tensorflow/compiler/xla/service/dynamic_index_splitter.h" @@ -80,6 +81,7 @@ Status InterpreterCompiler::RunHloOptimization(HloModule* hlo_module) { HloPassPipeline pipeline("Interpreter"); pipeline.AddPass(); + pipeline.AddPass(); pipeline.AddPass(); pipeline.AddPass( hlo_module->mutable_entry_computation_layout(), diff --git a/tensorflow/compiler/xla/service/layout_assignment.cc b/tensorflow/compiler/xla/service/layout_assignment.cc index 2dd9e055503e01f5c1ace5e6cd3dc64012828b71..e90bdb640d71296243e35ce19b412e3fcd83c78f 100644 --- a/tensorflow/compiler/xla/service/layout_assignment.cc +++ b/tensorflow/compiler/xla/service/layout_assignment.cc @@ -588,48 +588,40 @@ Status LayoutAssignment::AddMandatoryConstraints( TF_RETURN_IF_ERROR(constraints->SetOperandLayout( body_layout.result_shape(), instruction, 0)); } else if (instruction->opcode() == HloOpcode::kConditional) { - // The layout of the true and false computations must match, and must + // The layout of the branch computations must match, and must // be the layout of the kConditional instruction. - TF_RET_CHECK(instruction->operand_count() == 3); - - HloComputation* true_computation = instruction->true_computation(); - HloComputation* false_computation = instruction->false_computation(); - const HloInstruction* true_operand = instruction->operand(1); - const HloInstruction* false_operand = instruction->operand(2); - - TF_RET_CHECK(true_computation->num_parameters() == 1); - TF_RET_CHECK(false_computation->num_parameters() == 1); - ComputationLayout& true_computation_layout = - FindOrDie(computation_layouts_, true_computation); - ComputationLayout& false_computation_layout = - FindOrDie(computation_layouts_, false_computation); - - DCHECK(ShapeUtil::Compatible(true_operand->shape(), - true_computation_layout.parameter_shape(0))); - DCHECK(ShapeUtil::Compatible( - false_operand->shape(), false_computation_layout.parameter_shape(0))); - if (true_computation_layout.result_layout() != - false_computation_layout.result_layout()) { - // We assign layouts in DFS fashion, so the true and false computations - // might have negotiated a different layout. But for the conditional - // instruction POV the layout must match, so we run again on the false - // computation, this time with proper computation layout. - VLOG(2) << "Reset %conditional false computation result layout: " - "false_computation=" - << false_computation->name() - << " conditional=" << instruction->name() << " shape=" - << true_computation_layout.result_layout().ToString(); - *false_computation_layout.mutable_result_layout() = - true_computation_layout.result_layout(); + ComputationLayout& branch0_computation_layout = + FindOrDie(computation_layouts_, instruction->branch_computation(0)); + for (int j = 0; j < instruction->branch_count(); ++j) { + TF_RET_CHECK(instruction->branch_computation(j)->num_parameters() == 1); + ComputationLayout& branch_computation_layout = + FindOrDie(computation_layouts_, instruction->branch_computation(j)); + + DCHECK(ShapeUtil::Compatible( + instruction->operand(j + 1)->shape(), + branch_computation_layout.parameter_shape(0))); + if (branch0_computation_layout.result_layout() != + branch_computation_layout.result_layout()) { + // We assign layouts in DFS fashion, so the br_0 and br_j + // computations might have negotiated a different layout. But for the + // case instruction POV the layout must match, so we run again + // on the br_j computation, this time with proper computation layout. + VLOG(2) << "Reset %conditional branch " << j + << " computation result layout: branch_computation=" + << instruction->branch_computation(j)->name() + << " case=" << instruction->name() << " shape=" + << branch0_computation_layout.result_layout().ToString(); + *branch_computation_layout.mutable_result_layout() = + branch0_computation_layout.result_layout(); + } + if (j == 0) { + TF_RETURN_IF_ERROR(constraints->SetInstructionLayout( + branch0_computation_layout.result_shape(), instruction)); + } + TF_RETURN_IF_ERROR(constraints->SetOperandLayout( + branch_computation_layout.parameter_shape(0), instruction, j + 1, + /*mandatory=*/true)); } - TF_RETURN_IF_ERROR(constraints->SetInstructionLayout( - true_computation_layout.result_shape(), instruction)); - TF_RETURN_IF_ERROR(constraints->SetOperandLayout( - true_computation_layout.parameter_shape(0), instruction, 1, - /*mandatory=*/true)); - TF_RETURN_IF_ERROR(constraints->SetOperandLayout( - false_computation_layout.parameter_shape(0), instruction, 2, - /*mandatory=*/true)); } } // Finally set the result layout to match ComputationLayout, if there is one. @@ -699,28 +691,21 @@ Status CheckWhileLayout(HloInstruction* while_inst, Status CheckConditionalLayout( HloInstruction* instruction, - const ComputationLayout& true_computation_layout, - const ComputationLayout& false_computation_layout) { - HloComputation* true_computation = instruction->true_computation(); - HloComputation* false_computation = instruction->false_computation(); - const HloInstruction* true_operand = instruction->operand(1); - const HloInstruction* false_operand = instruction->operand(2); - - TF_RET_CHECK(true_computation_layout.result_layout() == - false_computation_layout.result_layout()); - TF_RET_CHECK(true_computation_layout.result_layout().MatchesLayoutInShape( - instruction->shape())); - TF_RET_CHECK(true_computation_layout.result_layout().MatchesLayoutInShape( - true_computation->root_instruction()->shape())); - TF_RET_CHECK(false_computation_layout.result_layout().MatchesLayoutInShape( - instruction->shape())); - TF_RET_CHECK(false_computation_layout.result_layout().MatchesLayoutInShape( - false_computation->root_instruction()->shape())); - TF_RET_CHECK(true_computation_layout.parameter_layout(0).MatchesLayoutInShape( - true_operand->shape())); - TF_RET_CHECK( - false_computation_layout.parameter_layout(0).MatchesLayoutInShape( - false_operand->shape())); + absl::Span branch_computation_layouts) { + for (int j = 0; j < instruction->branch_count(); ++j) { + const HloInstruction* branch_operand = instruction->operand(j + 1); + TF_RET_CHECK(branch_computation_layouts[0].result_layout() == + branch_computation_layouts[j].result_layout()); + TF_RET_CHECK( + branch_computation_layouts[j].result_layout().MatchesLayoutInShape( + instruction->shape())); + TF_RET_CHECK( + branch_computation_layouts[j].result_layout().MatchesLayoutInShape( + instruction->branch_computation(j)->root_instruction()->shape())); + TF_RET_CHECK( + branch_computation_layouts[j].parameter_layout(0).MatchesLayoutInShape( + branch_operand->shape())); + } return Status::OK(); } @@ -937,13 +922,16 @@ Status LayoutAssignment::CheckLayouts(HloModule* module) { FindOrDie(computation_layouts_, instruction->while_condition()), FindOrDie(computation_layouts_, instruction->while_body()))); break; - case HloOpcode::kConditional: + case HloOpcode::kConditional: { + std::vector branch_computation_layouts; + for (auto branch_computation : instruction->branch_computations()) { + branch_computation_layouts.emplace_back( + FindOrDie(computation_layouts_, branch_computation)); + } TF_RETURN_IF_ERROR(CheckConditionalLayout( - instruction, - FindOrDie(computation_layouts_, instruction->true_computation()), - FindOrDie(computation_layouts_, - instruction->false_computation()))); + instruction, absl::MakeSpan(branch_computation_layouts))); break; + } default: break; } @@ -1019,16 +1007,6 @@ std::unique_ptr LayoutAssignment::ChooseOperandLayoutFromOutputLayout( Shape operand_shape = operand->shape(); *operand_shape.mutable_layout() = LayoutUtil::GetDefaultLayoutForShape(operand_shape); - if (ShapeUtil::ReshapeIsBitcast(operand_shape, output_shape_with_layout)) { - return absl::make_unique(operand_shape.layout()); - } - if (operand_shape.rank() == output_shape.rank()) { - *operand_shape.mutable_layout() = output_layout; - if (ShapeUtil::ReshapeIsBitcast(operand_shape, - output_shape_with_layout)) { - return absl::make_unique(output_layout); - } - } auto aligned_operand_shape = ShapeUtil::AlignLayouts(output_shape_with_layout, operand_shape); if (aligned_operand_shape) { @@ -1090,16 +1068,6 @@ std::unique_ptr LayoutAssignment::ChooseOutputLayoutFromOperandLayout( Shape output_shape = user->shape(); *output_shape.mutable_layout() = LayoutUtil::GetDefaultLayoutForShape(output_shape); - if (ShapeUtil::ReshapeIsBitcast(output_shape, operand_shape_with_layout)) { - return absl::make_unique(output_shape.layout()); - } - if (operand->shape().rank() == output_shape.rank()) { - *output_shape.mutable_layout() = operand_layout; - if (ShapeUtil::ReshapeIsBitcast(output_shape, - operand_shape_with_layout)) { - return absl::make_unique(operand_layout); - } - } auto aligned_user_shape = ShapeUtil::AlignLayouts(operand_shape_with_layout, output_shape); if (aligned_user_shape) { @@ -2058,6 +2026,7 @@ bool LayoutAssignment::InstructionCanChangeLayout( case HloOpcode::kRemainder: case HloOpcode::kReverse: case HloOpcode::kRoundNearestAfz: + case HloOpcode::kRsqrt: case HloOpcode::kScatter: case HloOpcode::kSelect: case HloOpcode::kSelectAndScatter: @@ -2068,9 +2037,11 @@ bool LayoutAssignment::InstructionCanChangeLayout( case HloOpcode::kSin: case HloOpcode::kSlice: case HloOpcode::kSort: + case HloOpcode::kSqrt: case HloOpcode::kSubtract: case HloOpcode::kTanh: case HloOpcode::kTriangularSolve: + case HloOpcode::kCholesky: case HloOpcode::kTupleSelect: case HloOpcode::kWhile: return false; diff --git a/tensorflow/compiler/xla/service/pattern_matcher.h b/tensorflow/compiler/xla/service/pattern_matcher.h index 9e3d1060210790f60243195a1c1dff13f1fc7fc5..7164bfc4cd48ea945519dadece92d8df2e88d02a 100644 --- a/tensorflow/compiler/xla/service/pattern_matcher.h +++ b/tensorflow/compiler/xla/service/pattern_matcher.h @@ -2053,10 +2053,12 @@ XLA_UNOP_PATTERN(RecvDone) XLA_UNOP_PATTERN(ReducePrecision) XLA_UNOP_PATTERN(Reshape) XLA_UNOP_PATTERN(Reverse) +XLA_UNOP_PATTERN(Rsqrt) XLA_UNOP_PATTERN(SendDone) XLA_UNOP_PATTERN(Sign) XLA_UNOP_PATTERN(Sin) XLA_UNOP_PATTERN(Slice) +XLA_UNOP_PATTERN(Sqrt) XLA_UNOP_PATTERN(Tanh) XLA_UNOP_PATTERN(Transpose) #undef XLA_UNOP_PATTERN diff --git a/tensorflow/compiler/xla/service/shape_inference.cc b/tensorflow/compiler/xla/service/shape_inference.cc index 3f4456c1bbf0f620609459256424b9cb30a04e13..20d41592e61b23000b66aa54fe56421838ad0dbc 100644 --- a/tensorflow/compiler/xla/service/shape_inference.cc +++ b/tensorflow/compiler/xla/service/shape_inference.cc @@ -262,6 +262,8 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, case HloOpcode::kExpm1: case HloOpcode::kLog: case HloOpcode::kLog1p: + case HloOpcode::kRsqrt: + case HloOpcode::kSqrt: case HloOpcode::kTanh: if (!ShapeUtil::ElementIsFloating(shape) && !ShapeUtil::ElementIsComplex(shape)) { @@ -1289,16 +1291,12 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, TF_RETURN_IF_ERROR( ExpectArray(scale_shape, "scale input of batch norm inference")); - TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(operand_shape) == - Status::OK()); - TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(offset_shape) == - Status::OK()); - TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(scale_shape) == - Status::OK()); - TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(mean_shape) == - Status::OK()); - TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(variance_shape) == - Status::OK()); + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(operand_shape)); + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(offset_shape)); + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(scale_shape)); + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(mean_shape)); + TF_RETURN_IF_ERROR( + ShapeUtil::ValidateShapeWithOptionalLayout(variance_shape)); if (feature_index >= operand_shape.rank()) { return InvalidArgument( @@ -1909,6 +1907,14 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, /* static */ StatusOr ShapeInference::InferTriangularSolveShape( const Shape& a, const Shape& b, const TriangularSolveOptions& options) { + if ((!ShapeUtil::ElementIsFloating(a) && !ShapeUtil::ElementIsComplex(a)) || + a.element_type() != b.element_type()) { + return InvalidArgument( + "Expected element types in shape to be floating or complex and " + "identical for TriangularSolve; got %s and %s.", + PrimitiveType_Name(a.element_type()), + PrimitiveType_Name(b.element_type())); + } if (a.rank() < 2) { return InvalidArgument( "The 'a' argument to TriangularSolve must have rank >= 2, got shape %s", @@ -1950,6 +1956,27 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, return b; } +/* static */ StatusOr ShapeInference::InferCholeskyShape( + const Shape& a) { + if (!ShapeUtil::ElementIsFloating(a) && !ShapeUtil::ElementIsComplex(a)) { + return InvalidArgument( + "Expected element type in shape to be floating or complex for " + "Cholesky; got %s.", + PrimitiveType_Name(a.element_type())); + } + if (a.rank() < 2) { + return InvalidArgument( + "The 'a' argument to Cholesky must have rank >= 2, got shape %s", + a.ToString()); + } + if (a.dimensions(a.rank() - 2) != a.dimensions(a.rank() - 1)) { + return InvalidArgument( + "The two minor dimensions of 'a' must have equal size, got %s.", + a.ToString()); + } + return a; +} + /* static */ StatusOr ShapeInference::InferAllReduceShape( absl::Span operand_shapes) { for (const Shape* operand_shape : operand_shapes) { @@ -2531,59 +2558,55 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, } /* static */ StatusOr ShapeInference::InferConditionalShape( - const Shape& predicate, const Shape& true_operand, - const Shape& false_operand, const ProgramShape& true_computation, - const ProgramShape& false_computation) { - if (!ShapeUtil::Equal(predicate, ShapeUtil::MakeShape(PRED, {}))) { - return InvalidArgument("Predicate must be a boolean; got %s.", - ShapeUtil::HumanString(predicate)); - } - - if (true_computation.parameters_size() != 1) { - return InvalidArgument("true_computation must take 1 argument; got %d.", - true_computation.parameters_size()); - } - if (!ShapeUtil::Compatible(true_computation.parameters(0), true_operand)) { - auto true_shape_string = [&]() { - return StrFormat("true_operand: %s; true_computation: %s", - ShapeUtil::HumanString(true_operand), - ShapeUtil::HumanString(true_computation)); - }; - return InvalidArgument( - "true_operand must match the shape of the only parameter of " - "true_computation: got %s.", - true_shape_string()); + const Shape& branch_index, + absl::Span branch_computations, + absl::Span branch_operands) { + if (!ShapeUtil::Equal(branch_index, ShapeUtil::MakeShape(PRED, {})) && + !ShapeUtil::Equal(branch_index, ShapeUtil::MakeShape(S32, {}))) { + return InvalidArgument("branch_index must be bool or int32; got %s.", + ShapeUtil::HumanString(branch_index)); + } + if (branch_index.element_type() == PRED) { + TF_RET_CHECK(2 == branch_computations.size()); + } else { + TF_RET_CHECK(!branch_computations.empty()); } + TF_RET_CHECK(branch_computations.size() == branch_operands.size()); - if (false_computation.parameters_size() != 1) { - return InvalidArgument("false_computation must take 1 argument; got %d.", - false_computation.parameters_size()); - } - if (!ShapeUtil::Compatible(false_computation.parameters(0), false_operand)) { - auto false_shape_string = [&]() { - return StrFormat("false_operand: %s; false_computation: %s", - ShapeUtil::HumanString(false_operand), - ShapeUtil::HumanString(false_computation)); - }; - return InvalidArgument( - "false_operand must match the shape of the only parameter of " - "false_computation: got %s.", - false_shape_string()); - } - if (!ShapeUtil::Compatible(true_computation.result(), - false_computation.result())) { - auto shape_string = [&]() { - return StrFormat( - "true_computation result: %s; false_computation result: %s.", - ShapeUtil::HumanString(true_computation.result()), - ShapeUtil::HumanString(false_computation.result())); - }; - return InvalidArgument( - "the result of true_computation and false_computation must have the " - "same shape: got %s.", - shape_string()); + for (int j = 0; j < branch_computations.size(); ++j) { + if (branch_computations[j].parameters_size() != 1) { + return InvalidArgument( + "branch computation %d must take 1 argument; got %d.", j, + branch_computations[j].parameters_size()); + } + if (!ShapeUtil::Compatible(branch_computations[j].parameters(0), + branch_operands[j])) { + auto shape_string = [&]() { + return StrFormat("operand: %s; computation: %s", + ShapeUtil::HumanString(branch_operands[j]), + ShapeUtil::HumanString(branch_computations[j])); + }; + return InvalidArgument( + "branch operand %d must match the shape of the only parameter of " + "branch computation %d: got %s.", + j, j, shape_string()); + } + + if (!ShapeUtil::Compatible(branch_computations[0].result(), + branch_computations[j].result())) { + auto shape_string = [&]() { + return StrFormat( + "branch 0 computation result: %s; branch %d computation result: %s", + ShapeUtil::HumanString(branch_computations[0].result()), j, + ShapeUtil::HumanString(branch_computations[j].result())); + }; + return InvalidArgument( + "the result of branch 0 computation and branch %d computation must " + "have the same shape: got %s.", + j, shape_string()); + } } - return true_computation.result(); + return branch_computations[0].result(); } /* static */ StatusOr ShapeInference::InferBroadcastShape( diff --git a/tensorflow/compiler/xla/service/shape_inference.h b/tensorflow/compiler/xla/service/shape_inference.h index acb071ab18824472153fc608b812ad2d9c52651e..590a664224e6786bf387494139c66a69a43a5247 100644 --- a/tensorflow/compiler/xla/service/shape_inference.h +++ b/tensorflow/compiler/xla/service/shape_inference.h @@ -120,6 +120,9 @@ class ShapeInference { static StatusOr InferTriangularSolveShape( const Shape& a, const Shape& b, const TriangularSolveOptions& options); + // Infers the shape produced by the given triangular solve operation. + static StatusOr InferCholeskyShape(const Shape& a); + // Infers the shape produced by a cross replica sum with the given operand // shapes. static StatusOr InferAllReduceShape( @@ -205,11 +208,11 @@ class ShapeInference { const ProgramShape& body, const Shape& init); - // Infers the shape produced by a conditional operation. + // Infers the shape produced by a predicated or indexed conditional operation. static StatusOr InferConditionalShape( - const Shape& predicate, const Shape& true_operand, - const Shape& false_operand, const ProgramShape& true_computation, - const ProgramShape& false_computation); + const Shape& branch_index, + absl::Span branch_computations, + absl::Span branch_operands); // Infers the shape produced by a broadcast operation. static StatusOr InferBroadcastShape( diff --git a/tensorflow/compiler/xla/service/shape_inference_test.cc b/tensorflow/compiler/xla/service/shape_inference_test.cc index f400ef51f07b006eef2ea674feff1dd72f836e77..6f8cc6136bb7f13afc6332fdf52211772a9ad040 100644 --- a/tensorflow/compiler/xla/service/shape_inference_test.cc +++ b/tensorflow/compiler/xla/service/shape_inference_test.cc @@ -1582,79 +1582,166 @@ TEST_F(ShapeInferenceTest, Rank1Transpose) { ShapeUtil::Compatible(inferred_shape, ShapeUtil::MakeShape(F32, {5}))); } -TEST_F(ShapeInferenceTest, Conditional) { +TEST_F(ShapeInferenceTest, ConditionalPred) { auto inferred_status0 = ShapeInference::InferConditionalShape( - pred_, vector_32_, vector_64_, - ShapeUtil::MakeProgramShape({vector_32_}, f32_), - ShapeUtil::MakeProgramShape({vector_64_}, f32_)); + pred_, + {ShapeUtil::MakeProgramShape({vector_32_}, f32_), + ShapeUtil::MakeProgramShape({vector_64_}, f32_)}, + {vector_32_, vector_64_}); EXPECT_IS_OK(inferred_status0.status()); EXPECT_TRUE(ShapeUtil::Equal(f32_, inferred_status0.ValueOrDie())); auto inferred_status1 = ShapeInference::InferConditionalShape( - pred_, matrix_32_48_, vector_32_, - ShapeUtil::MakeProgramShape({matrix_32_48_}, vector_64_), - ShapeUtil::MakeProgramShape({vector_32_}, vector_64_)); + pred_, + {ShapeUtil::MakeProgramShape({matrix_32_48_}, vector_64_), + ShapeUtil::MakeProgramShape({vector_32_}, vector_64_)}, + {matrix_32_48_, vector_32_}); EXPECT_IS_OK(inferred_status1.status()); EXPECT_TRUE(ShapeUtil::Equal(vector_64_, inferred_status1.ValueOrDie())); auto tuple_f32_v32 = ShapeUtil::MakeTupleShape({f32_, vector_32_}); auto inferred_status2 = ShapeInference::InferConditionalShape( - pred_, matrix_32_48_, tuple_f32_v32, - ShapeUtil::MakeProgramShape({matrix_32_48_}, vector_32_), - ShapeUtil::MakeProgramShape({tuple_f32_v32}, vector_32_)); + pred_, + {ShapeUtil::MakeProgramShape({matrix_32_48_}, vector_32_), + ShapeUtil::MakeProgramShape({tuple_f32_v32}, vector_32_)}, + {matrix_32_48_, tuple_f32_v32}); EXPECT_IS_OK(inferred_status2.status()); EXPECT_TRUE(ShapeUtil::Equal(vector_32_, inferred_status2.ValueOrDie())); auto inferred_status_error0 = ShapeInference::InferConditionalShape( - s32_, vector_32_, vector_64_, - ShapeUtil::MakeProgramShape({vector_32_}, f32_), - ShapeUtil::MakeProgramShape({vector_64_}, f32_)); + f32_, + {ShapeUtil::MakeProgramShape({vector_32_}, f32_), + ShapeUtil::MakeProgramShape({vector_64_}, f32_)}, + {vector_32_, vector_64_}); EXPECT_FALSE(inferred_status_error0.ok()); EXPECT_THAT(inferred_status_error0.status().error_message(), - HasSubstr("Predicate must be a boolean")); + HasSubstr("must be bool or int32")); auto inferred_status_error1 = ShapeInference::InferConditionalShape( - pred_, ShapeUtil::MakeTupleShape({f32_, vector_32_}), matrix_32_48_, - ShapeUtil::MakeProgramShape({f32_, vector_32_}, vector_32_), - ShapeUtil::MakeProgramShape({matrix_32_48_}, vector_32_)); + pred_, + {ShapeUtil::MakeProgramShape({f32_, vector_32_}, vector_32_), + ShapeUtil::MakeProgramShape({matrix_32_48_}, vector_32_)}, + {ShapeUtil::MakeTupleShape({f32_, vector_32_}), matrix_32_48_}); EXPECT_FALSE(inferred_status_error1.ok()); EXPECT_THAT(inferred_status_error1.status().error_message(), - HasSubstr("true_computation must take 1 argument")); + HasSubstr("branch computation 0 must take 1 argument")); auto inferred_status_error2 = ShapeInference::InferConditionalShape( - pred_, vector_32_, vector_64_, - ShapeUtil::MakeProgramShape({vector_64_}, f32_), - ShapeUtil::MakeProgramShape({vector_64_}, f32_)); + pred_, + {ShapeUtil::MakeProgramShape({vector_64_}, f32_), + ShapeUtil::MakeProgramShape({vector_64_}, f32_)}, + {vector_32_, vector_64_}); EXPECT_FALSE(inferred_status_error2.ok()); EXPECT_THAT(inferred_status_error2.status().error_message(), - HasSubstr("true_operand must match the shape of the only " - "parameter of true_computation")); + HasSubstr("branch operand 0 must match the shape of the only " + "parameter of branch computation 0")); auto inferred_status_error3 = ShapeInference::InferConditionalShape( - pred_, matrix_32_48_, ShapeUtil::MakeTupleShape({f32_, vector_32_}), - ShapeUtil::MakeProgramShape({matrix_32_48_}, vector_32_), - ShapeUtil::MakeProgramShape({f32_, vector_32_}, vector_32_)); + pred_, + {ShapeUtil::MakeProgramShape({matrix_32_48_}, vector_32_), + ShapeUtil::MakeProgramShape({f32_, vector_32_}, vector_32_)}, + {matrix_32_48_, ShapeUtil::MakeTupleShape({f32_, vector_32_})}); EXPECT_FALSE(inferred_status_error3.ok()); EXPECT_THAT(inferred_status_error3.status().error_message(), - HasSubstr("false_computation must take 1 argument")); + HasSubstr("branch computation 1 must take 1 argument")); auto inferred_status_error4 = ShapeInference::InferConditionalShape( - pred_, vector_32_, vector_64_, - ShapeUtil::MakeProgramShape({vector_32_}, f32_), - ShapeUtil::MakeProgramShape({vector_32_}, f32_)); + pred_, + {ShapeUtil::MakeProgramShape({vector_32_}, f32_), + ShapeUtil::MakeProgramShape({vector_32_}, f32_)}, + {vector_32_, vector_64_}); EXPECT_FALSE(inferred_status_error4.ok()); EXPECT_THAT(inferred_status_error4.status().error_message(), - HasSubstr("false_operand must match the shape of the only " - "parameter of false_computation")); + HasSubstr("branch operand 1 must match the shape of the only " + "parameter of branch computation 1")); auto inferred_status_error5 = ShapeInference::InferConditionalShape( - pred_, vector_32_, vector_64_, - ShapeUtil::MakeProgramShape({vector_32_}, f32_), - ShapeUtil::MakeProgramShape({vector_64_}, vector_32_)); + pred_, + {ShapeUtil::MakeProgramShape({vector_32_}, f32_), + ShapeUtil::MakeProgramShape({vector_64_}, vector_32_)}, + {vector_32_, vector_64_}); EXPECT_FALSE(inferred_status_error5.ok()); EXPECT_THAT(inferred_status_error5.status().error_message(), - HasSubstr("the result of true_computation and false_computation " - "must have the same shape")); + HasSubstr("the result of branch 0 computation and branch 1 " + "computation must have the same shape")); +} + +TEST_F(ShapeInferenceTest, ConditionalIndexed) { + auto r0s32 = ShapeUtil::MakeShape(S32, {}); + auto inferred_status0 = ShapeInference::InferConditionalShape( + r0s32, + {ShapeUtil::MakeProgramShape({vector_32_}, f32_), + ShapeUtil::MakeProgramShape({vector_64_}, f32_), + ShapeUtil::MakeProgramShape({vector_64_}, f32_)}, + {vector_32_, vector_64_, vector_64_}); + EXPECT_IS_OK(inferred_status0.status()); + EXPECT_TRUE(ShapeUtil::Equal(f32_, inferred_status0.ValueOrDie())); + + auto inferred_status1 = ShapeInference::InferConditionalShape( + r0s32, + {ShapeUtil::MakeProgramShape({matrix_32_48_}, vector_64_), + ShapeUtil::MakeProgramShape({vector_32_}, vector_64_), + ShapeUtil::MakeProgramShape({matrix_32_48_}, vector_64_)}, + {matrix_32_48_, vector_32_, matrix_32_48_}); + EXPECT_IS_OK(inferred_status1.status()); + EXPECT_TRUE(ShapeUtil::Equal(vector_64_, inferred_status1.ValueOrDie())); + + auto tuple_f32_v32 = ShapeUtil::MakeTupleShape({f32_, vector_32_}); + auto inferred_status2 = ShapeInference::InferConditionalShape( + r0s32, {ShapeUtil::MakeProgramShape({tuple_f32_v32}, vector_32_)}, + {tuple_f32_v32}); + EXPECT_IS_OK(inferred_status2.status()); + EXPECT_TRUE(ShapeUtil::Equal(vector_32_, inferred_status2.ValueOrDie())); + + auto inferred_status_error0 = ShapeInference::InferConditionalShape( + pred_, + {ShapeUtil::MakeProgramShape({vector_32_}, f32_), + ShapeUtil::MakeProgramShape({vector_32_}, f32_), + ShapeUtil::MakeProgramShape({vector_64_}, f32_)}, + {vector_32_, vector_32_, vector_64_}); + EXPECT_FALSE(inferred_status_error0.ok()); + EXPECT_THAT(inferred_status_error0.status().error_message(), + HasSubstr("2 == branch_computations.size()")); + + auto inferred_status_error1 = ShapeInference::InferConditionalShape( + r0s32, + {ShapeUtil::MakeProgramShape({matrix_32_48_}, vector_32_), + ShapeUtil::MakeProgramShape({f32_, vector_32_}, vector_32_), + ShapeUtil::MakeProgramShape({matrix_32_48_}, vector_32_)}, + {matrix_32_48_, ShapeUtil::MakeTupleShape({f32_, vector_32_}), + matrix_32_48_}); + EXPECT_FALSE(inferred_status_error1.ok()); + EXPECT_THAT(inferred_status_error1.status().error_message(), + HasSubstr("branch computation 1 must take 1 argument")); + + auto inferred_status_error2 = ShapeInference::InferConditionalShape( + r0s32, + {ShapeUtil::MakeProgramShape({r0s32}, f32_), + ShapeUtil::MakeProgramShape({vector_32_}, f32_), + ShapeUtil::MakeProgramShape({vector_32_}, f32_)}, + {r0s32, vector_32_, vector_64_}); + EXPECT_FALSE(inferred_status_error2.ok()); + EXPECT_THAT(inferred_status_error2.status().error_message(), + HasSubstr("branch operand 2 must match the shape of the only " + "parameter of branch computation 2")); + + auto inferred_status_error3 = ShapeInference::InferConditionalShape( + r0s32, + {ShapeUtil::MakeProgramShape({vector_32_}, f32_), + ShapeUtil::MakeProgramShape({vector_32_}, f32_), + ShapeUtil::MakeProgramShape({vector_32_}, f32_), + ShapeUtil::MakeProgramShape({vector_64_}, vector_32_)}, + {vector_32_, vector_32_, vector_32_, vector_64_}); + EXPECT_FALSE(inferred_status_error3.ok()); + EXPECT_THAT(inferred_status_error3.status().error_message(), + HasSubstr("the result of branch 0 computation and branch 3 " + "computation must have the same shape")); + + auto inferred_status_error4 = + ShapeInference::InferConditionalShape(r0s32, {}, {}); + EXPECT_FALSE(inferred_status_error4.ok()); + EXPECT_THAT(inferred_status_error4.status().error_message(), + HasSubstr("!branch_computations.empty()")); } TEST_F(ShapeInferenceTest, BadSlice) { diff --git a/tensorflow/compiler/xla/service/while_loop_analysis.cc b/tensorflow/compiler/xla/service/while_loop_analysis.cc index c93a9ba3176002a34fe84a29e62075de4d19168f..40f268f889bd538f73b8d6597502161670a3844d 100644 --- a/tensorflow/compiler/xla/service/while_loop_analysis.cc +++ b/tensorflow/compiler/xla/service/while_loop_analysis.cc @@ -14,15 +14,18 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/service/while_loop_analysis.h" +#include "absl/base/casts.h" #include "tensorflow/compiler/xla/service/hlo_evaluator.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/pattern_matcher.h" namespace xla { using absl::nullopt; using absl::optional; +namespace m = match; // Finds and returns the non-constant operand in instr. // @@ -48,34 +51,29 @@ static optional GetGTEOperandIndex(const HloInstruction* instr, const HloInstruction* gte_operand) { VLOG(2) << "GetGTEOperandIndex(" << instr->ToString() << ", " << gte_operand->ToString() << ")"; - optional tuple_idx; + + // Among the operands of `instr`, find one that is a get-tuple-element op. + auto gte_it = c_find_if(instr->operands(), [](const HloInstruction* instr) { + return instr->opcode() == HloOpcode::kGetTupleElement; + }); + if (gte_it == instr->operands().end()) { + VLOG(2) << "instr does not have a gte operand."; + return nullopt; + } + + // All operands of `instr` must be either constants or of the form + // get-tuple-element(gte_operand, tuple_idx) + // for the same value tuple_idx. + int64 tuple_idx = (*gte_it)->tuple_index(); for (const HloInstruction* operand : instr->operands()) { - if (operand->IsConstant()) { - continue; - } - // Look through copies. - // TODO(b/68830972): We wouldn't need this if for loop matching on the GPU - // would run before copy insertion. - if (operand->opcode() == HloOpcode::kCopy) { - operand = operand->operand(0); - } - if (operand->opcode() != HloOpcode::kGetTupleElement) { - VLOG(2) << "instr uses something other than gte(gte_operand): " - << operand->ToString(); - return nullopt; - } - if (operand->operand(0) != gte_operand) { - VLOG(2) << "instr has gte whose operand is not gte_operand: " - << operand->ToString(); + if (!Match(operand, m::Constant()) && + !Match(operand, + m::GetTupleElement(m::Op().Is(gte_operand), tuple_idx))) { + VLOG(2) + << "instr uses something other than a constant or gte(gte_operand, " + << tuple_idx << "): " << operand->ToString(); return nullopt; } - if (tuple_idx && tuple_idx != operand->tuple_index()) { - VLOG(2) << "instr has operands with conflicting gte indices, " - << *tuple_idx << " vs " << operand->tuple_index(); - return nullopt; - } - - tuple_idx = operand->tuple_index(); } return tuple_idx; } @@ -166,8 +164,171 @@ static optional GetLoopInductionVarTupleIdx( return indvar_tuple_idx; } +// Converts the given literal to a scalar int64, if possible. +// +// Fails if the literal is not an integral type or if the value it contains +// cannot be represented in an int64. +static optional LiteralAsScalarInt64(const Literal& l) { + if (!ShapeUtil::IsEffectiveScalar(l.shape())) { + VLOG(2) << "literal is not an effective scalar: " << l.ToString(); + return nullopt; + } + switch (l.shape().element_type()) { + case S8: + return l.GetFirstElement(); + case S16: + return l.GetFirstElement(); + case S32: + return l.GetFirstElement(); + case S64: + return l.GetFirstElement(); + case U8: + return l.GetFirstElement(); + case U16: + return l.GetFirstElement(); + case U32: + return l.GetFirstElement(); + case U64: { + uint64 v = l.GetFirstElement(); + if (v > static_cast(std::numeric_limits::max())) { + VLOG(2) << "uint64 literal is out of range for int64: " << v; + return nullopt; + } + return v; + } + default: + VLOG(2) << "literal is of non-integral type " << l.shape().ToString(); + return nullopt; + } +} + +// Computes a + b, returning nullopt if it overflows. +optional CheckedAdd(int64 a, int64 b) { + // Overflow occurred iff `a` and `b` have the same sign and `a + b` has a + // different sign, see Hacker's Delignt 2nd Ed. pp 28. + uint64 aa = absl::bit_cast(a); + uint64 bb = absl::bit_cast(b); + int64 result = absl::bit_cast(aa + bb); + if (a >= 0 == b >= 0 && result >= 0 != a >= 0) { + return nullopt; + } + return result; +} + +// Computes a - b, returning nullopt if it overflows. +optional CheckedSubtract(int64 a, int64 b) { + uint64 aa = absl::bit_cast(a); + uint64 bb = absl::bit_cast(b); + int64 result = absl::bit_cast(aa - bb); + // Overflow occurred iff `a` and `b` have different signs and the sign of + // `a - b` is the same as that of `b`, see Hacker's Delight 2nd Ed. pp 29. + if (a >= 0 != b >= 0 && result >= 0 == b >= 0) { + return nullopt; + } + return result; +} + +// Check if +// - `i` is initialized to a scalar constant K (namely, `indvar_init`), +// - the while condition does `i < N` or `i <= N`, and +// - the while body does `i++`. +// If so, it's trivial to compute the loop bound. +static optional PatternMatchLoopTripCount(HloInstruction* while_op, + int64 indvar_tuple_idx, + const Literal& indvar_init) { + // First, find the scalar constant K that `i` is initialized to. + optional indvar_init_val = LiteralAsScalarInt64(indvar_init); + if (!indvar_init_val) { + VLOG(2) << "Pattern-match failed: induction variable init is not a " + "constant scalar representable as an int64: " + << indvar_init.ToString(); + return nullopt; + } + + // Check that `i` goes as `i++` in the while body. + // + // TODO(jlebar): We could also handle i-- and other idioms. + auto* while_body = while_op->while_body(); + auto* while_body_indvar_update = + while_body->root_instruction()->operand(indvar_tuple_idx); + auto* while_body_indvar = NonConstantOperand(while_body_indvar_update); + if (!Match(while_body_indvar_update, + m::AddAnyOrder(m::Op().Is(while_body_indvar), + m::ConstantEffectiveScalar(1)))) { + VLOG(2) << "Pattern-match failed: induction variable does not go as i++: " + << while_body_indvar_update->ToString(); + return nullopt; + } + + // Check that we do op(i, N) or op(N, i) as the while condition. Capture the + // value N. + auto* while_cond = while_op->while_condition(); + auto* while_cond_root = while_cond->root_instruction(); + auto* while_cond_indvar = NonConstantOperand(while_cond_root); + HloInstruction* while_cond_bound = nullptr; + if (!Match(while_cond_root, + m::Op().WithBinaryOperandsAnyOrder( + m::Op().Is(while_cond_indvar), + m::ConstantEffectiveScalar(&while_cond_bound)))) { + VLOG(2) << "Pattern-match failed: while condition is not of the form " + "op(i, N) or op(N, i)."; + return nullopt; + } + // Note: If this succeeds, the constant `N` is representable as an int64 -- + // that is, if it's an XLA U64, it fits within an int64. + optional while_cond_bound_val = + LiteralAsScalarInt64(while_cond_bound->literal()); + if (!while_cond_bound_val) { + VLOG(2) << "Pattern-match failed: while condition induction variable is " + "not a constant scalar representable as an int64."; + return nullopt; + } + + // Handle `i = K; i < N; ++i`. + if (Match(while_cond_root, + m::Op() + .WithOpcode(HloOpcode::kLt) + .WithOperand(0, m::Op().Is(while_cond_indvar)))) { + VLOG(2) << "Pattern-match succeeded: loop condition is i < N: " + << while_cond_root->ToString(); + optional trips = + CheckedSubtract(*while_cond_bound_val, *indvar_init_val); + if (trips) { + return std::max(int64{0}, *trips); + } else { + VLOG(2) << "Pattern-match failed: Trip count exceeds INT64_MAX."; + return nullopt; + } + } + + // Handle `i = K; i <= N; ++i`. + if (Match(while_cond_root, + m::Op() + .WithOpcode(HloOpcode::kLe) + .WithOperand(0, m::Op().Is(while_cond_indvar)))) { + VLOG(2) << "Pattern-match succeeded: loop condition is i <= N: " + << while_cond_root->ToString(); + optional trips = + CheckedSubtract(*while_cond_bound_val, *indvar_init_val); + if (!trips) { + VLOG(2) << "Pattern-match failed: Trip count exceeds INT64_MAX"; + return nullopt; + } + trips = CheckedAdd(*trips, 1); + if (!trips) { + VLOG(2) << "Pattern-match failed: Trip count exceeds INT64_MAX"; + return nullopt; + } + return std::max(0, *trips); + } + + VLOG(2) << "Pattern-match failed: while condition follows unknown pattern: " + << while_cond_root->ToString(); + return nullopt; +} + optional ComputeWhileLoopTripCount(HloInstruction* while_op, - int64 max_value_returned) { + int64 max_brute_force_iters) { VLOG(2) << "Getting trip count for loop " << while_op->ToString(); // The loop's induction variable is found at @@ -188,23 +349,30 @@ optional ComputeWhileLoopTripCount(HloInstruction* while_op, auto* indvar_init = while_init->mutable_operand(*indvar_tuple_idx); StatusOr indvar_init_result = evaluator.Evaluate(indvar_init); if (!indvar_init_result.ok()) { - VLOG(2) << "Couldn't evaluate induction variable init: " - << indvar_init_result.status(); + VLOG(2) << "Couldn't evaluate induction variable init, " + << indvar_init_result.status() << ", " << indvar_init->ToString(); return nullopt; } + Literal indvar_iter_val = std::move(indvar_init_result).ValueOrDie(); + + // First, try to pattern-match. + if (auto trip_count = PatternMatchLoopTripCount(while_op, *indvar_tuple_idx, + indvar_iter_val)) { + return trip_count; + } + // If our pattern-match failed, try brute-forcing the loop trip count. auto* while_body = while_op->while_body(); auto* while_body_indvar_update = while_body->root_instruction()->operand(*indvar_tuple_idx); auto* while_body_indvar = NonConstantOperand(while_body_indvar_update); - // The initial value of the induction variable. - Literal indvar_iter_val = std::move(indvar_init_result).ValueOrDie(); - for (int64 trip_count = 0; trip_count != max_value_returned + 1; + auto* while_cond = while_op->while_condition(); + auto* while_cond_root = while_cond->root_instruction(); + auto* while_cond_indvar = NonConstantOperand(while_cond_root); + + for (int64 trip_count = 0; trip_count != max_brute_force_iters + 1; ++trip_count) { - auto* while_cond = while_op->while_condition(); - auto* while_cond_root = while_cond->root_instruction(); - auto* while_cond_indvar = NonConstantOperand(while_cond_root); StatusOr result = evaluator.EvaluateWithSubstitutions( while_cond_root, {{while_cond_indvar, &indvar_iter_val}}); if (!result.ok()) { diff --git a/tensorflow/compiler/xla/service/while_loop_analysis.h b/tensorflow/compiler/xla/service/while_loop_analysis.h index ac69a727bd6b403672a676400993fb7d8afc0a55..9bb784a544b0c8bbf6d9e96742c3ab35c609b113 100644 --- a/tensorflow/compiler/xla/service/while_loop_analysis.h +++ b/tensorflow/compiler/xla/service/while_loop_analysis.h @@ -22,11 +22,14 @@ limitations under the License. namespace xla { // Returns the precise trip count of the loop if it's statically known, -// nullopt otherwise. max_value_returned limits the number of steps that are -// evaluated while trying to brute force a loop trip count, trip counts larger -// than max_value_returned result in nullopt. -absl::optional ComputeWhileLoopTripCount(HloInstruction *while_op, - int64 max_value_returned = 128); +// nullopt otherwise. +// +// max_brute_force_iters limits the number of steps that are evaluated while +// trying to brute force a loop trip count. trip counts larger than +// max_brute_force_iters may be returned if we can pattern-match the loop +// condition. +absl::optional ComputeWhileLoopTripCount( + HloInstruction *while_op, int64 max_brute_force_iters = 128); // Returns an upper bound on the trip count of the loop if it's statically // known, nullopt otherwise. diff --git a/tensorflow/compiler/xla/service/while_loop_trip_count_annotator.cc b/tensorflow/compiler/xla/service/while_loop_trip_count_annotator.cc new file mode 100644 index 0000000000000000000000000000000000000000..03bb6792fe74e9eb90278cbd4152e609a7904c80 --- /dev/null +++ b/tensorflow/compiler/xla/service/while_loop_trip_count_annotator.cc @@ -0,0 +1,40 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/while_loop_trip_count_annotator.h" +#include "tensorflow/compiler/xla/service/while_loop_analysis.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { + +StatusOr WhileLoopTripCountAnnotator::Run(HloModule* module) { + bool changed = false; + for (const HloComputation* comp : module->computations()) { + for (HloInstruction* instr : comp->instructions()) { + if (instr->opcode() != HloOpcode::kWhile) { + continue; + } + if (auto trip_count = ComputeWhileLoopTripCount(instr)) { + WhileLoopBackendConfig config; + config.mutable_known_trip_count()->set_n(*trip_count); + TF_RETURN_IF_ERROR(instr->set_backend_config(config)); + changed = true; + } + } + } + return changed; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/while_loop_trip_count_annotator.h b/tensorflow/compiler/xla/service/while_loop_trip_count_annotator.h new file mode 100644 index 0000000000000000000000000000000000000000..7cda2f10cefba821bccc1b5d3b5a33cd7a68e004 --- /dev/null +++ b/tensorflow/compiler/xla/service/while_loop_trip_count_annotator.h @@ -0,0 +1,48 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_TRIP_COUNT_ANNOTATOR_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_TRIP_COUNT_ANNOTATOR_H_ + +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" +#include "tensorflow/compiler/xla/statusor.h" + +namespace xla { + +// Pass that annotates `while` loops with known trip counts. +// +// The annotation is stored as a backend-config on the while loop node. +// +// This pass should run after all passes that might semantically modify a while +// loop, e.g. by unrolling it. Otherwise, a loop could end up with a +// backend-config that doesn't match its true trip-count. +// +// This pass does some pattern-matching on loop bodies and conditions, so it +// should run after most HLO simplifications and before fusion and layout +// assignment, which make pattern matching much more difficult by e.g. +// introducing `copy` nodes. +class WhileLoopTripCountAnnotator : public HloModulePass { + public: + ~WhileLoopTripCountAnnotator() override {} + absl::string_view name() const override { + return "while-loop-trip-count-annotator"; + } + StatusOr Run(HloModule* module) override; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_TRIP_COUNT_ANNOTATOR_H_ diff --git a/tensorflow/compiler/xla/service/while_loop_trip_count_annotator_test.cc b/tensorflow/compiler/xla/service/while_loop_trip_count_annotator_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..5c19cbc015dab7721dc60e015ca7125cc628fc0a --- /dev/null +++ b/tensorflow/compiler/xla/service/while_loop_trip_count_annotator_test.cc @@ -0,0 +1,207 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/while_loop_trip_count_annotator.h" +#include "tensorflow/compiler/xla/service/pattern_matcher.h" +#include "tensorflow/compiler/xla/service/while_loop_simplifier.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/core/lib/core/status_test_util.h" + +namespace xla { +namespace { + +class TripCountAnnotatorTest : public HloTestBase {}; + +TEST_F(TripCountAnnotatorTest, KnownSmallTripCount) { + const char* kModuleStr = R"( + HloModule test + Body { + param = (s32[]) parameter(0) + i = s32[] get-tuple-element(param), index=0 + one = s32[] constant(1) + i_plus_one = s32[] add(i, one) + ROOT tuple = (s32[]) tuple(i_plus_one) + } + + Cond { + param = (s32[]) parameter(0) + i = s32[] get-tuple-element(param), index=0 + trip_count = s32[] constant(10) + ROOT done = pred[] less-than(i, trip_count) + } + + ENTRY test { + i_start = s32[] constant(0) + initial_tuple = (s32[]) tuple(i_start) + ROOT while = (s32[]) while(initial_tuple), condition=Cond, body=Body + })"; + + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + WhileLoopTripCountAnnotator pass; + TF_ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, m.get())); + ASSERT_TRUE(changed); + + TF_ASSERT_OK_AND_ASSIGN(auto config, + m->entry_computation() + ->root_instruction() + ->backend_config()); + EXPECT_EQ(10, config.known_trip_count().n()); +} + +TEST_F(TripCountAnnotatorTest, KnownLargeTripCount) { + const char* kModuleStr = R"( + HloModule test + Body { + param = (s32[]) parameter(0) + i = s32[] get-tuple-element(param), index=0 + one = s32[] constant(1) + i_plus_one = s32[] add(i, one) + ROOT tuple = (s32[]) tuple(i_plus_one) + } + + Cond { + param = (s32[]) parameter(0) + i = s32[] get-tuple-element(param), index=0 + trip_count = s32[] constant(1000000) + ROOT done = pred[] less-than(i, trip_count) + } + + ENTRY test { + i_start = s32[] constant(0) + initial_tuple = (s32[]) tuple(i_start) + ROOT while = (s32[]) while(initial_tuple), condition=Cond, body=Body + })"; + + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + WhileLoopTripCountAnnotator pass; + TF_ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, m.get())); + ASSERT_TRUE(changed); + + TF_ASSERT_OK_AND_ASSIGN(auto config, + m->entry_computation() + ->root_instruction() + ->backend_config()); + EXPECT_EQ(1000000, config.known_trip_count().n()); +} + +TEST_F(TripCountAnnotatorTest, NonzeroStart) { + const char* kModuleStr = R"( + HloModule test + Body { + param = (s32[]) parameter(0) + i = s32[] get-tuple-element(param), index=0 + one = s32[] constant(1) + i_plus_one = s32[] add(i, one) + ROOT tuple = (s32[]) tuple(i_plus_one) + } + + Cond { + param = (s32[]) parameter(0) + i = s32[] get-tuple-element(param), index=0 + trip_count = s32[] constant(1000000) + ROOT done = pred[] less-than(i, trip_count) + } + + ENTRY test { + i_start = s32[] constant(10) + initial_tuple = (s32[]) tuple(i_start) + ROOT while = (s32[]) while(initial_tuple), condition=Cond, body=Body + })"; + + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + WhileLoopTripCountAnnotator pass; + TF_ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, m.get())); + ASSERT_TRUE(changed); + + TF_ASSERT_OK_AND_ASSIGN(auto config, + m->entry_computation() + ->root_instruction() + ->backend_config()); + EXPECT_EQ(999990, config.known_trip_count().n()); +} + +TEST_F(TripCountAnnotatorTest, LessThanOrEqualTo) { + const char* kModuleStr = R"( + HloModule test + Body { + param = (s32[]) parameter(0) + i = s32[] get-tuple-element(param), index=0 + one = s32[] constant(1) + i_plus_one = s32[] add(i, one) + ROOT tuple = (s32[]) tuple(i_plus_one) + } + + Cond { + param = (s32[]) parameter(0) + i = s32[] get-tuple-element(param), index=0 + trip_count = s32[] constant(1000000) + ROOT done = pred[] less-than-or-equal-to(i, trip_count) + } + + ENTRY test { + i_start = s32[] constant(10) + initial_tuple = (s32[]) tuple(i_start) + ROOT while = (s32[]) while(initial_tuple), condition=Cond, body=Body + })"; + + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + WhileLoopTripCountAnnotator pass; + TF_ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, m.get())); + ASSERT_TRUE(changed); + + TF_ASSERT_OK_AND_ASSIGN(auto config, + m->entry_computation() + ->root_instruction() + ->backend_config()); + EXPECT_EQ(999991, config.known_trip_count().n()); +} + +TEST_F(TripCountAnnotatorTest, Int64Overflow) { + // for(i = INT64_MIN; i < INT64_MAX; ++i) + // + // We store the trip count as an int64, so this loop is unanalyzable. + const char* kModuleStr = R"( + HloModule test + Body { + param = (s64[]) parameter(0) + i = s64[] get-tuple-element(param), index=0 + one = s64[] constant(1) + i_plus_one = s64[] add(i, one) + ROOT tuple = (s64[]) tuple(i_plus_one) + } + + Cond { + param = (s64[]) parameter(0) + i = s64[] get-tuple-element(param), index=0 + trip_count = s64[] constant(9223372036854775807) // 2^63-1 + ROOT done = pred[] less-than-or-equal-to(i, trip_count) + } + + ENTRY test { + i_start = s64[] constant(-9223372036854775808) // -2^63 + initial_tuple = (s64[]) tuple(i_start) + ROOT while = (s64[]) while(initial_tuple), condition=Cond, body=Body + })"; + + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + WhileLoopTripCountAnnotator pass; + TF_ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, m.get())); + EXPECT_FALSE(changed); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/shape_test.cc b/tensorflow/compiler/xla/shape_test.cc index 526abafea5cc244418a4ec05db7da6203716b483..dbdafcc0a1f7348af8394598363d570118cdd87e 100644 --- a/tensorflow/compiler/xla/shape_test.cc +++ b/tensorflow/compiler/xla/shape_test.cc @@ -35,6 +35,8 @@ class ShapeTest : public ::testing::Test { const Shape opaque_ = ShapeUtil::MakeOpaqueShape(); const Shape token_ = ShapeUtil::MakeTokenShape(); const Shape scalar_ = ShapeUtil::MakeShape(F32, {}); + const Shape scalar_with_tile_ = + ShapeUtil::MakeShapeWithLayout(F32, {}, {}, {Tile({256})}); const Shape matrix_ = ShapeUtil::MakeShape(U32, {1, 2}); const Shape matrix2_ = ShapeUtil::MakeShapeWithLayout(S32, {3, 4}, {0, 1}); const Shape tuple_ = @@ -66,6 +68,8 @@ TEST_F(ShapeTest, ShapeToString) { EXPECT_EQ("opaque[]", opaque_.ToString(/*print_layout=*/true)); EXPECT_EQ("f32[]", scalar_.ToString(/*print_layout=*/true)); + EXPECT_EQ("f32[]{:T(256)}", + scalar_with_tile_.ToString(/*print_layout=*/true)); EXPECT_EQ("u32[1,2]{1,0}", matrix_.ToString(/*print_layout=*/true)); EXPECT_EQ("s32[3,4]{0,1}", matrix2_.ToString(/*print_layout=*/true)); EXPECT_EQ("(opaque[], f32[], u32[1,2]{1,0}, s32[3,4]{0,1})", diff --git a/tensorflow/compiler/xla/shape_util.cc b/tensorflow/compiler/xla/shape_util.cc index d045fc7a9e291258640eca75166e116cf7390a7b..acaa9cae7c2c2745a3ed413ca9f00b5bf0187a0c 100644 --- a/tensorflow/compiler/xla/shape_util.cc +++ b/tensorflow/compiler/xla/shape_util.cc @@ -462,10 +462,14 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( shape.is_dynamic_dimension(i) ? "<=" : "", shape.dimensions(i)); } result += "]"; - if (!IsScalar(shape) && shape.IsArray()) { - if (LayoutUtil::HasLayout(shape)) { - StrAppend(&result, LayoutUtil::HumanString(shape.layout())); + if (IsScalar(shape)) { + string layout_str = LayoutUtil::HumanString(shape.layout()); + // Don't print "{}" as layout for scalars. + if (layout_str != "{}") { + StrAppend(&result, layout_str); } + } else if (shape.IsArray() && LayoutUtil::HasLayout(shape)) { + StrAppend(&result, LayoutUtil::HumanString(shape.layout())); } return result; } diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD index 72c5ed58e284c650e9daa13dec714d60f347c4e3..15d05f6a6e6bcfb79712fc4c3b36bef6113546fb 100644 --- a/tensorflow/compiler/xla/tests/BUILD +++ b/tensorflow/compiler/xla/tests/BUILD @@ -552,6 +552,7 @@ xla_test( xla_test( name = "conditional_test", srcs = ["conditional_test.cc"], + shard_count = 2, deps = [ "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:global_data", @@ -674,19 +675,20 @@ xla_test( ) xla_test( - name = "exhaustive_f32_elementwise_op_test", - srcs = ["exhaustive_f32_elementwise_op_test.cc"], + name = "exhaustive_op_test", + srcs = ["exhaustive_op_test.cc"], real_hardware_only = True, # Very slow on the interpreter. shard_count = 48, tags = [ "optonly", # This is a big test that we skip for capacity reasons in OSS testing. - "nooss", + "no_oss", ], deps = [ ":client_library_test_base", ":literal_test_util", "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/client/lib:constants", "//tensorflow/compiler/xla/client/lib:math", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", @@ -1600,6 +1602,39 @@ xla_test( ], ) +xla_test( + name = "multi_device_all_reduce_test", + srcs = ["multi_device_all_reduce_test.cc"], + backends = ["gpu"], + tags = [ + "manual", + "multi_gpu", + "no_oss", + "notap", + ], + deps = [ + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_parser", + "//tensorflow/compiler/xla/service:hlo_runner", + "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/compiler/xla/tests:test_utils", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:lib", + "//tensorflow/core:test", + "@com_google_absl//absl/memory", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", + ], +) + xla_test( name = "bitcast_convert_test", srcs = ["bitcast_convert_test.cc"], @@ -2183,3 +2218,23 @@ xla_test( "//tensorflow/core:test", ], ) + +xla_test( + name = "cholesky_test", + srcs = ["cholesky_test.cc"], + tags = ["optonly"], + deps = [ + "//tensorflow/compiler/xla:array2d", + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/client/lib:arithmetic", + "//tensorflow/compiler/xla/client/lib:matrix", + "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:test", + ], +) diff --git a/tensorflow/compiler/xla/client/lib/cholesky_test.cc b/tensorflow/compiler/xla/tests/cholesky_test.cc similarity index 54% rename from tensorflow/compiler/xla/client/lib/cholesky_test.cc rename to tensorflow/compiler/xla/tests/cholesky_test.cc index 095dd4fbf8b7c90047c4428b50c626c16e9c1e94..272d5784362dd347061e7178ff48f9fab4ffd822 100644 --- a/tensorflow/compiler/xla/client/lib/cholesky_test.cc +++ b/tensorflow/compiler/xla/tests/cholesky_test.cc @@ -13,8 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/xla/client/lib/cholesky.h" - +#include #include #include #include @@ -32,27 +31,27 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/status_test_util.h" +namespace xla { namespace { -using xla::int64; - -using CholeskyTest = xla::ClientLibraryTestBase; +using CholeskyTest = ClientLibraryTestBase; -XLA_TEST_F(CholeskyTest, Simple) { - xla::XlaBuilder builder(TestName()); +XLA_TEST_F(CholeskyTest, Lower) { + XlaBuilder builder(TestName()); - xla::Array2D a_vals({ - {4, 6, 8, 10}, - {6, 45, 54, 63}, - {8, 54, 146, 166}, + float nan = std::numeric_limits::quiet_NaN(); + Array2D a_vals({ + {4, nan, nan, nan}, + {6, 45, nan, nan}, + {8, 54, 146, nan}, {10, 63, 166, 310}, }); - xla::XlaOp a; + XlaOp a; auto a_data = CreateR2Parameter(a_vals, 0, "a", &builder, &a); - xla::Cholesky(a, /*block_size=*/2); + LowerTriangle(Cholesky(a, /*lower=*/true)); - xla::Array2D expected({ + Array2D expected({ {2, 0, 0, 0}, {3, 6, 0, 0}, {4, 7, 9, 0}, @@ -60,34 +59,62 @@ XLA_TEST_F(CholeskyTest, Simple) { }); ComputeAndCompareR2(&builder, expected, {a_data.get()}, - xla::ErrorSpec(1e-4, 1e-4)); + ErrorSpec(1e-4, 1e-4)); +} + +XLA_TEST_F(CholeskyTest, Upper) { + XlaBuilder builder(TestName()); + + float nan = std::numeric_limits::quiet_NaN(); + Array2D a_vals({ + {4, 6, 8, 10}, + {nan, 45, 54, 63}, + {nan, nan, 146, 166}, + {nan, nan, nan, 310}, + }); + + XlaOp a; + auto a_data = CreateR2Parameter(a_vals, 0, "a", &builder, &a); + UpperTriangle(Cholesky(a, /*lower=*/false)); + + Array2D expected({ + {2, 3, 4, 5}, + {0, 6, 7, 8}, + {0, 0, 9, 10}, + {0, 0, 0, 11}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get()}, + ErrorSpec(1e-4, 1e-4)); } XLA_TEST_F(CholeskyTest, Simple2) { - xla::XlaBuilder builder(TestName()); + XlaBuilder builder(TestName()); - xla::Array2D a_vals({ + Array2D a_vals({ {16, 24, 8, 12}, {24, 61, 82, 48}, {8, 82, 456, 106}, {12, 48, 106, 62}, }); - xla::XlaOp a; + XlaOp a; auto a_data = CreateR2Parameter(a_vals, 0, "a", &builder, &a); - xla::Cholesky(a); + LowerTriangle(Cholesky(a, /*lower=*/true)); - xla::Array2D expected( - {{4, 0, 0, 0}, {6, 5, 0, 0}, {2, 14, 16, 0}, {3, 6, 1, 4}}); + Array2D expected({{4, 0, 0, 0}, // + {6, 5, 0, 0}, // + {2, 14, 16, 0}, // + {3, 6, 1, 4}}); ComputeAndCompareR2(&builder, expected, {a_data.get()}, - xla::ErrorSpec(1e-4, 1e-4)); + ErrorSpec(1e-4, 1e-4)); } XLA_TEST_F(CholeskyTest, SimpleBatched) { - xla::XlaBuilder builder(TestName()); + XlaBuilder builder(TestName()); - xla::Array3D a_vals({ + Array3D a_vals({ { {4, 6, 8, 10}, {6, 45, 54, 63}, @@ -102,65 +129,78 @@ XLA_TEST_F(CholeskyTest, SimpleBatched) { }, }); - xla::XlaOp a; + XlaOp a; auto a_data = CreateR3Parameter(a_vals, 0, "a", &builder, &a); - xla::Cholesky(a); + LowerTriangle(Cholesky(a, /*lower=*/true)); - xla::Array3D expected({ + Array3D expected({ { {2, 0, 0, 0}, {3, 6, 0, 0}, {4, 7, 9, 0}, {5, 8, 10, 11}, }, - {{4, 0, 0, 0}, {6, 5, 0, 0}, {2, 14, 16, 0}, {3, 6, 1, 4}}, + {{4, 0, 0, 0}, // + {6, 5, 0, 0}, // + {2, 14, 16, 0}, // + {3, 6, 1, 4}}, }); ComputeAndCompareR3(&builder, expected, {a_data.get()}, - xla::ErrorSpec(1e-4, 1e-4)); + ErrorSpec(1e-4, 1e-4)); } -using CholeskyTestCase = std::tuple; +using CholeskyTestCase = std::tuple; class RandomCholeskyTest - : public xla::ClientLibraryTestBase, + : public ClientLibraryTestBase, public ::testing::WithParamInterface {}; XLA_TEST_P(RandomCholeskyTest, Random) { - xla::XlaBuilder builder(TestName()); + XlaBuilder builder(TestName()); auto test_params = GetParam(); std::vector dimensions = {std::get<0>(test_params), std::get<1>(test_params), std::get<1>(test_params)}; - xla::Shape shape = xla::ShapeUtil::MakeShape(xla::F32, dimensions); + bool lower = std::get<2>(test_params); + Shape shape = ShapeUtil::MakeShape(F32, dimensions); TF_ASSERT_OK_AND_ASSIGN( - auto literal, - xla::LiteralUtil::CreateRandomLiteral(shape, 0.0, 1.0)); + auto literal, LiteralUtil::CreateRandomLiteral(shape, 0.0, 1.0)); - auto input = xla::Parameter(&builder, 0, shape, "input"); + auto input = Parameter(&builder, 0, shape, "input"); // Form a random positive definite matrix. - auto matrix = xla::BatchDot(input, TransposeInMinorDims(input), - xla::PrecisionConfig::HIGHEST); + auto matrix = + BatchDot(input, TransposeInMinorDims(input), PrecisionConfig::HIGHEST); - auto cholesky = xla::Cholesky(matrix, /*block_size=*/4); + auto cholesky = Triangle(Cholesky(matrix, lower), lower); // Verify that ||matrix - cholesky * cholesky_t||_2 ~= 0 - auto verification = xla::BatchDot(cholesky, TransposeInMinorDims(cholesky), - xla::PrecisionConfig::HIGHEST); + XlaOp verification; + if (lower) { + verification = BatchDot(cholesky, TransposeInMinorDims(cholesky), + PrecisionConfig::HIGHEST); + } else { + verification = BatchDot(TransposeInMinorDims(cholesky), cholesky, + PrecisionConfig::HIGHEST); + } auto delta = matrix - verification; - xla::Reduce(delta * delta, xla::ConstantR0(&builder, 0.0), - CreateScalarAddComputation(xla::F32, &builder), {0, 1, 2}); + Reduce(delta * delta, ConstantR0(&builder, 0.0), + CreateScalarAddComputation(F32, &builder), {0, 1, 2}); TF_ASSERT_OK_AND_ASSIGN(auto input_data, client_->TransferToServer(literal)); ComputeAndCompareR0(&builder, 0.0, {input_data.get()}, - xla::ErrorSpec(1e-4, 1e-4)); + ErrorSpec(1e-4, 1e-4)); } INSTANTIATE_TEST_SUITE_P(RandomCholeskyTestInstance, RandomCholeskyTest, - ::testing::Values(CholeskyTestCase{1, 1}, - CholeskyTestCase{1, 2}, - CholeskyTestCase{10, 5}, - CholeskyTestCase{2, 20})); + ::testing::Values(CholeskyTestCase{1, 1, true}, + CholeskyTestCase{1, 2, true}, + CholeskyTestCase{1, 50, true}, + CholeskyTestCase{1, 50, false}, + CholeskyTestCase{10, 5, true}, + CholeskyTestCase{5, 10, false}, + CholeskyTestCase{2, 20, true})); } // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/conditional_test.cc b/tensorflow/compiler/xla/tests/conditional_test.cc index 32cac499c7439af80bafb88ac61b0b078f589599..f75c3fb01e2c854475537ca4b413f381cf74355c 100644 --- a/tensorflow/compiler/xla/tests/conditional_test.cc +++ b/tensorflow/compiler/xla/tests/conditional_test.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/client/xla_computation.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" @@ -169,6 +170,11 @@ class ConditionalOpTest : public ClientLibraryTestBase { ErrorSpec error_spec_{0.001}; }; +// Test fixture to run indexed conditional (switch/case) tests with varying +// number of branches. +class CaseOpTest : public ConditionalOpTest, + public ::testing::WithParamInterface {}; + // Test true and false computations that do not take any parameters. XLA_TEST_F(ConditionalOpTest, Parameters0) { XlaBuilder builder(TestName()); @@ -182,6 +188,36 @@ XLA_TEST_F(ConditionalOpTest, Parameters0) { ComputeAndCompareR0(&builder, 56.0f, {pred_arg.get()}, error_spec_); } +// Test branch computations that do not take any parameters. +XLA_TEST_P(CaseOpTest, Parameters0) { + int num_branches = GetParam(); + for (int bi = -1; bi <= num_branches; ++bi) { + SCOPED_TRACE(bi); + XlaBuilder builder(TestName()); + XlaOp branch_index; + auto branch_index_arg = CreateR0Parameter(bi, 0, "branch_index_arg", + &builder, &branch_index); + auto operand = Tuple(&builder, {}); + + std::vector operands(num_branches, operand); + std::vector branches; + branches.reserve(num_branches); + std::vector branches_p(num_branches); + for (int i = 0; i < num_branches; ++i) { + branches.emplace_back( + CreateR0ConstantComputation(static_cast(i) * 10)); + branches_p[i] = &branches[i]; + } + Conditional(branch_index, branches_p, operands); + + float expected = 10 * static_cast((bi < 0 || bi >= num_branches) + ? num_branches - 1 + : bi); + ComputeAndCompareR0(&builder, expected, {branch_index_arg.get()}, + error_spec_); + } +} + // Test true and false computations that take in 1 parameter. XLA_TEST_F(ConditionalOpTest, Parameters1) { XlaBuilder builder(TestName()); @@ -195,6 +231,45 @@ XLA_TEST_F(ConditionalOpTest, Parameters1) { ComputeAndCompareR0(&builder, 12.0f, {pred_arg.get()}, error_spec_); } +// Test branch computations that take in 1 parameter. +XLA_TEST_P(CaseOpTest, Parameters1) { + int num_branches = GetParam(); + for (int bi = -1; bi <= num_branches; ++bi) { + SCOPED_TRACE(bi); + XlaBuilder builder(TestName()); + XlaOp branch_index; + auto branch_index_arg = CreateR0Parameter(bi, 0, "branch_index_arg", + &builder, &branch_index); + + auto make_branch = [&builder, this](int i) { + auto sb = builder.CreateSubBuilder(absl::StrCat("branch_", i)); + Add(ConstantR0(sb.get(), static_cast(i)), + Parameter(sb.get(), 0, r0f32_, "p0")); + return sb->BuildAndNoteError(); + }; + std::vector branches; + branches.reserve(num_branches); + std::vector branches_p(num_branches); + std::vector operands; + operands.reserve(num_branches); + std::vector expecteds(num_branches); + for (int i = 0; i < num_branches; ++i) { + branches.emplace_back(make_branch(i)); + branches_p[i] = &branches[i]; + auto fi = static_cast(i); + operands.emplace_back(ConstantR0(&builder, 10 * fi + 7)); + expecteds[i] = 10 * fi + 7 + fi; + } + + Conditional(branch_index, branches_p, operands); + float expected = (bi < 0 || bi >= num_branches) + ? expecteds[num_branches - 1] + : expecteds[bi]; + ComputeAndCompareR0(&builder, expected, {branch_index_arg.get()}, + error_spec_); + } +} + // Test conditional with two different computations in the true and false cases // that take in different arguments. XLA_TEST_F(ConditionalOpTest, DiffComputationsDiffArgs) { @@ -331,6 +406,46 @@ XLA_TEST_F(ConditionalOpTest, Parameters2ArrayTrueBranch) { error_spec_); } +// Test branch computations that take in 2 array parameters. +XLA_TEST_P(CaseOpTest, Parameters2Array) { + int num_branches = GetParam(); + for (int bi = -1; bi <= num_branches; ++bi) { + SCOPED_TRACE(bi); + XlaBuilder builder(TestName()); + XlaOp branch_index; + auto branch_index_arg = + CreateR0Parameter(bi, 0, "pred", &builder, &branch_index); + auto operand1 = ConstantR1(&builder, {24.0f, 56.0f}); + auto operand2 = ConstantR1(&builder, {10.0f, 11.0f}); + auto operands = Tuple(&builder, {operand1, operand2}); + auto make_branch = [&builder, this](int i) { + auto sb = builder.CreateSubBuilder(absl::StrCat("branch_", i)); + auto p = Parameter(sb.get(), 0, tuple_2_r1s2f32_, "p0"); + Add(Mul(ConstantR0(sb.get(), static_cast(i)), + GetTupleElement(p, 0)), + GetTupleElement(p, 1)); + return sb->BuildAndNoteError(); + }; + std::vector branches; + branches.reserve(num_branches); + std::vector branches_p(num_branches); + for (int i = 0; i < num_branches; ++i) { + branches.emplace_back(make_branch(i)); + branches_p[i] = &branches[i]; + } + Conditional(branch_index, branches_p, + std::vector(num_branches, operands)); + auto modified_bi = static_cast( + (bi < 0 || bi >= num_branches) ? num_branches - 1 : bi); + ComputeAndCompareR1( + &builder, {24.0f * modified_bi + 10, 56.0f * modified_bi + 11}, + {branch_index_arg.get()}, error_spec_); + } +} + +INSTANTIATE_TEST_SUITE_P(CaseOpTest_Instantiation, CaseOpTest, + ::testing::Values(1, 2, 3, 4, 5)); + // Test true and false computations that take in 2 array parameters and // predicate is false. XLA_TEST_F(ConditionalOpTest, Parameters2ArrayFalseBranch) { @@ -582,8 +697,8 @@ XLA_TEST_F(ConditionalOpTest, ShapeMismatch) { auto result = builder.Build(); EXPECT_FALSE(result.ok()); EXPECT_THAT(result.status().error_message(), - ::testing::HasSubstr("true_operand must match the shape of the " - "only parameter of true_computation")); + ::testing::HasSubstr("operand 0 must match the shape of the " + "only parameter of branch computation 0")); } XLA_TEST_F(ConditionalOpTest, SwappedInputsInSequentialConditionals) { diff --git a/tensorflow/compiler/xla/tests/convolution_test.cc b/tensorflow/compiler/xla/tests/convolution_test.cc index 9db9f2563b636c4f929585eb13a9c7f809833eda..cfee9c0f8a4c908d5dbdd5345ed7f839dfa4dee2 100644 --- a/tensorflow/compiler/xla/tests/convolution_test.cc +++ b/tensorflow/compiler/xla/tests/convolution_test.cc @@ -1945,7 +1945,7 @@ XLA_TEST_F(ConvolutionTest, ConvolveF32BackwardInputGroupedConvolution) { class ConvolutionHloTest : public HloTestBase {}; -XLA_TEST_F(ConvolutionHloTest, DISABLED_ON_CPU(ConvolveF64Forward)) { +XLA_TEST_F(ConvolutionHloTest, ConvolveF64Forward) { constexpr char kHlo[] = R"( HloModule TestModule @@ -1957,7 +1957,7 @@ ENTRY Test { EXPECT_TRUE(RunAndCompare(kHlo, ErrorSpec{0.001})); } -XLA_TEST_F(ConvolutionHloTest, DISABLED_ON_CPU(ConvolveF32ForwardReversed)) { +XLA_TEST_F(ConvolutionHloTest, ConvolveF32ForwardReversed) { constexpr char kHlo[] = R"( HloModule TestModule @@ -1969,7 +1969,7 @@ ENTRY Test { EXPECT_TRUE(RunAndCompare(kHlo, ErrorSpec{0.001})); } -XLA_TEST_F(ConvolutionHloTest, DISABLED_ON_CPU(ConvolveF64BackwardFilter)) { +XLA_TEST_F(ConvolutionHloTest, ConvolveF64BackwardFilter) { constexpr char kHlo[] = R"( HloModule TestModule @@ -1981,7 +1981,7 @@ ENTRY Test { EXPECT_TRUE(RunAndCompare(kHlo, ErrorSpec{0.001})); } -XLA_TEST_F(ConvolutionHloTest, DISABLED_ON_CPU(ConvolveF64BackwardInput)) { +XLA_TEST_F(ConvolutionHloTest, ConvolveF64BackwardInput) { constexpr char kHlo[] = R"( HloModule TestModule diff --git a/tensorflow/compiler/xla/tests/dot_operation_test.cc b/tensorflow/compiler/xla/tests/dot_operation_test.cc index 6ee2178a227a12b7baa933f036a44db8ec630a4c..414d0b14a6b4f0307851fcc717c5e8a74a33782b 100644 --- a/tensorflow/compiler/xla/tests/dot_operation_test.cc +++ b/tensorflow/compiler/xla/tests/dot_operation_test.cc @@ -1344,5 +1344,22 @@ ENTRY TransposeOutput { EXPECT_TRUE(RunAndCompare(hlo_string, ErrorSpec{4e-3, 4e-3})); } +XLA_TEST_F(DotOperationTextTest, MatrixVectorComplex) { + absl::string_view hlo_string = + R"( +HloModule MatrixVectorComplex + +ENTRY MatrixVectorComplex { + p0 = c64[5,5] parameter(0) + p1 = c64[5,1] parameter(1) + p2 = c64[5,1] parameter(2) + dot = c64[5,1] dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} + ROOT add = c64[5,1] add(dot, p2) +} +)"; + + EXPECT_TRUE(RunAndCompare(hlo_string, ErrorSpec{4e-3, 4e-3})); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/exhaustive_f32_elementwise_op_test.cc b/tensorflow/compiler/xla/tests/exhaustive_f32_elementwise_op_test.cc deleted file mode 100644 index b961e6102692cb3b90976d621c62cb4cf18a9b6b..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/tests/exhaustive_f32_elementwise_op_test.cc +++ /dev/null @@ -1,237 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include -#include "absl/base/casts.h" -#include "tensorflow/compiler/xla/client/lib/math.h" -#include "tensorflow/compiler/xla/client/xla_builder.h" -#include "tensorflow/compiler/xla/tests/client_library_test_base.h" -#include "tensorflow/compiler/xla/tests/literal_test_util.h" -#include "tensorflow/compiler/xla/tests/test_macros.h" - -namespace xla { -namespace { - -class ExhaustiveF32ElementwiseOpTest - : public ClientLibraryTestBase, - public ::testing::WithParamInterface> { - protected: - ErrorSpec error_spec_{0.0001, 0.0001}; - - bool IsClose(float expected, float actual) { - float abs_err = std::abs(expected - actual); - float rel_err = abs_err / std::abs(expected); - return abs_err < error_spec_.abs || rel_err < error_spec_.rel || - (std::isnan(expected) && std::isnan(actual)) || - (std::isinf(expected) && std::isinf(actual) && - (expected > 0) == (actual > 0)); - } - - template - void ExhaustivelyTestF32Op(EnqueueOpTy enqueue_op, - float (*evaluate_op)(float), - std::pair known_incorrect_range) { - SetFastMathDisabled(true); - - int64 begin, end; - std::tie(begin, end) = GetParam(); - int64 input_size = end - begin; - - if (begin >= known_incorrect_range.first && - end <= known_incorrect_range.second) { - LOG(INFO) << absl::StreamFormat( - "Skipping this shard, as the range under test, [%d, %d), falls " - "entirely within the known-incorrect range [%d, %d).", - begin, end, known_incorrect_range.first, - known_incorrect_range.second); - return; - } - - LOG(INFO) << "Checking range [" << begin << ", " << end << ")"; - - XlaBuilder builder(TestName()); - - auto ith_input_elem = [&](int64 i) -> float { - i += begin; - // If the operation is known to be buggy on a specific input clamp that - // input to 0 under the assumption that the op is at least correct on 0. - if (i >= known_incorrect_range.first && - i < known_incorrect_range.second) { - return 0; - } - return absl::bit_cast(i); - }; - - Literal input_literal = - LiteralUtil::CreateFromDimensions(F32, {input_size}); - absl::Span input_arr = input_literal.data(); - for (int64 i = 0; i < input_size; i++) { - input_arr[i] = ith_input_elem(i); - } - auto input = Parameter(&builder, 0, input_literal.shape(), "input"); - enqueue_op(&builder, input); - TF_ASSERT_OK_AND_ASSIGN(XlaComputation comp, builder.Build()); - - // Build and run the computation using the LocalClient API, rather than the - // plain Client API, which is used by ClientLibraryTestBase. This is - // because the plain Client API results does more memcpys to/from Literals, - // and that's slow given that we're touching a lot of data here. - // - // Copy debug options from ClientLibraryTestBase. In particular, we're - // interested in disabling constant folding. - ExecutableBuildOptions build_opts; - *build_opts.mutable_debug_options() = *mutable_debug_options(); - TF_ASSERT_OK_AND_ASSIGN( - auto executable, - client_->Compile(comp, {&input_literal.shape()}, build_opts)); - - TF_ASSERT_OK_AND_ASSIGN( - ScopedShapedBuffer input_data, - client_->LiteralToShapedBuffer(input_literal, /*device_ordinal=*/0)); - - ExecutableRunOptions run_opts; - run_opts.set_allocator(client_->backend().memory_allocator()); - run_opts.set_intra_op_thread_pool( - client_->backend().eigen_intra_op_thread_pool_device()); - TF_ASSERT_OK_AND_ASSIGN(ScopedShapedBuffer result, - executable->Run({&input_data}, run_opts)); - - TF_ASSERT_OK_AND_ASSIGN(Literal result_literal, - client_->ShapedBufferToLiteral(result)); - - // We essentially reimplement LiteralTestUtil::Near here because - // a) this streamlined implementation is much faster, and - // b) we can print out better error messages (namely, we can print out - // which floating-point value input failed, while LiteralTestUtil::Near - // can only print out the input index that failed). - // c) we need special handling of certain inputs. For example, we say that - // a denormal input has multiple correct outputs (namely, f(x) and f(0)) - // and just needs to be close to one of them. - absl::Span result_arr = result_literal.data(); - ASSERT_EQ(result_arr.size(), input_arr.size()); - int64 mismatches = 0; - // Hoisting this out of the loop is a nice speedup on shards that have many - // denormals. - const float expected_at_zero = evaluate_op(0); - for (int64 i = 0; i < input_arr.size(); ++i) { - float input = ith_input_elem(i); - float actual = result_arr[i]; - float expected = evaluate_op(input); - if (IsClose(expected, actual)) { - continue; - } - - constexpr int64 kMaxMismatchesPrinted = 1000; - if (std::fpclassify(input) == FP_SUBNORMAL) { - // For denormal inputs, we accept answers that are close to either - // - evaluate_op(input) OR - // - evaluate_op(0). - if (IsClose(expected_at_zero, actual)) { - continue; - } - ++mismatches; - if (mismatches < kMaxMismatchesPrinted || VLOG_IS_ON(2)) { - // Use %0.9g because that's guaranteed to print an f32 to full - // precision. - LOG(ERROR) << absl::StreamFormat( - "Mismatch on denormal value %0.9g (0x%08x). Expected either " - "%0.9g (0x%08x) (evaluated at true value) or %0.9g (0x%08x) " - "(evaluated at zero), but got %0.9g (0x%08x).", - input, absl::bit_cast(input), // - expected, absl::bit_cast(expected), // - expected_at_zero, absl::bit_cast(expected_at_zero), - actual, absl::bit_cast(actual)); - } - } else { - mismatches++; - if (mismatches < kMaxMismatchesPrinted || VLOG_IS_ON(2)) { - LOG(ERROR) << absl::StreamFormat( - "Mismatch on %0.9g (0x%08x). Expected %0.9g (0x%08x), but got " - "%0.9g (0x%08x).", - input, absl::bit_cast(input), // - expected, absl::bit_cast(expected), // - actual, absl::bit_cast(actual)); - } - } - - if (mismatches == kMaxMismatchesPrinted && !VLOG_IS_ON(2)) { - LOG(ERROR) << "Not printing any more mismatches; pass " - "--vmodule=exhaustive_f32_elementwise_op_test=2 to see " - "all of them."; - } - } - EXPECT_EQ(mismatches, 0); - } -}; - -XLA_TEST_P(ExhaustiveF32ElementwiseOpTest, LogF32) { -#if !defined(XLA_TEST_BACKEND_CPU) && !defined(XLA_TEST_BACKEND_GPU) - error_spec_ = ErrorSpec{0.001, 0.001}; -#endif - ExhaustivelyTestF32Op( - [](XlaBuilder* builder, const XlaOp& input) { Log(input); }, std::log, - /*known_incorrect_range=*/{0, 0}); -} - -XLA_TEST_P(ExhaustiveF32ElementwiseOpTest, ExpF32) { -#ifdef XLA_TEST_BACKEND_CPU - // TODO(b/73142289): The vectorized Exp implementation gives results outside - // our error spec in this range (these numbers are bitwise representations of - // floats expressed as a zero extended int64): - std::pair known_incorrect_range = {1107296256 + 11583654, - 1107296256 + 11629080}; -#else - std::pair known_incorrect_range = {0, 0}; -#endif - - ExhaustivelyTestF32Op( - [](XlaBuilder* builder, const XlaOp& input) { Exp(input); }, std::exp, - known_incorrect_range); -} - -XLA_TEST_P(ExhaustiveF32ElementwiseOpTest, TanhF32) { - ExhaustivelyTestF32Op( - [](XlaBuilder* builder, const XlaOp& input) { Tanh(input); }, std::tanh, - /*known_incorrect_range=*/{0, 0}); -} - -XLA_TEST_P(ExhaustiveF32ElementwiseOpTest, ErfF32) { - ExhaustivelyTestF32Op( - [](XlaBuilder* builder, const XlaOp& input) { Erf(input); }, std::erf, - /*known_incorrect_range=*/{0, 0}); -} - -XLA_TEST_P(ExhaustiveF32ElementwiseOpTest, ErfcF32) { - ExhaustivelyTestF32Op( - [](XlaBuilder* builder, const XlaOp& input) { Erfc(input); }, std::erfc, - /*known_incorrect_range=*/{0, 0}); -} - -std::vector> CreateExhaustiveParameters() { - // We break up the 2^32-element space into small'ish chunks to keep peak - // memory usage low. - std::vector> result; - const int64 step = 1 << 25; - for (int64 i = 0; i < (1l << 32); i += step) { - result.push_back({i, i + step}); - } - return result; -} - -INSTANTIATE_TEST_CASE_P(ExhaustiveF32ElementwiseOpTestInstance, - ExhaustiveF32ElementwiseOpTest, - ::testing::ValuesIn(CreateExhaustiveParameters())); -} // namespace -} // namespace xla diff --git a/tensorflow/compiler/xla/tests/exhaustive_op_test.cc b/tensorflow/compiler/xla/tests/exhaustive_op_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..58bb9a217b805a142869149c19d7bcfc91a1aee1 --- /dev/null +++ b/tensorflow/compiler/xla/tests/exhaustive_op_test.cc @@ -0,0 +1,646 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include "absl/base/casts.h" +#include "tensorflow/compiler/xla/client/lib/constants.h" +#include "tensorflow/compiler/xla/client/lib/math.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/tests/client_library_test_base.h" +#include "tensorflow/compiler/xla/tests/literal_test_util.h" +#include "tensorflow/compiler/xla/tests/test_macros.h" + +namespace xla { +namespace { + +using Eigen::half; + +template +T EvaluatePolynomial(T x, const std::array& coeffs) { + T result = 0; + for (T c : coeffs) { + result = result * x + c; + } + return result; +} + +// There's no std::erfinv, so we have to implement it ourselves. This follows +// Wichura 1998, https://www.jstor.org/stable/2347330 which, notably, is a +// different implementation from that in math.cc. +float HostErfInv(float x) { + std::array kPolyA = { + 8.8709406962545514830200e2, 1.1819493347062294404278e4, + 2.3782041382114385731252e4, 1.6235862515167575384252e4, + 4.8548868893843886794648e3, 6.9706266534389598238465e2, + 4.7072688112383978012285e1, 1.1975323115670912564578e0, + }; + std::array kPolyB = { + 5.2264952788528545610e3, 2.8729085735721942674e4, 3.9307895800092710610e4, + 2.1213794301586595867e4, 5.3941960214247511077e3, 6.8718700749205790830e2, + 4.2313330701600911252e1, 1.0000000000000000000e0, + }; + std::array kPolyC = { + 7.74545014278341407640e-4, 2.27238449892691845833e-2, + 2.41780725177450611770e-1, 1.27045825245236838258e0, + 3.64784832476320460504e0, 5.76949722146069140550e0, + 4.63033784615654529590e0, 1.42343711074968357734e0, + }; + std::array kPolyD = { + 1.4859850019840355905497876e-9, 7.7441459065157709165577218e-4, + 2.1494160384252876777097297e-2, 2.0945065210512749128288442e-1, + 9.7547832001787427186894837e-1, 2.3707661626024532365971225e0, + 2.9036514445419946173133295e0, 1.4142135623730950488016887e0, + }; + std::array kPolyE = { + 2.01033439929228813265e-7, 2.71155556874348757815e-5, + 1.24266094738807843860e-3, 2.65321895265761230930e-2, + 2.96560571828504891230e-1, 1.78482653991729133580e0, + 5.46378491116411436990e0, 6.65790464350110377720e0, + }; + std::array kPolyF = { + 2.891024605872965461538222e-15, 2.010321207683943062279931e-7, + 2.611088405080593625138020e-5, 1.112800997078859844711555e-3, + 2.103693768272068968719679e-2, 1.936480946950659106176712e-1, + 8.482908416595164588112026e-1, 1.414213562373095048801689e0, + }; + + if (std::abs(x) > 1 || std::isnan(x)) { + return std::numeric_limits::quiet_NaN(); + } + if (std::abs(x) == 1) { + return std::copysign(std::numeric_limits::infinity(), x); + } + + float unsigned_result = [&] { + float y = std::abs(x); + if (y <= 0.85) { + double r = 0.180625 - 0.25 * y * y; + return (y * EvaluatePolynomial(r, kPolyA)) / + EvaluatePolynomial(r, kPolyB); + } else { + double r = std::sqrt(std::log(2.0) - std::log1p(-y)); + if (r <= 5.0) { + r -= 1.6; + return EvaluatePolynomial(r, kPolyC) / EvaluatePolynomial(r, kPolyD); + } else { + r -= 5; + return EvaluatePolynomial(r, kPolyE) / EvaluatePolynomial(r, kPolyF); + } + } + }(); + return std::copysign(unsigned_result, x); +} + +// Digamma implementation using a polynomial from Cephes. Notably this is a +// different implementation from the one in math.cc. +float HostDigamma(float x) { + // Euler-Mascheroni constant + float kGamma = 0.57721566490153286061; + float kPi = M_PI; + + std::array kPoly = { + -4.16666666666666666667E-3, + 3.96825396825396825397E-3, + -8.33333333333333333333E-3, + 8.33333333333333333333E-2, + }; + + float reflection = 0; + if (x <= 0) { + float floor = std::floor(x); + if (x == floor) { + return std::numeric_limits::quiet_NaN(); + } + // Compute reflection term, pi * cot(pi * x). + reflection = x - floor; + if (reflection == 0.5) { + reflection = 0; + } else { + if (reflection > 0.5) { + reflection = x - (floor + 1.0f); + } + reflection = kPi / std::tan(kPi * reflection); + } + x = 1 - x; + } + + float result = 0; + if (x <= 10 && x == std::floor(x)) { + // Special case for integers <= 10. + for (int i = 1; i < x; ++i) { + result += 1.0f / i; + } + result -= kGamma; + } else { + float w = 0; + for (; x < 10; ++x) { + w += 1.0f / x; + } + if (x < 1e8) { + float z = 1.0f / (x * x); + result = z * EvaluatePolynomial(z, kPoly); + } + result = std::log(x) - 0.5f / x - result - w; + } + + // Compute the final, reflected value. + return result - reflection; +} + +// For f32, f16, and bf16, we need 9, 5, and 4 decimal places of precision to be +// guaranteed that we're printing the full number. +// +// (The general formula is, given a floating-point number with S significand +// bits, the number of decimal digits needed to print it to full precision is +// +// ceil(1 + S * log_10(2)) ~= ceil(1 + S * 0.30103). +// +// See https://people.eecs.berkeley.edu/~wkahan/Math128/BinDecBin.pdf.) +string StringifyNum(float x) { + return absl::StrFormat("%0.9g (0x%08x)", x, absl::bit_cast(x)); +} + +string StringifyNum(half x) { + return absl::StrFormat("%0.5g (0x%04x)", static_cast(x), + absl::bit_cast(x)); +} + +string StringifyNum(bfloat16 x) { + return absl::StrFormat("%0.4g (0x%04x)", static_cast(x), + absl::bit_cast(x)); +} + +// Test parameter is a tuple containing +// - primitive type under test, +// - (begin, end) range under test, as zero-extended int64s bitcast to the +// primtive type under test. +class ExhaustiveOpTest + : public ClientLibraryTestBase, + public ::testing::WithParamInterface< + std::tuple>> { + public: + ExhaustiveOpTest() + : ty_(std::get<0>(GetParam())), platform_(client_->platform()->Name()) {} + + void Run(std::function enqueue_op, + float (*evaluate_op)(float)) { + SetFastMathDisabled(true); + + // Run all HLO passes. In particular, constant folding is disabled by + // default for tests, but we need to run it in order to tickle some bugs. + mutable_debug_options()->clear_xla_disable_hlo_passes(); + + PrimitiveType ty; + std::tie(ty, std::ignore) = GetParam(); + + switch (ty) { + case F32: + SetDefaultErrSpec(0.0001, 0.0001); + RunImpl(enqueue_op, evaluate_op); + break; + case F16: + SetDefaultErrSpec(0.001, 0.001); + RunImpl(enqueue_op, evaluate_op); + break; + case BF16: + SetDefaultErrSpec(0.001, 0.01); + RunImpl(enqueue_op, evaluate_op); + break; + default: + LOG(FATAL) << "Unhandled type."; + } + } + + void SetDefaultErrSpec(float abs_err, float rel_err) { + if (!abs_err_.has_value()) { + abs_err_ = abs_err; + } + if (!rel_err_.has_value()) { + rel_err_ = rel_err; + } + } + + template + void RunImpl(std::function enqueue_op, + float (*evaluate_op)(float)) { + static_assert( + sizeof(T) == sizeof(IntegralT), + "IntegralT must be an unsigned integer type of the same width as T."); + + PrimitiveType ty; + std::pair test_range; + std::tie(ty, test_range) = GetParam(); + int64 begin, end; + std::tie(begin, end) = test_range; + + if (begin >= known_incorrect_begin_ && end <= known_incorrect_end_) { + LOG(INFO) << absl::StreamFormat( + "Skipping this shard, as the range under test, [%d, %d), falls " + "entirely within the known-incorrect range [%d, %d).", + begin, end, known_incorrect_begin_, known_incorrect_end_); + return; + } + + LOG(INFO) << "Checking range [" << begin << ", " << end << ")"; + + int64 input_size = end - begin; + Literal input_literal = LiteralUtil::CreateFromDimensions(ty, {input_size}); + absl::Span input_arr = input_literal.data(); + for (int64 i = 0; i < input_size; i++) { + IntegralT input_val = i + begin; + // If the operation is known to be buggy on a specific input clamp that + // input to 0 under the assumption that the op is at least correct on 0. + if (input_val >= known_incorrect_begin_ && + input_val < known_incorrect_end_) { + input_arr[i] = T{0}; + } else { + input_arr[i] = absl::bit_cast(input_val); + } + } + + TF_ASSERT_OK_AND_ASSIGN(Literal result_literal, + BuildAndRunComputation(enqueue_op, input_literal)); + ExpectNear(input_literal, result_literal, evaluate_op); + } + + StatusOr BuildAndRunComputation( + const std::function& enqueue_op, + const Literal& input_literal) { + XlaBuilder builder(TestName()); + auto input = Parameter(&builder, 0, input_literal.shape(), "input"); + enqueue_op(input); + TF_ASSIGN_OR_RETURN(XlaComputation comp, builder.Build()); + + // Build and run the computation using the LocalClient API, rather than the + // plain Client API, which is used by ClientLibraryTestBase. This is + // because the plain Client API results does more memcpys to/from Literals, + // and that's slow given that we're touching a lot of data here. + // + // Copy debug options from ClientLibraryTestBase. In particular, we're + // interested in disabling constant folding. + ExecutableBuildOptions build_opts; + *build_opts.mutable_debug_options() = *mutable_debug_options(); + TF_ASSIGN_OR_RETURN( + auto executable, + client_->Compile(comp, {&input_literal.shape()}, build_opts)); + + TF_ASSIGN_OR_RETURN( + ScopedShapedBuffer input_data, + client_->LiteralToShapedBuffer(input_literal, /*device_ordinal=*/0)); + + ExecutableRunOptions run_opts; + run_opts.set_allocator(client_->backend().memory_allocator()); + run_opts.set_intra_op_thread_pool( + client_->backend().eigen_intra_op_thread_pool_device()); + TF_ASSIGN_OR_RETURN(ScopedShapedBuffer result, + executable->Run({&input_data}, run_opts)); + + TF_ASSIGN_OR_RETURN(Literal result_literal, + client_->ShapedBufferToLiteral(result)); + return std::move(result_literal); + } + + template + bool IsClose(T expected, T actual) { + float expected_f32 = static_cast(expected); + float actual_f32 = static_cast(actual); + float abs_err = std::abs(expected_f32 - actual_f32); + float rel_err = abs_err / std::abs(expected_f32); + if (strict_signed_zeros_ && actual == T{0} && expected == T{0}) { + // Check sign of zero. + return std::signbit(actual_f32) == std::signbit(expected_f32); + } + return abs_err < *abs_err_ || rel_err < *rel_err_ || + (std::isnan(expected_f32) && std::isnan(actual_f32)) || + (std::isinf(expected_f32) && std::isinf(actual_f32) && + (expected_f32 > 0) == (actual_f32 > 0)); + } + + template + void ExpectNear(const Literal& input_literal, const Literal& result_literal, + float (*evaluate_op)(float)) { + // We essentially reimplement LiteralTestUtil::Near here because + // a) this streamlined implementation is much faster, and + // b) we can print out better error messages (namely, we can print out + // which floating-point value input failed, while LiteralTestUtil::Near + // can only print out the input index that failed). + // c) we need special handling of certain inputs. For example, we say that + // a denormal input has multiple correct outputs (namely, f(x) and f(0)) + // and just needs to be close to one of them. + absl::Span input_arr = input_literal.data(); + absl::Span result_arr = result_literal.data(); + ASSERT_EQ(result_arr.size(), input_arr.size()); + int64 mismatches = 0; + // Hoisting these out of the loop is a nice speedup on shards that have many + // denormals. + const T expected_at_pos_zero = static_cast(evaluate_op(0)); + const T expected_at_neg_zero = static_cast(evaluate_op(-0.0)); + for (int64 i = 0; i < input_arr.size(); ++i) { + T input = input_arr[i]; + float input_f32 = static_cast(input); + T actual = result_arr[i]; + T expected = static_cast(evaluate_op(input_f32)); + + if (IsClose(expected, actual)) { + continue; + } + + // Easy case: If `input` is not denormal and !IsClose(expected, actual), + // print an error. + // + // (This doesn't correctly detect f16 and bfloat16 denormals! This seems + // to be OK for now, but at some point we may need to implement fpclassify + // for half and bfloat.) + if (std::fpclassify(input_f32) != FP_SUBNORMAL) { + PrintMismatch(&mismatches, [&] { + return absl::StrFormat("Mismatch on %s. Expected %s, but got %s.", + StringifyNum(input), StringifyNum(expected), + StringifyNum(actual)); + }); + continue; + } + + // Otherwise, `input` is denormal. For denormal inputs, we accept answers + // that are close to any of: + // + // - evaluate_op(input) + // - evaluate_op(+/-0), where the sign of 0 equal to the sign of + // `input`, + // - if relaxed_denormal_signs_, evaluate_op(-/+0), where the sign of + // 0 is the opposite of `input`. + T sign_preserving_ftz_expected = + std::signbit(input_f32) ? expected_at_neg_zero : expected_at_pos_zero; + T sign_nonpreserving_ftz_expected = + std::signbit(input_f32) ? expected_at_pos_zero : expected_at_neg_zero; + if (IsClose(sign_preserving_ftz_expected, actual) || + (relaxed_denormal_signs_ && + IsClose(sign_nonpreserving_ftz_expected, actual))) { + continue; + } + + if (relaxed_denormal_signs_) { + PrintMismatch(&mismatches, [&] { + return absl::StrFormat( + "Mismatch on denormal value %s. Expected one of:\n" + " %10s (evaluated at full-precision value)\n" + " %10s (evaluated after flushing to sign-preserving zero)\n" + " %10s (evaluated after flushing to non-sign-preserving " + "zero)\n" + "but got %s.", + StringifyNum(input), StringifyNum(expected), + StringifyNum(sign_preserving_ftz_expected), + StringifyNum(sign_nonpreserving_ftz_expected), + StringifyNum(actual)); + }); + } else { + PrintMismatch(&mismatches, [&] { + return absl::StrFormat( + "Mismatch on denormal value %s. Expected one of:\n" + " %10s (evaluated at full-precision value)\n" + " %10s (evaluated after flushing to sign-preserving zero)\n" + "but got %s.", + StringifyNum(input), StringifyNum(expected), + StringifyNum(sign_preserving_ftz_expected), StringifyNum(actual)); + }); + } + } + EXPECT_EQ(mismatches, 0); + } + + template + void PrintMismatch(int64* mismatches, const ErrorGenerator& err_generator) { + // We send a few mismatches to gunit so they show up nicely in test logs. + // Then we send more to LOG(ERROR). The remainder we squelch unless we're + // at vlog level 2. + constexpr int64 kMaxMismatchesLoggedToGunit = 10; + constexpr int64 kMaxMismatchesLoggedToErr = 1000; + + (*mismatches)++; + if (*mismatches < kMaxMismatchesLoggedToGunit) { + FAIL() << err_generator(); + } else if (*mismatches < kMaxMismatchesLoggedToErr || VLOG_IS_ON(2)) { + LOG(ERROR) << err_generator(); + } else if (*mismatches == kMaxMismatchesLoggedToErr) { + LOG(ERROR) << "Not printing any more mismatches; pass " + "--vmodule=exhaustive_f32__op_test=2 to see " + "all of them."; + } + } + + // The following members are set during construction so testcases can read + // these values and use them e.g. to influence the values given to the mutable + // members below. + + // The primitive type under test. + const PrimitiveType ty_; + + // The platform under test. + const string platform_; + + // Tests can set the following variables for control over execution. This is + // safe because each XLA_TEST_P instantiates a new instance of this class. + + // Testing will ignore the given range (encoded as bitwise representations of + // the type under test zero-extended to int64). + int64 known_incorrect_begin_ = 0; + int64 known_incorrect_end_ = 0; + + // If unset, reasonable defaults will be used depending on the type under + // test. + absl::optional abs_err_; + absl::optional rel_err_; + + // If true, will consider -0 not near to +0 and vice versa. Note that + // +epsilon may still be considered close to -0, depending on the error spec; + // this only covers the case when both `expected` and `actual` are equal to 0. + bool strict_signed_zeros_ = false; + + // If true, allows denormals to be flushed to non-sign-preserving 0. + // + // For example, normally we'd expect sqrt(-denormal) to be either nan (sqrt of + // a negative number) or -inf (flush the denormal to sign-perserving zero, + // then sqrt(-0)). But with this as true, we'll also accept 0 (sqrt(0)). + // + // XLA:GPU preserves denormal signs, but other backends don't. + bool relaxed_denormal_signs_ = platform_ != "CUDA"; +}; + +XLA_TEST_P(ExhaustiveOpTest, Log) { + if (platform_ != "Host" && platform_ != "CUDA" && ty_ == F32) { + abs_err_ = 0.001; + rel_err_ = 0.001; + } + + Run(Log, std::log); +} + +XLA_TEST_P(ExhaustiveOpTest, Log1p) { + if (platform_ != "Host" && platform_ != "CUDA" && ty_ == F32) { + abs_err_ = 0.001; + rel_err_ = 0.001; + } + + Run(Log1p, std::log1p); +} + +XLA_TEST_P(ExhaustiveOpTest, Exp) { + if (platform_ == "Host" && ty_ == F32) { + // TODO(b/73142289): The vectorized Exp implementation gives results outside + // our error spec in this range. + known_incorrect_begin_ = 1107296256 + 11583654; + known_incorrect_end_ = 1107296256 + 11629080; + } else if (platform_ == "Host" && ty_ == BF16) { + // TODO(jlebar): Is this a rounding error? Why doesn't it occur on XLA:GPU? + // + // Mismatch on 88.5 (0x42b1). + // Expected 2.72491739e+38 (0x7f4d), but got inf (0x7f80). + known_incorrect_begin_ = 0x42b1; + known_incorrect_end_ = 0x42b2; + } + + Run(Exp, std::exp); +} + +XLA_TEST_P(ExhaustiveOpTest, Expm1) { + // Expm1 has the same erroneous behavior on CPU as Exp. + if (platform_ == "Host" && ty_ == F32) { + // TODO(b/73142289): The vectorized Exp implementation gives results outside + // our error spec in this range. + known_incorrect_begin_ = 1107296256 + 11583654; + known_incorrect_end_ = 1107296256 + 11629080; + } else if (platform_ == "Host" && ty_ == BF16) { + // TODO(jlebar): Is this a rounding error? Why doesn't it occur on XLA:GPU? + // + // Mismatch on 88.5 (0x42b1). + // Expected 2.72491739e+38 (0x7f4d), but got inf (0x7f80). + known_incorrect_begin_ = 0x42b1; + known_incorrect_end_ = 0x42b2; + } + + Run(Expm1, std::expm1); +} + +// It feels a little overkill to exhaustively test sqrt and pow(x, 0.5), but +// this *did* find a bug, namely that some backends were assuming sqrt(x) == +// pow(x, 0.5), but this is not true for x == -inf. +XLA_TEST_P(ExhaustiveOpTest, PowOneHalf) { + Run([](XlaOp x) { return Pow(x, ScalarLike(x, 0.5)); }, + +[](float x) { return std::pow(x, 0.5f); }); +} + +XLA_TEST_P(ExhaustiveOpTest, Rsqrt) { + Run( + Rsqrt, +[](float x) { return 1 / std::sqrt(x); }); +} + +XLA_TEST_P(ExhaustiveOpTest, Sqrt) { + if (platform_ == "Host" || platform_ == "CUDA") { + strict_signed_zeros_ = true; + } + + Run(Sqrt, std::sqrt); +} + +// TODO(jlebar): Add remaining trig functions. Don't forget Atan2! +// TODO(jlebar): Test trig functions over complex inputs. +XLA_TEST_P(ExhaustiveOpTest, Tanh) { Run(Tanh, std::tanh); } + +XLA_TEST_P(ExhaustiveOpTest, Erf) { Run(Erf, std::erf); } +XLA_TEST_P(ExhaustiveOpTest, Erfc) { Run(Erfc, std::erfc); } +XLA_TEST_P(ExhaustiveOpTest, ErfInv) { Run(ErfInv, HostErfInv); } +XLA_TEST_P(ExhaustiveOpTest, Digamma) { + if (platform_ != "Host" && platform_ != "CUDA") { + // TODO(b/123956399): This is a fairly high error, significantly higher than + // we see on CPU/GPU. + rel_err_ = 0.01; + abs_err_ = 0.01; + } + + if (platform_ == "CUDA") { + // On GPU we get a wrong answer for the denormal inputs +/-2.93873588e-39 + // (0x00200000 and 0x80200000). These should return -/+inf (at least + // according to our reference implementation!) but XLA:GPU returns + // -/+3.40282326e+38 (0xff7ffffe and 0x7f7ffffe). + // + // I deem this an acceptable result, as XLA:GPU flushes denormals, and as + // the results we get here are very close to MAX_FLOAT. We just hardcode + // these results, as this is better than ignoring these inputs altogether. + auto host_digamma_with_gpu_ftz_errors = +[](float x) { + if (absl::bit_cast(x) == 0x00200000 || + absl::bit_cast(x) == 0x80200000) { + return std::copysign(std::numeric_limits::max(), -x); + } + return HostDigamma(x); + }; + Run(Digamma, host_digamma_with_gpu_ftz_errors); + } else { + Run(Digamma, HostDigamma); + } +} +XLA_TEST_P(ExhaustiveOpTest, Lgamma) { + // Our implementation gets within 0.0001 rel error except for ~20 denormal + // inputs on GPU. Anyway 0.001 rel error should be good enough for lgamma. + if (platform_ == "CUDA" && (ty_ == F32 || ty_ == F16)) { + rel_err_ = 0.001; + } + if (platform_ != "Host" && platform_ != "CUDA") { + // TODO(b/123956399): This is a fairly high error, significantly higher than + // we see on CPU/GPU. + rel_err_ = 0.01; + abs_err_ = 0.01; + + // Overflows for to inf for input 4.08500343e+36 (0x7c44af8e). + if (ty_ == F32) { + known_incorrect_begin_ = 0x7c44af8e; + known_incorrect_end_ = 0x7c44af8e + 1; + } + } + Run(Lgamma, std::lgamma); +} + +XLA_TEST_P(ExhaustiveOpTest, Round) { Run(Round, std::round); } + +std::vector> CreateExhaustiveF32Ranges() { + // We break up the 2^32-element space into small'ish chunks to keep peak + // memory usage low. + std::vector> result; + const int64 step = 1 << 25; + for (int64 i = 0; i < (1l << 32); i += step) { + result.push_back({i, i + step}); + } + return result; +} + +INSTANTIATE_TEST_SUITE_P( + F32, ExhaustiveOpTest, + ::testing::Combine(::testing::Values(F32), + ::testing::ValuesIn(CreateExhaustiveF32Ranges()))); + +#if !defined(XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16) +INSTANTIATE_TEST_SUITE_P( + F16, ExhaustiveOpTest, + ::testing::Combine(::testing::Values(F16), + ::testing::Values(std::make_pair(0, 1 << 16)))); +#endif + +#if defined(XLA_BACKEND_SUPPORTS_BFLOAT16) +INSTANTIATE_TEST_SUITE_P( + BF16, ExhaustiveOpTest, + ::testing::Combine(::testing::Values(BF16), + ::testing::Values(std::make_pair(0, 1 << 16)))); +#endif + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.cc b/tensorflow/compiler/xla/tests/hlo_test_base.cc index 0151981ef16aabe9e363bc4d7f9ba96d4a1f170f..62e2b465cfe67e9504cea96abf2d9ddfd2ce1a54 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.cc +++ b/tensorflow/compiler/xla/tests/hlo_test_base.cc @@ -207,13 +207,14 @@ Literal HloTestBase::ExecuteAndTransfer(std::unique_ptr module, StatusOr> HloTestBase::ExecuteReplicated( std::unique_ptr module, absl::Span arguments, - int64 num_replicas) { + int64 num_replicas, bool use_threads) { HloRunner::ReplicatedExecuteOptions options; options.num_replicas = num_replicas; for (auto argument : arguments) { options.arguments.push_back(argument); } - return test_runner_.ExecuteReplicated(std::move(module), options); + return test_runner_.ExecuteReplicated(std::move(module), options, + use_threads); } StatusOr> HloTestBase::MakeReferenceModule( diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.h b/tensorflow/compiler/xla/tests/hlo_test_base.h index 3c2bcbb5df5ce94dd37f63d0c0e609f3ad2b60aa..df9c29a186f36a538f71082e7d2ca6509e0758b3 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.h +++ b/tensorflow/compiler/xla/tests/hlo_test_base.h @@ -174,9 +174,13 @@ class HloTestBase : public ::testing::Test { absl::Span arguments); // Executes the given module on multiple replicas. + // + // use_threads indicates whether this replicated computation will be executed + // with a thread-per-replica, vs using an implicitly async call such as + // Executable::ExecuteOnStreams. StatusOr> ExecuteReplicated( std::unique_ptr module, absl::Span arguments, - int64 num_replicas); + int64 num_replicas, bool use_threads); // Executes the given hlo module on two backends and compares results. // diff --git a/tensorflow/compiler/xla/tests/multi_device_all_reduce_test.cc b/tensorflow/compiler/xla/tests/multi_device_all_reduce_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..1513d89ba9c95b3097229b268d22832dee3e98cd --- /dev/null +++ b/tensorflow/compiler/xla/tests/multi_device_all_reduce_test.cc @@ -0,0 +1,56 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/compiler/xla/service/hlo_parser.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/tests/test_macros.h" + +namespace xla { +namespace { + +class MultiDeviceAllReduceTest : public HloTestBase {}; + +XLA_TEST_F(MultiDeviceAllReduceTest, TwoReplicasOneOperand) { + const char* module_str = R"( + HloModule test + + add { + x = f32[] parameter(0) + y = f32[] parameter(1) + add = f32[] add(x, y) + } + + ENTRY test_computation { + p = f32[3] parameter(0) + ROOT crs = f32[3] all-reduce(p), to_apply=add + })"; + auto config = GetModuleConfigForTest(); + config.set_replica_count(2); + auto module = ParseHloString(module_str, config).ValueOrDie(); + auto literal = LiteralUtil::CreateR1({1, 2, 3}); + auto expected = LiteralUtil::CreateR1({2, 4, 6}); + TF_ASSERT_OK_AND_ASSIGN(std::vector results, + ExecuteReplicated(std::move(module), {&literal}, 2, + /*use_threads=*/true)); + EXPECT_EQ(expected, results[0]); + EXPECT_EQ(expected, results[1]); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/test_utils.cc b/tensorflow/compiler/xla/tests/test_utils.cc index 67d2258928f75c078588c9425359f9468f4463ed..4ac3dbd80cfaf2340d8f79cef3e9e02058cf919c 100644 --- a/tensorflow/compiler/xla/tests/test_utils.cc +++ b/tensorflow/compiler/xla/tests/test_utils.cc @@ -112,6 +112,31 @@ void PopulateWithFloatingPointData(Literal* literal, std::minstd_rand0* engine, } } +template +void PopulateWithComplexData(Literal* result, std::minstd_rand0* engine, + bool no_duplicates) { + using InnerFloatT = typename ComplexT::value_type; + CHECK(engine != nullptr); + CHECK_EQ(result->shape().element_type(), + primitive_util::NativeToPrimitiveType()); + Shape floating_point_shape = ShapeUtil::ChangeElementType( + result->shape(), primitive_util::NativeToPrimitiveType()); + Literal real_lit(floating_point_shape); + Literal imaginary_lit(floating_point_shape); + + PopulateWithFloatingPointData(&real_lit, engine, no_duplicates); + PopulateWithFloatingPointData(&imaginary_lit, engine, + no_duplicates); + + absl::Span real_data = real_lit.data(); + absl::Span imaginary_data = + imaginary_lit.data(); + absl::Span result_data = result->data(); + for (int i = 0; i < real_lit.data().size(); i++) { + result_data[i] = ComplexT(real_data[i], imaginary_data[i]); + } +} + template <> void PopulateWithFloatingPointData(Literal* literal, std::minstd_rand0* engine, @@ -220,6 +245,12 @@ StatusOr MakeFakeLiteralInternal(const Shape& shape, case U64: PopulateWithRandomIntegralData(&literal, engine, no_duplicates); break; + case C64: + PopulateWithComplexData(&literal, engine, no_duplicates); + break; + case C128: + PopulateWithComplexData(&literal, engine, no_duplicates); + break; case PRED: { std::uniform_int_distribution generator(0, 1); TF_CHECK_OK( diff --git a/tensorflow/compiler/xla/util.h b/tensorflow/compiler/xla/util.h index f2fd17dc99455a921bf875aad2a3661b4d456823..1754ae0e44f3420bf7eb7cfb3b558dd476b31455 100644 --- a/tensorflow/compiler/xla/util.h +++ b/tensorflow/compiler/xla/util.h @@ -260,6 +260,16 @@ Status Unavailable(const absl::FormatSpec& format, return WithLogBacktrace( tensorflow::errors::Unavailable(absl::StrFormat(format, args...))); } +template +Status Unknown(const absl::FormatSpec& format, const Args&... args) { + return WithLogBacktrace( + tensorflow::errors::Unknown(absl::StrFormat(format, args...))); +} +template +Status Internal(const absl::FormatSpec& format, const Args&... args) { + return WithLogBacktrace( + tensorflow::errors::Internal(absl::StrFormat(format, args...))); +} template Status InvalidArgumentStrCat(Args&&... concat) { diff --git a/tensorflow/compiler/xla/xla_data.proto b/tensorflow/compiler/xla/xla_data.proto index 226299a7186ef0acb41f6d01fdeffeee06f13d4d..6e5772a7396bae1674ec4e7393ba03506c9381e4 100644 --- a/tensorflow/compiler/xla/xla_data.proto +++ b/tensorflow/compiler/xla/xla_data.proto @@ -16,6 +16,7 @@ limitations under the License. syntax = "proto3"; package xla; + option cc_enable_arenas = true; // Primitive types are the individual values that can be held in rectangular @@ -55,7 +56,7 @@ enum PrimitiveType { F64 = 12; // Complex values of fixed width. - C64 = 15; // Paired F32 (real, imag), as in std::complex. + C64 = 15; // Paired F32 (real, imag), as in std::complex. C128 = 18; // Paired F64 (real, imag), as in std::complex. // A tuple is a polymorphic sequence; e.g. a shape that holds different @@ -199,7 +200,7 @@ message ShapeProto { // in this field represents an upper bound on the size of the dimension. repeated int64 dimensions = 3; - // For tuples only, the shapes of constitutent shapes in the tuple sequence. + // For tuples only, the shapes of constituent shapes in the tuple sequence. repeated ShapeProto tuple_shapes = 4; // The layout used to back this shape. @@ -367,7 +368,7 @@ message LiteralProto { repeated uint64 u64s = 7; repeated float f32s = 8; repeated double f64s = 9; - repeated float c64s = 12; // Stored as interleaved real, imag floats. + repeated float c64s = 12; // Stored as interleaved real, imag floats. repeated double c128s = 18; // Stored as interleaved real, imag doubles. repeated LiteralProto tuple_literals = 10; // The F16s, BF16s, U16s and S16s are encoded in little endian byte order @@ -510,7 +511,7 @@ message ConvolutionDimensionNumbers { repeated int64 output_spatial_dimensions = 12; // Next = 13 -}; +} enum FftType { FFT = 0; // Forward FFT; complex in, complex out. @@ -529,7 +530,7 @@ message DotDimensionNumbers { repeated int64 lhs_batch_dimensions = 3; // The dimension numbers that represent the 'rhs' batch dimensions. repeated int64 rhs_batch_dimensions = 4; -}; +} enum RandomDistribution { RNG_INVALID = 0; @@ -565,6 +566,12 @@ message TriangularSolveOptions { Transpose transpose_a = 4; } +message CholeskyOptions { + // If true, uses the lower triangle of `a`. If false, uses the upper triangle + // of `a`. + bool lower = 1; +} + message OpSharding { enum Type { // This sharding is replicated across all devices (implies maximal, @@ -636,3 +643,20 @@ message ParameterReplication { // the HLO instruction's shape. repeated bool replicated_at_leaf_buffers = 1; } + +// A backend-config for kWhile loops that stores the loop's trip count, if it is +// known. +// +// This is useful for backends that can implement a `for i in 0..N` loop more +// efficiently than a `while` loop. For example, on GPUs, we can implement a +// `for i in 0..N` loop by enqueueing the kernels for the loop body N times, +// whereas implementing a `while` loop requires a host-device sync on each +// iteration. +message WhileLoopBackendConfig { + message KnownTripCount { + int64 n = 1; + } + // This indirection lets us distinguish between known-trip-count == 0 and + // unknown-trip-count. + KnownTripCount known_trip_count = 1; +} diff --git a/tensorflow/compiler/xrt/ops/xrt_state_ops.cc b/tensorflow/compiler/xrt/ops/xrt_state_ops.cc index 8832270fb2730d1ba64fa069b38f4a04b61773ef..87546fce4e4e7e38ef934d32ff95a60a4ad5492a 100644 --- a/tensorflow/compiler/xrt/ops/xrt_state_ops.cc +++ b/tensorflow/compiler/xrt/ops/xrt_state_ops.cc @@ -53,7 +53,7 @@ The shapes can differ from the corresponding input one, as long as the total number of elements matches. In other words, it is possible to feed an input tensor with shape {8} and have a corresponding shape {2,2,2}. layouts: A vector holding the requested layout in minor-to-major sequence. -If empty, the default layout wil be used. +If empty, the default layout will be used. For a tuple, the layouts vector holds a linearized minor-to-major numbers for all the tuple leaves, in the order they appear within the tuple. The elements within the layouts sequence corresponding to a given tuple diff --git a/tensorflow/contrib/bigtable/kernels/bigtable_lib.cc b/tensorflow/contrib/bigtable/kernels/bigtable_lib.cc index 98f906408c230a4382ffafe412ee9990d4384930..a81c3ec5c239789e37badf0f9629fab6db450b34 100644 --- a/tensorflow/contrib/bigtable/kernels/bigtable_lib.cc +++ b/tensorflow/contrib/bigtable/kernels/bigtable_lib.cc @@ -27,7 +27,7 @@ Status GrpcStatusToTfStatus(const ::grpc::Status& status) { status.error_code() == ::grpc::StatusCode::OUT_OF_RANGE) { grpc_code = ::grpc::StatusCode::INTERNAL; } - return Status(static_cast<::tensorflow::error::Code>(status.error_code()), + return Status(static_cast<::tensorflow::error::Code>(grpc_code), strings::StrCat("Error reading from Cloud Bigtable: ", status.error_message())); } diff --git a/tensorflow/contrib/bigtable/python/ops/bigtable_api.py b/tensorflow/contrib/bigtable/python/ops/bigtable_api.py index fa64055dfd65a134afdf46cebccb7f7d96106502..736cf3da49e934d49d0587d729cff6eaaed8f254 100644 --- a/tensorflow/contrib/bigtable/python/ops/bigtable_api.py +++ b/tensorflow/contrib/bigtable/python/ops/bigtable_api.py @@ -475,15 +475,17 @@ class BigtableTable(object): """ if timestamp is None: timestamp = -1 # Bigtable server provided timestamp. - for tensor_type in nest.flatten(dataset.output_types): + for tensor_type in nest.flatten( + dataset_ops.get_legacy_output_types(dataset)): if tensor_type != dtypes.string: raise ValueError("Not all elements of the dataset were `tf.string`") - for shape in nest.flatten(dataset.output_shapes): + for shape in nest.flatten(dataset_ops.get_legacy_output_shapes(dataset)): if not shape.is_compatible_with(tensor_shape.scalar()): raise ValueError("Not all elements of the dataset were scalars") if len(column_families) != len(columns): raise ValueError("len(column_families) != len(columns)") - if len(nest.flatten(dataset.output_types)) != len(columns) + 1: + if len(nest.flatten( + dataset_ops.get_legacy_output_types(dataset))) != len(columns) + 1: raise ValueError("A column name must be specified for every component of " "the dataset elements. (e.g.: len(columns) != " "len(dataset.output_types))") diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/estimator.py b/tensorflow/contrib/boosted_trees/estimator_batch/estimator.py index a178820841c4c8bcb7f5742babdb6d0f4825de31..5ffbb9067081d7440ab5e11290697b822051bee5 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/estimator.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/estimator.py @@ -84,12 +84,10 @@ class GradientBoostedDecisionTreeClassifier(estimator.Estimator): output_leaf_index: whether to output leaf indices along with predictions during inference. The leaf node indexes are available in predictions dict by the key 'leaf_index'. It is a Tensor of rank 2 and its shape is - [batch_size, num_trees]. - For example, - result_iter = classifier.predict(...) - for result_dict in result_iter: - # access leaf index list by result_dict["leaf_index"] - # which contains one leaf index per tree + [batch_size, num_trees]. For example, result_iter = + classifier.predict(...) + for result_dict in result_iter: # access leaf index list by + result_dict["leaf_index"] # which contains one leaf index per tree override_global_step_value: If after the training is done, global step value must be reset to this value. This should be used to reset global step to a number > number of steps used to train the current ensemble. @@ -179,8 +177,8 @@ class GradientBoostedDecisionTreeRegressor(estimator.Estimator): `[batch_size, label_dimension]`). num_trees: An int, number of trees to build. feature_columns: A list of feature columns. - label_name: String, name of the key in label dict. Can be null if label - is a tensor (single headed models). + label_name: String, name of the key in label dict. Can be null if label is + a tensor (single headed models). weight_column_name: Name of the column for weights, or None if not weighted. model_dir: Directory for model exports, etc. @@ -195,11 +193,11 @@ class GradientBoostedDecisionTreeRegressor(estimator.Estimator): opposed to contrib) version of tensorflow. output_leaf_index: whether to output leaf indices along with predictions during inference. The leaf node indexes are available in predictions - dict by the key 'leaf_index'. For example, - result_dict = classifier.predict(...) - for example_prediction_result in result_dict: - # access leaf index list by example_prediction_result["leaf_index"] - # which contains one leaf index per tree + dict by the key 'leaf_index'. For example, result_dict = + classifier.predict(...) + for example_prediction_result in result_dict: # access leaf index list + by example_prediction_result["leaf_index"] # which contains one leaf + index per tree override_global_step_value: If after the training is done, global step value must be reset to this value. This should be used to reset global step to a number > number of steps used to train the current ensemble. @@ -286,11 +284,11 @@ class GradientBoostedDecisionTreeEstimator(estimator.Estimator): opposed to contrib) version of tensorflow. output_leaf_index: whether to output leaf indices along with predictions during inference. The leaf node indexes are available in predictions - dict by the key 'leaf_index'. For example, - result_dict = classifier.predict(...) - for example_prediction_result in result_dict: - # access leaf index list by example_prediction_result["leaf_index"] - # which contains one leaf index per tree + dict by the key 'leaf_index'. For example, result_dict = + classifier.predict(...) + for example_prediction_result in result_dict: # access leaf index list + by example_prediction_result["leaf_index"] # which contains one leaf + index per tree override_global_step_value: If after the training is done, global step value must be reset to this value. This should be used to reset global step to a number > number of steps used to train the current ensemble. @@ -353,10 +351,9 @@ class GradientBoostedDecisionTreeRanker(estimator.Estimator): layer. It can also be a function that computes the number of examples based on the depth of the layer that's being built. head: `Head` instance. - ranking_model_pair_keys: Keys to distinguish between features - for left and right part of the training pairs for ranking. For example, - for an Example with features "a.f1" and "b.f1", the keys would be - ("a", "b"). + ranking_model_pair_keys: Keys to distinguish between features for left and + right part of the training pairs for ranking. For example, for an + Example with features "a.f1" and "b.f1", the keys would be ("a", "b"). num_trees: An int, number of trees to build. feature_columns: A list of feature columns. weight_column_name: Name of the column for weights, or None if not @@ -376,12 +373,10 @@ class GradientBoostedDecisionTreeRanker(estimator.Estimator): output_leaf_index: whether to output leaf indices along with predictions during inference. The leaf node indexes are available in predictions dict by the key 'leaf_index'. It is a Tensor of rank 2 and its shape is - [batch_size, num_trees]. - For example, - result_iter = classifier.predict(...) - for result_dict in result_iter: - # access leaf index list by result_dict["leaf_index"] - # which contains one leaf index per tree + [batch_size, num_trees]. For example, result_iter = + classifier.predict(...) + for result_dict in result_iter: # access leaf index list by + result_dict["leaf_index"] # which contains one leaf index per tree override_global_step_value: If after the training is done, global step value must be reset to this value. This should be used to reset global step to a number > number of steps used to train the current ensemble. @@ -417,12 +412,12 @@ class GradientBoostedDecisionTreeRanker(estimator.Estimator): config=config, feature_engineering_fn=feature_engineering_fn) + # When using this estimator, make sure to regularize the hessian (at least l2, # min_node_weight)! # TODO(nponomareva): extend to take multiple quantiles in one go. class GradientBoostedDecisionTreeQuantileRegressor(estimator.Estimator): - """An estimator that does quantile regression and returns quantile estimates. - """ + """An estimator that does quantile regression and returns quantile estimates.""" def __init__(self, learner_config, @@ -449,8 +444,8 @@ class GradientBoostedDecisionTreeQuantileRegressor(estimator.Estimator): layer. It can also be a function that computes the number of examples based on the depth of the layer that's being built. quantiles: a list of quantiles for the loss, each between 0 and 1. - label_dimension: Dimension of regression label. This is the size - of the last dimension of the labels `Tensor` (typically, this has shape + label_dimension: Dimension of regression label. This is the size of the + last dimension of the labels `Tensor` (typically, this has shape `[batch_size, label_dimension]`). When label_dimension>1, it is recommended to use multiclass strategy diagonal hessian or full hessian. num_trees: An int, number of trees to build. @@ -469,11 +464,11 @@ class GradientBoostedDecisionTreeQuantileRegressor(estimator.Estimator): opposed to contrib) version of tensorflow. output_leaf_index: whether to output leaf indices along with predictions during inference. The leaf node indexes are available in predictions - dict by the key 'leaf_index'. For example, - result_dict = classifier.predict(...) - for example_prediction_result in result_dict: - # access leaf index list by example_prediction_result["leaf_index"] - # which contains one leaf index per tree + dict by the key 'leaf_index'. For example, result_dict = + classifier.predict(...) + for example_prediction_result in result_dict: # access leaf index list + by example_prediction_result["leaf_index"] # which contains one leaf + index per tree override_global_step_value: If after the training is done, global step value must be reset to this value. This should be used to reset global step to a number > number of steps used to train the current ensemble. @@ -519,6 +514,7 @@ class GradientBoostedDecisionTreeQuantileRegressor(estimator.Estimator): config=config, feature_engineering_fn=feature_engineering_fn) + # ================== New Estimator interface=================================== # The estimators below use new core Estimator interface and must be used with # new feature columns and heads. @@ -534,10 +530,8 @@ def core_multiclass_head( def loss_fn(labels, logits): result = losses.per_example_maxent_loss( - labels=labels, - logits=logits, - weights=weight_column, - num_classes=n_classes) + # Don't pass the weights: head already multiplies by them. + labels=labels, logits=logits, weights=None, num_classes=n_classes) return result[0] # pylint:disable=protected-access @@ -564,7 +558,8 @@ def core_quantile_regression_head( result = losses.per_example_quantile_regression_loss( labels=labels, predictions=logits, - weights=weight_column, + # Don't pass the weights: head already multiplies by them. + weights=None, quantile=quantiles) return result[0] @@ -623,11 +618,11 @@ class CoreGradientBoostedDecisionTreeEstimator(core_estimator.Estimator): the bias. output_leaf_index: whether to output leaf indices along with predictions during inference. The leaf node indexes are available in predictions - dict by the key 'leaf_index'. For example, - result_dict = classifier.predict(...) - for example_prediction_result in result_dict: - # access leaf index list by example_prediction_result["leaf_index"] - # which contains one leaf index per tree + dict by the key 'leaf_index'. For example, result_dict = + classifier.predict(...) + for example_prediction_result in result_dict: # access leaf index list + by example_prediction_result["leaf_index"] # which contains one leaf + index per tree num_quantiles: Number of quantiles to build for numeric feature values. """ @@ -685,10 +680,9 @@ class CoreGradientBoostedDecisionTreeRanker(core_estimator.Estimator): layer. It can also be a function that computes the number of examples based on the depth of the layer that's being built. head: `Head` instance. - ranking_model_pair_keys: Keys to distinguish between features - for left and right part of the training pairs for ranking. For example, - for an Example with features "a.f1" and "b.f1", the keys would be - ("a", "b"). + ranking_model_pair_keys: Keys to distinguish between features for left and + right part of the training pairs for ranking. For example, for an + Example with features "a.f1" and "b.f1", the keys would be ("a", "b"). num_trees: An int, number of trees to build. feature_columns: A list of feature columns. weight_column_name: Name of the column for weights, or None if not @@ -703,12 +697,10 @@ class CoreGradientBoostedDecisionTreeRanker(core_estimator.Estimator): output_leaf_index: whether to output leaf indices along with predictions during inference. The leaf node indexes are available in predictions dict by the key 'leaf_index'. It is a Tensor of rank 2 and its shape is - [batch_size, num_trees]. - For example, - result_iter = classifier.predict(...) - for result_dict in result_iter: - # access leaf index list by result_dict["leaf_index"] - # which contains one leaf index per tree + [batch_size, num_trees]. For example, result_iter = + classifier.predict(...) + for result_dict in result_iter: # access leaf index list by + result_dict["leaf_index"] # which contains one leaf index per tree num_quantiles: Number of quantiles to build for numeric feature values. Raises: @@ -748,8 +740,7 @@ class CoreGradientBoostedDecisionTreeRanker(core_estimator.Estimator): # TODO(nponomareva): extend to take multiple quantiles in one go. class CoreGradientBoostedDecisionTreeQuantileRegressor( core_estimator.Estimator): - """An estimator that does quantile regression and returns quantile estimates. - """ + """An estimator that does quantile regression and returns quantile estimates.""" def __init__(self, learner_config, @@ -775,8 +766,8 @@ class CoreGradientBoostedDecisionTreeQuantileRegressor( layer. It can also be a function that computes the number of examples based on the depth of the layer that's being built. quantiles: a list of quantiles for the loss, each between 0 and 1. - label_dimension: Dimension of regression label. This is the size - of the last dimension of the labels `Tensor` (typically, this has shape + label_dimension: Dimension of regression label. This is the size of the + last dimension of the labels `Tensor` (typically, this has shape `[batch_size, label_dimension]`). When label_dimension>1, it is recommended to use multiclass strategy diagonal hessian or full hessian. num_trees: An int, number of trees to build. @@ -795,11 +786,11 @@ class CoreGradientBoostedDecisionTreeQuantileRegressor( the bias. output_leaf_index: whether to output leaf indices along with predictions during inference. The leaf node indexes are available in predictions - dict by the key 'leaf_index'. For example, - result_dict = classifier.predict(...) - for example_prediction_result in result_dict: - # access leaf index list by example_prediction_result["leaf_index"] - # which contains one leaf index per tree + dict by the key 'leaf_index'. For example, result_dict = + classifier.predict(...) + for example_prediction_result in result_dict: # access leaf index list + by example_prediction_result["leaf_index"] # which contains one leaf + index per tree num_quantiles: Number of quantiles to build for numeric feature values. """ if len(quantiles) > 1: @@ -814,7 +805,9 @@ class CoreGradientBoostedDecisionTreeQuantileRegressor( params={ 'head': core_quantile_regression_head( - quantiles[0], label_dimension=label_dimension), + quantiles[0], + label_dimension=label_dimension, + weight_column=weight_column_name), 'feature_columns': feature_columns, 'learner_config': diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/model.py b/tensorflow/contrib/boosted_trees/estimator_batch/model.py index a6e422847d3914188bca9e6dff797ba1ffb06749..eecf3c5aeb6c6785cae3fd5808954a73db6190d6 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/model.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/model.py @@ -25,6 +25,7 @@ from tensorflow.contrib.boosted_trees.estimator_batch import estimator_utils from tensorflow.contrib.boosted_trees.estimator_batch import trainer_hooks from tensorflow.contrib.boosted_trees.python.ops import model_ops from tensorflow.contrib.boosted_trees.python.training.functions import gbdt_batch +from tensorflow.core.protobuf import config_pb2 from tensorflow.python.framework import ops from tensorflow.python.ops import state_ops from tensorflow.python.training import training_util @@ -88,6 +89,12 @@ def model_builder(features, if config is None: raise ValueError("Missing estimator RunConfig.") + if config.session_config is not None: + session_config = config.session_config + session_config.allow_soft_placement = True + else: + session_config = config_pb2.ConfigProto(allow_soft_placement=True) + config = config.replace(session_config=session_config) center_bias = params["center_bias"] diff --git a/tensorflow/contrib/boosted_trees/kernels/split_handler_ops.cc b/tensorflow/contrib/boosted_trees/kernels/split_handler_ops.cc index 6d78e27e8f69ea289b686af8402bd91967f997f4..65276242abaf96de8b1936365278b18f8bba93a9 100644 --- a/tensorflow/contrib/boosted_trees/kernels/split_handler_ops.cc +++ b/tensorflow/contrib/boosted_trees/kernels/split_handler_ops.cc @@ -538,7 +538,6 @@ class BuildSparseInequalitySplitsOp : public OpKernel { partition_boundaries[non_empty_partitions[root_idx]]; float best_gain = std::numeric_limits::lowest(); - int32 best_dimension_idx = 0; bool default_right = false; int32 best_element_idx = 0; @@ -571,7 +570,6 @@ class BuildSparseInequalitySplitsOp : public OpKernel { // Iterate through dimensions. for (int j = 0; j < dimension_boundaries.size() - 1; ++j) { const DimensionBoundary& dimension_and_start = dimension_boundaries[j]; - const int32 dimension_id = dimension_and_start.dimension_id; int start_index = dimension_and_start.start_index; // Even for the last dimension, we always have additional dummy @@ -630,7 +628,6 @@ class BuildSparseInequalitySplitsOp : public OpKernel { best_right_node_stats = right_stats_default_left; best_element_idx = element_idx; default_right = false; - best_dimension_idx = dimension_id; } } // Consider calculating the default direction only when there were @@ -648,7 +645,6 @@ class BuildSparseInequalitySplitsOp : public OpKernel { best_right_node_stats = right_stats_default_right; best_element_idx = element_idx; default_right = true; - best_dimension_idx = dimension_id; } } } diff --git a/tensorflow/contrib/boosted_trees/python/ops/model_ops.py b/tensorflow/contrib/boosted_trees/python/ops/model_ops.py index ad6ff0a861af896ef0dd254bd47752d76378d63a..f9945959812f030f76cb481cfcf91cba1f352fc1 100644 --- a/tensorflow/contrib/boosted_trees/python/ops/model_ops.py +++ b/tensorflow/contrib/boosted_trees/python/ops/model_ops.py @@ -96,18 +96,18 @@ class TreeEnsembleVariable(tracking.TrackableResource): self._init_op = None super(TreeEnsembleVariable, self).__init__() - def create_resource(self): + def _create_resource(self): return gen_model_ops.decision_tree_ensemble_resource_handle_op( self._container, shared_name=self._name, name=self._name) - def initialize(self): + def _initialize(self): return gen_model_ops.create_tree_ensemble_variable( self.resource_handle, self._stamp_token, self._tree_ensemble_config) @property def initializer(self): if self._init_op is None: - self._init_op = self.initialize() + self._init_op = self._initialize() return self._init_op def is_initialized(self): diff --git a/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py b/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py index aff7105e94729942efc6e3e9d3ae23b733e8f5ed..82f9b17b3308d6a521c79ee7a6f48f6c3813a769 100644 --- a/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py +++ b/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py @@ -120,8 +120,8 @@ class QuantileAccumulator(tracking.TrackableResource): name = _PATTERN.sub("", name) with ops.name_scope(name, "QuantileAccumulator") as name: self._name = name - self._resource_handle = self.create_resource() - self._init_op = self.initialize() + self._resource_handle = self._create_resource() + self._init_op = self._initialize() is_initialized_op = self.is_initialized() resources.register_resource(self.resource_handle, self._init_op, is_initialized_op) @@ -129,11 +129,11 @@ class QuantileAccumulator(tracking.TrackableResource): self._init_op, name) ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, self._saveable) - def create_resource(self): + def _create_resource(self): return gen_quantile_ops.quantile_stream_resource_handle_op( container=self._container, shared_name=self._name, name=self._name) - def initialize(self): + def _initialize(self): return gen_quantile_ops.create_quantile_accumulator( self.resource_handle, self._init_stamp_token, @@ -145,7 +145,7 @@ class QuantileAccumulator(tracking.TrackableResource): @property def initializer(self): if self._init_op is None: - self._init_op = self.initialize() + self._init_op = self._initialize() return self._init_op def is_initialized(self): diff --git a/tensorflow/contrib/boosted_trees/python/ops/stats_accumulator_ops.py b/tensorflow/contrib/boosted_trees/python/ops/stats_accumulator_ops.py index 2a0a206d97bbf01ac382531df31a66d429842bbb..1f6bbbf5740ec3c47697ea600eef030aa257707f 100644 --- a/tensorflow/contrib/boosted_trees/python/ops/stats_accumulator_ops.py +++ b/tensorflow/contrib/boosted_trees/python/ops/stats_accumulator_ops.py @@ -144,8 +144,8 @@ class StatsAccumulator(tracking.TrackableResource): name = _PATTERN.sub("", name) with ops.name_scope(name, "StatsAccumulator") as name: self._name = name - self._resource_handle = self.create_resource() - self._init_op = self.initialize() + self._resource_handle = self._create_resource() + self._init_op = self._initialize() is_initialized_op = self.is_initialized() resources.register_resource(self.resource_handle, self.initializer, is_initialized_op) @@ -153,7 +153,7 @@ class StatsAccumulator(tracking.TrackableResource): self.resource_handle, self.initializer, self._is_scalar, name) ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, self._saveable) - def create_resource(self): + def _create_resource(self): if self._is_scalar: return ( gen_stats_accumulator_ops.stats_accumulator_scalar_resource_handle_op( @@ -163,7 +163,7 @@ class StatsAccumulator(tracking.TrackableResource): gen_stats_accumulator_ops.stats_accumulator_tensor_resource_handle_op( self._container, self._name, name=self._name)) - def initialize(self): + def _initialize(self): if self._is_scalar: return gen_stats_accumulator_ops.create_stats_accumulator_scalar( self.resource_handle, self._stamp_token) @@ -175,7 +175,7 @@ class StatsAccumulator(tracking.TrackableResource): @property def initializer(self): if self._init_op is None: - self._init_op = self.initialize() + self._init_op = self._initialize() return self._init_op def is_initialized(self): diff --git a/tensorflow/contrib/checkpoint/__init__.py b/tensorflow/contrib/checkpoint/__init__.py index 7b3df962542a656af8052e9f2eae6e83744411f2..a416588691f580143aa4e5ee53ca1e5cab9c42e0 100644 --- a/tensorflow/contrib/checkpoint/__init__.py +++ b/tensorflow/contrib/checkpoint/__init__.py @@ -46,7 +46,6 @@ from __future__ import print_function from tensorflow.contrib.checkpoint.python.containers import UniqueNameTracker from tensorflow.contrib.checkpoint.python.python_state import NumpyState -from tensorflow.contrib.checkpoint.python.python_state import PythonStateWrapper from tensorflow.contrib.checkpoint.python.split_dependency import split_dependency from tensorflow.contrib.checkpoint.python.visualize import dot_graph_from_checkpoint from tensorflow.core.protobuf.trackable_object_graph_pb2 import TrackableObjectGraph as CheckpointableObjectGraph @@ -55,6 +54,7 @@ from tensorflow.python.training.tracking.base import Trackable as Checkpointable from tensorflow.python.training.tracking.data_structures import List from tensorflow.python.training.tracking.data_structures import Mapping from tensorflow.python.training.tracking.data_structures import NoDependency +from tensorflow.python.training.tracking.python_state import PythonState as PythonStateWrapper from tensorflow.python.training.tracking.tracking import AutoTrackable as Checkpointable from tensorflow.python.training.tracking.util import capture_dependencies from tensorflow.python.training.tracking.util import list_objects @@ -62,3 +62,4 @@ from tensorflow.python.training.tracking.util import object_metadata from tensorflow.python.util.all_util import remove_undocumented remove_undocumented(module_name=__name__) + diff --git a/tensorflow/contrib/checkpoint/python/python_state.py b/tensorflow/contrib/checkpoint/python/python_state.py index 737a6c30c1dce65dd7638ee52e6c26a8a40f8321..1ada05227ba566cd3dfbff406e8fed80dccde684 100644 --- a/tensorflow/contrib/checkpoint/python/python_state.py +++ b/tensorflow/contrib/checkpoint/python/python_state.py @@ -17,13 +17,10 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import abc -import functools -import six - import numpy from tensorflow.python.training.tracking import base +from tensorflow.python.training.tracking import python_state as core_python_state # pylint: disable=g-import-not-at-top try: @@ -129,29 +126,7 @@ class NumpyState(base.Trackable): super(NumpyState, self).__setattr__(name, value) -@six.add_metaclass(abc.ABCMeta) -class PythonStateWrapper(base.Trackable): - """Wraps a Python object for storage in an object-based checkpoint.""" - - @abc.abstractmethod - def _serialize(self): - """Callback for `PythonStringStateSaveable` to serialize the object.""" - - @abc.abstractmethod - def _deserialize(self, string_value): - """Callback for `PythonStringStateSaveable` to deserialize the object.""" - - def _gather_saveables_for_checkpoint(self): - """Specify callbacks for saving and restoring `array`.""" - return { - "py_state": functools.partial( - base.PythonStringStateSaveable, - state_callback=self._serialize, - restore_callback=self._deserialize) - } - - -class _NumpyWrapper(PythonStateWrapper): +class _NumpyWrapper(core_python_state.PythonState): """Wraps a NumPy array for storage in an object-based checkpoint.""" def __init__(self, array): @@ -162,7 +137,7 @@ class _NumpyWrapper(PythonStateWrapper): """ self.array = array - def _serialize(self): + def serialize(self): """Callback to serialize the array.""" string_file = BytesIO() try: @@ -172,7 +147,7 @@ class _NumpyWrapper(PythonStateWrapper): string_file.close() return serialized - def _deserialize(self, string_value): + def deserialize(self, string_value): """Callback to deserialize the array.""" string_file = BytesIO(string_value) try: diff --git a/tensorflow/contrib/cmake/external/grpc.cmake b/tensorflow/contrib/cmake/external/grpc.cmake index 793a0f3398a4ba38ae900144154b271da8f98eeb..30b4e2dbdee1117df12ae7ab8ce902e667234fb0 100644 --- a/tensorflow/contrib/cmake/external/grpc.cmake +++ b/tensorflow/contrib/cmake/external/grpc.cmake @@ -17,7 +17,7 @@ include (ExternalProject) set(GRPC_INCLUDE_DIRS ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/include) set(GRPC_URL https://github.com/grpc/grpc.git) set(GRPC_BUILD ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc) -set(GRPC_TAG 69b2d570428623d2d9097ec9e21b3e0a02268ab9) +set(GRPC_TAG 62688b6a05cc85b47fb77dd408611734253e47e2) if(WIN32) # We use unsecure gRPC because boringssl does not build on windows diff --git a/tensorflow/contrib/cmake/python_modules.txt b/tensorflow/contrib/cmake/python_modules.txt index 3d86ab9abbb4cc90c406edc6237c0d2abe440122..fd205a4b9b065a4756fbe3985694bb64b93b85e6 100644 --- a/tensorflow/contrib/cmake/python_modules.txt +++ b/tensorflow/contrib/cmake/python_modules.txt @@ -71,6 +71,7 @@ tensorflow/python/summary/writer tensorflow/python/tools tensorflow/python/tools/api tensorflow/python/tools/api/generator +tensorflow/python/tpu tensorflow/python/training tensorflow/python/training/tracking tensorflow/python/user_ops diff --git a/tensorflow/contrib/compiler/BUILD b/tensorflow/contrib/compiler/BUILD index e32097ceddfec95b8677fc762d641d09078e5343..79c61589112b739837b401010690e7f4ca917d07 100644 --- a/tensorflow/contrib/compiler/BUILD +++ b/tensorflow/contrib/compiler/BUILD @@ -23,6 +23,7 @@ py_library( ], srcs_version = "PY2AND3", deps = [ + ":xla", "//tensorflow/core:protos_all_py", "//tensorflow/python:framework_for_generated_wrappers", ], diff --git a/tensorflow/contrib/compiler/__init__.py b/tensorflow/contrib/compiler/__init__.py index c4937dadfb8be3211377f0ae7017b95e7642dab0..797e5e8164e231e8b3806d40b32774711879b050 100644 --- a/tensorflow/contrib/compiler/__init__.py +++ b/tensorflow/contrib/compiler/__init__.py @@ -19,3 +19,4 @@ from __future__ import division from __future__ import print_function from tensorflow.contrib.compiler import jit +from tensorflow.contrib.compiler import xla diff --git a/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py b/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py index 1cb477716dfc6a9cc793939059784f9d89bcdd8a..c4e37b41c8542fb718b109f64b9edbc1d62c0ebf 100644 --- a/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py +++ b/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py @@ -104,7 +104,7 @@ class _CudnnRNN(base_layer.Layer): # Inference subgraph for unidirectional RNN on, e.g., CPU or mobile. with tf.Graph().as_default(): - single_cell = lambda: tf.contrib.cudnn_rnn.CudnnCompatibleLSTM(num_units) + single_cell = lambda: tf.contrib.cudnn_rnn.CudnnCompatibleLSTMCell(num_units) # NOTE: Even if there's only one layer, the cell needs to be wrapped in # MultiRNNCell. @@ -124,7 +124,7 @@ class _CudnnRNN(base_layer.Layer): # Inference subgraph for bidirectional RNN with tf.Graph().as_default(): - single_cell = lambda: tf.contrib.cudnn_rnn.CudnnCompatibleLSTM(num_units) + single_cell = lambda: tf.contrib.cudnn_rnn.CudnnCompatibleLSTMCell(num_units) cells_fw = [single_cell() for _ in range(num_layers)] cells_bw = [single_cell() for _ in range(num_layers)] diff --git a/tensorflow/contrib/data/python/kernel_tests/assert_element_shape_test.py b/tensorflow/contrib/data/python/kernel_tests/assert_element_shape_test.py index 4db711c1f3f2815e7b8cf275af315c062ce4c02e..077571fcd2091b3b7216c57627a11989f3db1fdf 100644 --- a/tensorflow/contrib/data/python/kernel_tests/assert_element_shape_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/assert_element_shape_test.py @@ -43,10 +43,12 @@ class AssertElementShapeTest(test_base.DatasetTestBase): dataset = dataset_ops.Dataset.range(5).map(create_dataset) expected_shapes = (tensor_shape.TensorShape(2), tensor_shape.TensorShape((3, 4))) - self.assertEqual(expected_shapes, dataset.output_shapes) + self.assertEqual(expected_shapes, + dataset_ops.get_legacy_output_shapes(dataset)) result = dataset.apply(batching.assert_element_shape(expected_shapes)) - self.assertEqual(expected_shapes, result.output_shapes) + self.assertEqual(expected_shapes, + dataset_ops.get_legacy_output_shapes(result)) iterator = dataset_ops.make_initializable_iterator(result) init_op = iterator.initializer @@ -83,12 +85,14 @@ class AssertElementShapeTest(test_base.DatasetTestBase): dataset = dataset_ops.Dataset.range(5).map(create_unknown_shape_dataset) unknown_shapes = (tensor_shape.TensorShape(None), tensor_shape.TensorShape(None)) - self.assertEqual(unknown_shapes, dataset.output_shapes) + self.assertEqual(unknown_shapes, + dataset_ops.get_legacy_output_shapes(dataset)) expected_shapes = (tensor_shape.TensorShape(2), tensor_shape.TensorShape((3, 4))) result = dataset.apply(batching.assert_element_shape(expected_shapes)) - self.assertEqual(expected_shapes, result.output_shapes) + self.assertEqual(expected_shapes, + dataset_ops.get_legacy_output_shapes(result)) iterator = dataset_ops.make_initializable_iterator(result) init_op = iterator.initializer @@ -113,7 +117,8 @@ class AssertElementShapeTest(test_base.DatasetTestBase): dataset = dataset_ops.Dataset.range(3).map(create_unknown_shape_dataset) unknown_shapes = (tensor_shape.TensorShape(None), tensor_shape.TensorShape(None)) - self.assertEqual(unknown_shapes, dataset.output_shapes) + self.assertEqual(unknown_shapes, + dataset_ops.get_legacy_output_shapes(dataset)) wrong_shapes = (tensor_shape.TensorShape(2), tensor_shape.TensorShape((3, 10))) @@ -141,7 +146,8 @@ class AssertElementShapeTest(test_base.DatasetTestBase): # Partial shapes are merged with actual shapes: actual_shapes = (tensor_shape.TensorShape(2), tensor_shape.TensorShape((3, 4))) - self.assertEqual(actual_shapes, result.output_shapes) + self.assertEqual(actual_shapes, + dataset_ops.get_legacy_output_shapes(result)) iterator = dataset_ops.make_initializable_iterator(result) init_op = iterator.initializer @@ -178,12 +184,14 @@ class AssertElementShapeTest(test_base.DatasetTestBase): dataset = dataset_ops.Dataset.range(5).map(create_unknown_shape_dataset) unknown_shapes = (tensor_shape.TensorShape(None), tensor_shape.TensorShape(None)) - self.assertEqual(unknown_shapes, dataset.output_shapes) + self.assertEqual(unknown_shapes, + dataset_ops.get_legacy_output_shapes(dataset)) expected_shapes = (tensor_shape.TensorShape(2), tensor_shape.TensorShape((None, 4))) result = dataset.apply(batching.assert_element_shape(expected_shapes)) - self.assertEqual(expected_shapes, result.output_shapes) + self.assertEqual(expected_shapes, + dataset_ops.get_legacy_output_shapes(result)) iterator = dataset_ops.make_initializable_iterator(result) init_op = iterator.initializer @@ -208,7 +216,8 @@ class AssertElementShapeTest(test_base.DatasetTestBase): dataset = dataset_ops.Dataset.range(3).map(create_unknown_shape_dataset) unknown_shapes = (tensor_shape.TensorShape(None), tensor_shape.TensorShape(None)) - self.assertEqual(unknown_shapes, dataset.output_shapes) + self.assertEqual(unknown_shapes, + dataset_ops.get_legacy_output_shapes(dataset)) wrong_shapes = (tensor_shape.TensorShape(2), tensor_shape.TensorShape((None, 10))) diff --git a/tensorflow/contrib/data/python/ops/batching.py b/tensorflow/contrib/data/python/ops/batching.py index 8c60459ca81cd7a7e08d90339011c54275ea9c0b..f8bb942c0a54d0892f382b1779ff830ab04b8258 100644 --- a/tensorflow/contrib/data/python/ops/batching.py +++ b/tensorflow/contrib/data/python/ops/batching.py @@ -19,6 +19,7 @@ from __future__ import print_function from tensorflow.contrib.framework import with_shape from tensorflow.python.data.experimental.ops import batching +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.util import nest from tensorflow.python.util import deprecation @@ -215,14 +216,14 @@ def assert_element_shape(expected_shapes): return nest.pack_sequence_as(elements, checked_tensors) def _apply_fn(dataset): - output_shapes = _merge_output_shapes(dataset.output_shapes, - expected_shapes) + output_shapes = _merge_output_shapes( + dataset_ops.get_legacy_output_shapes(dataset), expected_shapes) # pylint: disable=protected-access return batching._RestructuredDataset( dataset.map(_check_shape), - dataset.output_types, + dataset_ops.get_legacy_output_types(dataset), output_shapes=output_shapes, - output_classes=dataset.output_classes) + output_classes=dataset_ops.get_legacy_output_classes(dataset)) return _apply_fn diff --git a/tensorflow/contrib/data/python/ops/sliding.py b/tensorflow/contrib/data/python/ops/sliding.py index 6708e01d08135a132b797e317cd2a241c3428f40..2897c72b445b7312a9c0b74b61078baf50880884 100644 --- a/tensorflow/contrib/data/python/ops/sliding.py +++ b/tensorflow/contrib/data/python/ops/sliding.py @@ -18,7 +18,6 @@ from __future__ import division from __future__ import print_function from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.data.util import structure from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops @@ -38,9 +37,7 @@ class _SlideDataset(dataset_ops.UnaryDataset): self._window_shift = ops.convert_to_tensor( window_shift, dtype=dtypes.int64, name="window_shift") - input_structure = structure.convert_legacy_structure( - input_dataset.output_types, input_dataset.output_shapes, - input_dataset.output_classes) + input_structure = dataset_ops.get_structure(input_dataset) self._structure = input_structure._batch(None) # pylint: disable=protected-access variant_tensor = ged_ops.experimental_sliding_window_dataset( self._input_dataset._variant_tensor, # pylint: disable=protected-access diff --git a/tensorflow/contrib/distribute/python/BUILD b/tensorflow/contrib/distribute/python/BUILD index 63879968bfbd06d7005e57724cbc4dff1dbcbb5c..88d6c7a6d274494be7ed8aa699678c8c0519fc6c 100644 --- a/tensorflow/contrib/distribute/python/BUILD +++ b/tensorflow/contrib/distribute/python/BUILD @@ -260,18 +260,8 @@ py_library( srcs = ["tpu_strategy.py"], visibility = ["//tensorflow:internal"], deps = [ - ":one_device_strategy", "//tensorflow/contrib/tpu:tpu_lib", - "//tensorflow/python:constant_op", - "//tensorflow/python:control_flow_ops", - "//tensorflow/python:dtypes", - "//tensorflow/python:framework_ops", - "//tensorflow/python:tensor_util", - "//tensorflow/python:util", - "//tensorflow/python/distribute:input_lib", - "//tensorflow/python/distribute:numpy_dataset", - "//tensorflow/python/distribute:reduce_util", - "//tensorflow/python/distribute:values", + "//tensorflow/python/distribute:tpu_strategy", ], ) @@ -515,6 +505,7 @@ cuda_py_test( name = "cross_device_ops_test", srcs = ["cross_device_ops_test.py"], additional_deps = [ + ":collective_all_reduce_strategy", ":combinations", ":multi_worker_test_base", ":mirrored_strategy", diff --git a/tensorflow/contrib/distribute/python/checkpoint_utils_test.py b/tensorflow/contrib/distribute/python/checkpoint_utils_test.py index 7ee50f03155636a487020d0a9178107a06775588..79369fc6b93b4491c9744653d8d64c5c8a4de30d 100644 --- a/tensorflow/contrib/distribute/python/checkpoint_utils_test.py +++ b/tensorflow/contrib/distribute/python/checkpoint_utils_test.py @@ -56,6 +56,12 @@ def _create_checkpoints(sess, checkpoint_dir): class CheckpointUtilsWithDistributionStrategyTest( test.TestCase, parameterized.TestCase): + def _get_test_object(self): + checkpoint_dir = self.get_temp_dir() + with self.cached_session() as session: + v1, v2 = _create_checkpoints(session, checkpoint_dir) + return checkpoint_dir, v1, v2 + @combinations.generate(combinations.combine( distribution=[combinations.default_strategy, combinations.one_device_strategy, @@ -66,9 +72,7 @@ class CheckpointUtilsWithDistributionStrategyTest( in_replica_mode=[True, False], mode=["graph"])) def testInitFromCheckpoint(self, distribution, in_replica_mode): - checkpoint_dir = self.get_temp_dir() - with self.cached_session() as session: - v1_value, v2_value = _create_checkpoints(session, checkpoint_dir) + checkpoint_dir, v1_value, v2_value = self._get_test_object() def init_and_verify(g): v1 = variable_scope.get_variable("new_var1", [1, 10]) @@ -91,6 +95,39 @@ class CheckpointUtilsWithDistributionStrategyTest( else: init_and_verify(g) + @combinations.generate( + combinations.combine( + distribution=[ + combinations.default_strategy, combinations.one_device_strategy, + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.mirrored_strategy_with_two_gpus, + combinations.core_mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_two_gpus + ], + in_replica_mode=[True, False], + mode=["graph"])) + def testInitFromDifferentNameObject(self, distribution, in_replica_mode): + checkpoint_dir, v1_value, _ = self._get_test_object() + + def init_and_verify(g): + v1 = variable_scope.get_variable("new_var1", [1, 10]) + # Use string add to create new object in each replica + prefix = "new_" + suffix = "var1" + new_var1 = prefix + suffix + checkpoint_utils.init_from_checkpoint(checkpoint_dir, { + "var1": new_var1, + }) + with self.test_session(graph=g) as session: + session.run(variables.global_variables_initializer()) + self.assertAllEqual(v1_value, self.evaluate(v1)) + + with ops.Graph().as_default() as g, distribution.scope(): + if in_replica_mode: + distribution.extended.call_for_each_replica(init_and_verify, [g]) + else: + init_and_verify(g) + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py index 19741627980c34d8c281f7aed6f1464d4a03393e..d4f76e3e7b937798c978f740e080f44a4a1cb418 100644 --- a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py +++ b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function from tensorflow.python.distribute import collective_all_reduce_strategy +from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib from tensorflow.python.distribute import distribute_lib from tensorflow.python.distribute.cluster_resolver import SimpleClusterResolver from tensorflow.python.distribute.cluster_resolver import TFConfigClusterResolver @@ -41,22 +42,34 @@ class CollectiveAllReduceStrategy(distribute_lib.DistributionStrategy): distributed environment. """ - def __init__(self, num_gpus_per_worker=0): + def __init__(self, + num_gpus_per_worker=0, + communication=cross_device_ops_lib.CollectiveCommunication.AUTO): """Initializes the object. Args: num_gpus_per_worker: number of local GPUs or GPUs per worker, the default is 0 meaning CPU only. + communication: optional Enum of type + `distribute.experimental.CollectiveCommunication`. This provides a way + for the user to override the choice of collective op communication. + Possible values include `AUTO`, `RING`, and `NCCL`. """ super(CollectiveAllReduceStrategy, self).__init__( - CollectiveAllReduceExtended(self, num_gpus_per_worker)) + CollectiveAllReduceExtended( + self, + num_gpus_per_worker=num_gpus_per_worker, + communication=communication)) class CollectiveAllReduceExtended( collective_all_reduce_strategy.CollectiveAllReduceExtended): """Implementation of CollectiveAllReduceStrategy.""" - def __init__(self, container_strategy, num_gpus_per_worker): + def __init__(self, + container_strategy, + num_gpus_per_worker, + communication): # Use TFConfigClusterResolver to parse TF_CONFIG. We don't want to change # the constructor's interface to allow customized cluster resolver. Use # SimpleClusterResolver to override num_accelerators. @@ -65,6 +78,9 @@ class CollectiveAllReduceExtended( cluster_spec=tfconfig.cluster_spec(), task_type=tfconfig.task_type, task_id=tfconfig.task_id, - num_accelerators=num_gpus_per_worker) + num_accelerators={"GPU": num_gpus_per_worker}, + rpc_layer=tfconfig.rpc_layer) super(CollectiveAllReduceExtended, self).__init__( - container_strategy, cluster_resolver=cluster_resolver) + container_strategy, + communication=communication, + cluster_resolver=cluster_resolver) diff --git a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py index ee7640dd1cea15e62ae9912ebedbd853778364a6..cef3cc60737f7bf7128fe03b1cc19399699dfdd5 100644 --- a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py +++ b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py @@ -30,6 +30,7 @@ from tensorflow.core.protobuf import rewriter_config_pb2 from tensorflow.python import keras from tensorflow.python.data.ops import dataset_ops from tensorflow.python.distribute import collective_all_reduce_strategy as core_collective_all_reduce_strategy +from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib from tensorflow.python.distribute import cross_device_utils from tensorflow.python.distribute import distribute_lib from tensorflow.python.distribute import multi_worker_util @@ -62,7 +63,9 @@ class MockCollectiveAllReduceStrategy(distribute_lib.DistributionStrategy): def __init__(self, cluster_resolver): super(MockCollectiveAllReduceStrategy, self).__init__( core_collective_all_reduce_strategy.CollectiveAllReduceExtended( - self, cluster_resolver=cluster_resolver)) + self, + communication=cross_device_ops_lib.CollectiveCommunication.AUTO, + cluster_resolver=cluster_resolver)) def create_test_objects(cluster_spec=None, @@ -79,11 +82,11 @@ def create_test_objects(cluster_spec=None, cluster_spec=multi_worker_util.normalize_cluster_spec(cluster_spec), task_type=task_type, task_id=task_id, - num_accelerators=num_gpus) + num_accelerators={'GPU': num_gpus}) target = 'grpc://' + cluster_spec[task_type][task_id] else: cluster_resolver = SimpleClusterResolver( - ClusterSpec({}), num_accelerators=num_gpus) + ClusterSpec({}), num_accelerators={'GPU': num_gpus}) target = '' strategy = MockCollectiveAllReduceStrategy(cluster_resolver) @@ -478,6 +481,24 @@ class DistributedCollectiveAllReduceStrategyTest( self.assertEqual(['CollectiveReduce'], new_rewrite_options.scoped_allocator_opts.enable_op) + @combinations.generate(combinations.combine(mode=['eager'])) + def testEnableCollectiveOps(self): + mock_called = [False] + + # pylint: disable=dangerous-default-value + def mock_enable_collective_ops(server_def, mock_called=mock_called): + self.assertEqual('worker', server_def.job_name) + self.assertEqual(1, server_def.task_index) + self.assertEqual('grpc', server_def.protocol) + mock_called[0] = True + + with test.mock.patch.object(context.context(), 'enable_collective_ops', + mock_enable_collective_ops): + strategy, _, _ = self._get_test_object( + task_type='worker', task_id=1, num_gpus=2, use_core_strategy=True) + self.assertTrue(strategy.extended._std_server_started) + self.assertTrue(mock_called[0]) + class DistributedCollectiveAllReduceStrategyTestWithChief( CollectiveAllReduceStrategyTestBase, parameterized.TestCase): diff --git a/tensorflow/contrib/distribute/python/cross_device_ops_test.py b/tensorflow/contrib/distribute/python/cross_device_ops_test.py index 2d56c25f1034acb29e460985a5eee0133397b97f..2b8e0197961ae37b67dc8958054a03e164242dcd 100644 --- a/tensorflow/contrib/distribute/python/cross_device_ops_test.py +++ b/tensorflow/contrib/distribute/python/cross_device_ops_test.py @@ -23,6 +23,7 @@ import itertools from absl.testing import parameterized import numpy as np +from tensorflow.contrib.distribute.python import collective_all_reduce_strategy from tensorflow.contrib.distribute.python import combinations from tensorflow.contrib.distribute.python import mirrored_strategy from tensorflow.contrib.distribute.python import multi_worker_test_base @@ -429,6 +430,9 @@ class MultiWorkerCrossDeviceOpsTest(multi_worker_test_base.MultiWorkerTestBase, self._testReductionAndBroadcast(cross_device_ops, distribution) +NUM_WORKERS = 3 + + class MultiWorkerCollectiveAllReduceTest( multi_worker_test_base.MultiWorkerTestBase, parameterized.TestCase): @@ -436,9 +440,9 @@ class MultiWorkerCollectiveAllReduceTest( @classmethod def setUpClass(cls): - """Create a local cluster with 2 workers.""" + """Create a local cluster with 3 workers.""" cls._cluster_spec = multi_worker_test_base.create_in_process_cluster( - num_workers=3, num_ps=0) + num_workers=NUM_WORKERS, num_ps=0) def setUp(self): super(MultiWorkerCollectiveAllReduceTest, self).setUp() @@ -446,7 +450,12 @@ class MultiWorkerCollectiveAllReduceTest( # collective key base for different tests. MultiWorkerCollectiveAllReduceTest.collective_key_base += 100000 - def _get_test_objects(self, task_type, task_id, num_gpus=0, local_mode=False): + def _get_test_objects(self, + task_type, + task_id, + num_gpus=0, + use_strategy_object=False, + local_mode=False): collective_keys = cross_device_utils.CollectiveKeys( group_key_start=10 * num_gpus + MultiWorkerCollectiveAllReduceTest.collective_key_base, @@ -455,16 +464,24 @@ class MultiWorkerCollectiveAllReduceTest( instance_key_with_id_start=num_gpus * 10000 + MultiWorkerCollectiveAllReduceTest.collective_key_base) if local_mode: - collective_all_reduce_ops = cross_device_ops_lib.CollectiveAllReduce( - 1, num_gpus, collective_keys=collective_keys) if num_gpus: devices = ["/device:GPU:%d" % i for i in range(num_gpus)] else: devices = ["/device:CPU:0"] - return collective_all_reduce_ops, devices, "" + + if use_strategy_object: + # Still using contrib CollectiveAllReduceStrategy because we can specify + # num_gpus in its constructor. + strategy = collective_all_reduce_strategy.CollectiveAllReduceStrategy( + num_gpus_per_worker=num_gpus) + strategy.extended._collective_keys = collective_keys + strategy.extended._cross_device_ops._collective_keys = collective_keys + return strategy, devices, "" + else: + collective_all_reduce_ops = cross_device_ops_lib.CollectiveAllReduce( + 1, num_gpus, collective_keys=collective_keys) + return collective_all_reduce_ops, devices, "" else: - collective_all_reduce_ops = cross_device_ops_lib.CollectiveAllReduce( - 3, num_gpus, collective_keys=collective_keys) if num_gpus: devices = [ "/job:%s/task:%d/device:GPU:%d" % (task_type, task_id, i) @@ -472,8 +489,23 @@ class MultiWorkerCollectiveAllReduceTest( ] else: devices = ["/job:%s/task:%d" % (task_type, task_id)] - return (collective_all_reduce_ops, devices, - "grpc://" + self._cluster_spec[task_type][task_id]) + + if use_strategy_object: + strategy = collective_all_reduce_strategy.CollectiveAllReduceStrategy( + num_gpus_per_worker=num_gpus) + strategy.configure( + cluster_spec=self._cluster_spec, + task_type=task_type, + task_id=task_id) + strategy.extended._collective_keys = collective_keys + strategy.extended._cross_device_ops._collective_keys = collective_keys + return (strategy, devices, + "grpc://" + self._cluster_spec[task_type][task_id]) + else: + collective_all_reduce_ops = cross_device_ops_lib.CollectiveAllReduce( + NUM_WORKERS, num_gpus, collective_keys=collective_keys) + return (collective_all_reduce_ops, devices, + "grpc://" + self._cluster_spec[task_type][task_id]) def _assert_values_equal(self, left, right, sess): if isinstance(left, list): @@ -493,9 +525,18 @@ class MultiWorkerCollectiveAllReduceTest( for l, r in zip(left_values, right_values): self.assertEqual(l, r) - def _test_reduction(self, task_type, task_id, num_gpus, local_mode=False): + def _test_reduction(self, + task_type, + task_id, + num_gpus, + use_strategy_object=False, + local_mode=False): collective_all_reduce, devices, master_target = self._get_test_objects( - task_type, task_id, num_gpus, local_mode=local_mode) + task_type, + task_id, + num_gpus, + use_strategy_object=use_strategy_object, + local_mode=local_mode) if local_mode: num_workers = 1 worker_device = None @@ -503,6 +544,27 @@ class MultiWorkerCollectiveAllReduceTest( num_workers = len(self._cluster_spec.get("chief", [])) + len( self._cluster_spec.get("worker", [])) worker_device = "/job:%s/task:%d" % (task_type, task_id) + + def _reduce(test_object, reduce_op, per_replica, destinations): + if use_strategy_object: + with test_object.scope(): + # Mimic the behavior that distribution strategy usually strips the + # wrapper if there is only one value. + if len(per_replica.values) == 1: + per_replica = per_replica.values[0] + return test_object.extended.reduce_to(reduce_op, per_replica, + destinations) + else: + return test_object.reduce(reduce_op, per_replica, destinations) + + def _batch_reduce(test_object, reduce_op, value_destination_pairs): + if use_strategy_object: + with test_object.scope(): + return test_object.extended.batch_reduce_to(reduce_op, + value_destination_pairs) + else: + return test_object.batch_reduce(reduce_op, value_destination_pairs) + with ops.Graph().as_default(), \ ops.device(worker_device), \ self.cached_session(target=master_target) as sess: @@ -527,26 +589,30 @@ class MultiWorkerCollectiveAllReduceTest( # test reduce() for destinations in all_destinations: self._assert_values_equal( - collective_all_reduce.reduce( + _reduce( + collective_all_reduce, reduce_util.ReduceOp.MEAN, per_replica, - destinations=destinations), - _fake_mirrored(mean, destinations), sess) + destinations=destinations), _fake_mirrored(mean, destinations), + sess) self._assert_values_equal( - collective_all_reduce.reduce( + _reduce( + collective_all_reduce, reduce_util.ReduceOp.MEAN, per_replica_2, - destinations=destinations), - _fake_mirrored(mean_2, destinations), sess) + destinations=destinations), _fake_mirrored( + mean_2, destinations), sess) self._assert_values_equal( - collective_all_reduce.reduce( + _reduce( + collective_all_reduce, reduce_util.ReduceOp.SUM, per_replica, destinations=destinations), _fake_mirrored(mean * len(devices) * num_workers, destinations), sess) self._assert_values_equal( - collective_all_reduce.reduce( + _reduce( + collective_all_reduce, reduce_util.ReduceOp.SUM, per_replica_2, destinations=destinations), @@ -556,17 +622,13 @@ class MultiWorkerCollectiveAllReduceTest( # test batch_reduce() for d1, d2 in itertools.product(all_destinations, all_destinations): self._assert_values_equal( - collective_all_reduce.batch_reduce(reduce_util.ReduceOp.MEAN, - [(per_replica, d1), - (per_replica_2, d2)]), - [ - _fake_mirrored(mean, d1), - _fake_mirrored(mean_2, d2) - ], sess) + _batch_reduce(collective_all_reduce, reduce_util.ReduceOp.MEAN, + [(per_replica, d1), (per_replica_2, d2)]), + [_fake_mirrored(mean, d1), + _fake_mirrored(mean_2, d2)], sess) self._assert_values_equal( - collective_all_reduce.batch_reduce(reduce_util.ReduceOp.SUM, - [(per_replica, d1), - (per_replica_2, d2)]), + _batch_reduce(collective_all_reduce, reduce_util.ReduceOp.SUM, + [(per_replica, d1), (per_replica_2, d2)]), [ _fake_mirrored(mean * len(devices) * num_workers, d1), _fake_mirrored(mean_2 * len(devices) * num_workers, d2) @@ -575,18 +637,36 @@ class MultiWorkerCollectiveAllReduceTest( return True @combinations.generate( - combinations.combine(mode=["graph"], num_gpus=[0, 1, 2], required_gpus=1)) - def testReductionDistributed(self, num_gpus): + combinations.combine( + mode=["graph"], + num_gpus=[0, 1, 2], + required_gpus=1, + use_strategy_object=[True, False])) + def testReductionDistributed(self, num_gpus, use_strategy_object): if context.num_gpus() < num_gpus: return - self._run_between_graph_clients(self._test_reduction, self._cluster_spec, - num_gpus) + self._run_between_graph_clients( + self._test_reduction, + self._cluster_spec, + num_gpus, + use_strategy_object=use_strategy_object) # Collective ops doesn't support strategy with one device. - def testReductionLocal(self, num_gpus=2): + @combinations.generate( + combinations.combine( + mode=["graph"], + num_gpus=[2], + required_gpus=2, + use_strategy_object=[True, False])) + def testReductionLocal(self, num_gpus, use_strategy_object): if context.num_gpus() < num_gpus: return - self._test_reduction(None, None, num_gpus, local_mode=True) + self._test_reduction( + None, + None, + num_gpus, + use_strategy_object=use_strategy_object, + local_mode=True) if __name__ == "__main__": diff --git a/tensorflow/contrib/distribute/python/estimator_integration_test.py b/tensorflow/contrib/distribute/python/estimator_integration_test.py index e17085628ba6d1dfc79839fd824801723f07a518..1ff1e7c1d255492e0535175dae7594d2ceb4010b 100644 --- a/tensorflow/contrib/distribute/python/estimator_integration_test.py +++ b/tensorflow/contrib/distribute/python/estimator_integration_test.py @@ -22,7 +22,6 @@ import shutil import tempfile from absl.testing import parameterized import numpy as np -import six from tensorflow.contrib.distribute.python import combinations from tensorflow.contrib.optimizer_v2 import adagrad @@ -117,7 +116,7 @@ class DNNLinearCombinedClassifierIntegrationTest(test.TestCase, scores = estimator.evaluate(eval_input_fn) self.assertEqual(num_steps, scores[ops.GraphKeys.GLOBAL_STEP]) - self.assertIn('loss', six.iterkeys(scores)) + self.assertIn('loss', scores) predictions = np.array([ x[prediction_keys.PredictionKeys.PREDICTIONS] diff --git a/tensorflow/contrib/distribute/python/estimator_training_test.py b/tensorflow/contrib/distribute/python/estimator_training_test.py index 3f55a8a1c8b88d1b8e4031547fa3fbe519983630..1b422ef2d19127a2a699ea01d224cab9fde158e0 100644 --- a/tensorflow/contrib/distribute/python/estimator_training_test.py +++ b/tensorflow/contrib/distribute/python/estimator_training_test.py @@ -34,6 +34,7 @@ from tensorflow.contrib.distribute.python import multi_worker_test_base from tensorflow.contrib.distribute.python import parameter_server_strategy from tensorflow.contrib.optimizer_v2 import adagrad from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib from tensorflow.python.distribute import distribute_coordinator as dc from tensorflow.python.distribute import estimator_training as dc_training from tensorflow.python.distribute.distribute_config import DistributeConfig @@ -250,7 +251,7 @@ class DistributeCoordinatorIntegrationTest( def _get_strategy_object(self, strategy_cls): if strategy_cls == mirrored_strategy.CoreMirroredStrategy: - return strategy_cls(mirrored_strategy.all_local_devices()) + return strategy_cls() else: return strategy_cls(num_gpus_per_worker=context.num_gpus()) @@ -268,6 +269,7 @@ class DistributeCoordinatorIntegrationTest( mirrored_strategy.MirroredStrategy, mirrored_strategy.CoreMirroredStrategy, parameter_server_strategy.ParameterServerStrategy, + collective_all_reduce_strategy.CollectiveAllReduceStrategy, ], required_gpus=[0, 1])) def test_complete_flow_standalone_client(self, train_distribute_cls, @@ -287,6 +289,34 @@ class DistributeCoordinatorIntegrationTest( cluster_spec) self._inspect_train_and_eval_events(estimator) + @combinations.generate( + combinations.combine( + mode=["graph"], + eval_distribute_class=[ + None, + mirrored_strategy.MirroredStrategy, + mirrored_strategy.CoreMirroredStrategy, + parameter_server_strategy.ParameterServerStrategy, + ], + required_gpus=[0, 1])) + def test_complete_flow_standalone_client_collective_nccl( + self, eval_distribute_class): + train_distribute = ( + collective_all_reduce_strategy.CollectiveAllReduceStrategy( + num_gpus_per_worker=context.num_gpus(), + communication=cross_device_ops_lib.CollectiveCommunication.NCCL)) + + if eval_distribute_class: + eval_distribute = self._get_strategy_object(eval_distribute_class) + else: + eval_distribute = None + + cluster_spec = copy.deepcopy(self._cluster_spec) + cluster_spec.pop("ps", None) + estimator = self._complete_flow(train_distribute, eval_distribute, + cluster_spec) + self._inspect_train_and_eval_events(estimator) + @combinations.generate( combinations.combine( mode=["graph"], @@ -342,12 +372,14 @@ class DistributeCoordinatorIntegrationTest( parameter_server_strategy.ParameterServerStrategy, ], eval_distribute_cls=[ - None, mirrored_strategy.MirroredStrategy, + None, + mirrored_strategy.MirroredStrategy, mirrored_strategy.CoreMirroredStrategy, parameter_server_strategy.ParameterServerStrategy, + collective_all_reduce_strategy.CollectiveAllReduceStrategy, ], required_gpus=[0, 1])) - def test_complete_flow_indepedent_worker_between_graph( + def test_complete_flow_independent_worker_between_graph( self, train_distribute_cls, eval_distribute_cls): if (context.num_gpus() < 2 and eval_distribute_cls == collective_all_reduce_strategy.CollectiveAllReduceStrategy): @@ -399,8 +431,8 @@ class DistributeCoordinatorIntegrationTest( mirrored_strategy.CoreMirroredStrategy ], required_gpus=[0, 1])) - def test_complete_flow_indepedent_worker_in_graph(self, train_distribute_cls, - eval_distribute_cls): + def test_complete_flow_independent_worker_in_graph(self, train_distribute_cls, + eval_distribute_cls): train_distribute = self._get_strategy_object(train_distribute_cls) if eval_distribute_cls: diff --git a/tensorflow/contrib/distribute/python/examples/BUILD b/tensorflow/contrib/distribute/python/examples/BUILD index 5f89df5824a8d03198987a6fa3d21e2330deedf0..58bede801ff13bb60ed4ada4810eb8ce2dbcb0a3 100644 --- a/tensorflow/contrib/distribute/python/examples/BUILD +++ b/tensorflow/contrib/distribute/python/examples/BUILD @@ -56,3 +56,14 @@ py_binary( "//third_party/py/numpy", ], ) + +py_binary( + name = "mnist_tf1_tpu", + srcs = [ + "mnist_tf1_tpu.py", + ], + deps = [ + "//tensorflow:tensorflow_py", + "//third_party/py/numpy", + ], +) diff --git a/tensorflow/contrib/distribute/python/examples/mnist_tf1_tpu.py b/tensorflow/contrib/distribute/python/examples/mnist_tf1_tpu.py new file mode 100644 index 0000000000000000000000000000000000000000..ae8194c576e67f7ba864f63885c9b028e4136e61 --- /dev/null +++ b/tensorflow/contrib/distribute/python/examples/mnist_tf1_tpu.py @@ -0,0 +1,188 @@ +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Run MNIST on multiple GPUs on using MirroredStrategy with eager execution. + +By default, runs on all available GPUs, or CPU if no GPUs are available. + +NOTE: Currently, this takes more time than when running MNIST in eager without +MirroredStrategy because of a number overheads. Therefore, this is just a +proof of concept right now and cannot be used to actually scale up training. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from absl import app +from absl import flags +import numpy as np +import tensorflow as tf + + +flags.DEFINE_string("tpu", None, "Name of the TPU to use.") +flags.DEFINE_integer("batch_size", 64, + "What should be the size of each batch?") +flags.DEFINE_integer("num_epochs", 10, "How many epochs to run?") +flags.DEFINE_float("learning_rate", 0.01, "Learning Rate") +flags.DEFINE_float("momentum", 0.5, "SGD momentum") + +FLAGS = flags.FLAGS +NUM_TRAIN_IMAGES = 60000 + + +def create_model(): + max_pool = tf.keras.layers.MaxPooling2D((2, 2), (2, 2), padding="same") + # The model consists of a sequential chain of layers, so tf.keras.Sequential + # (a subclass of tf.keras.Model) makes for a compact description. + return tf.keras.Sequential([ + tf.keras.layers.Reshape( + target_shape=[28, 28, 1], + input_shape=(28, 28,)), + tf.keras.layers.Conv2D(2, 5, padding="same", activation=tf.nn.relu), + max_pool, + tf.keras.layers.Conv2D(4, 5, padding="same", activation=tf.nn.relu), + max_pool, + tf.keras.layers.Flatten(), + tf.keras.layers.Dense(32, activation=tf.nn.relu), + tf.keras.layers.Dropout(0.4), + tf.keras.layers.Dense(10)]) + + +def compute_loss(logits, labels): + loss = tf.reduce_sum( + tf.nn.sparse_softmax_cross_entropy_with_logits( + logits=logits, labels=labels)) + # Scale loss by global batch size. + return loss * (1. / FLAGS.batch_size) + + +def mnist_datasets(): + (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data() + # Numpy defaults to dtype=float64; TF defaults to float32. Stick with float32. + x_train, x_test = x_train / np.float32(255), x_test / np.float32(255) + y_train, y_test = y_train.astype(np.int64), y_test.astype(np.int64) + # TODO(priyag): `strategy.make_numpy_iterator` can be used directly instead of + # converting to datasets. + train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)) + test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)) + return train_dataset, test_dataset + + +def main(argv): + """Run a CNN model on MNIST data to demonstrate DistributedStrategies.""" + del argv # Unused. + tf.disable_v2_behavior() + + cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver( + tpu=FLAGS.tpu) + strategy = tf.contrib.distribute.TPUStrategy(cluster_resolver) + + with strategy.scope(): + train_ds, test_ds = mnist_datasets() + train_ds = train_ds.shuffle(NUM_TRAIN_IMAGES).batch(FLAGS.batch_size) + test_ds = test_ds.batch(FLAGS.batch_size) + + model = create_model() + optimizer = tf.keras.optimizers.SGD(FLAGS.learning_rate, FLAGS.momentum) + training_loss = tf.keras.metrics.Mean("training_loss", dtype=tf.float32) + training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy( + "training_accuracy", dtype=tf.float32) + test_loss = tf.keras.metrics.Mean("test_loss", dtype=tf.float32) + test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy( + "test_accuracy", dtype=tf.float32) + + def train_step(inputs): # pylint: disable=missing-docstring + images, labels = inputs + with tf.GradientTape() as tape: + logits = model(images, training=True) + loss = compute_loss(logits, labels) + grads = tape.gradient(loss, model.variables) + update_vars = optimizer.apply_gradients(zip(grads, model.variables)) + update_loss = training_loss.update_state(loss) + update_accuracy = training_accuracy.update_state(labels, logits) + + with tf.control_dependencies([update_vars, update_loss, update_accuracy]): + return tf.identity(loss) + + def test_step(inputs): + images, labels = inputs + logits = model(images, training=False) + loss = compute_loss(logits, labels) + update_loss = test_loss.update_state(loss) + update_accuracy = test_accuracy.update_state(labels, logits) + + with tf.control_dependencies([update_loss, update_accuracy]): + return tf.identity(loss) + + train_iterator = strategy.make_dataset_iterator(train_ds) + test_iterator = strategy.make_dataset_iterator(test_ds) + + dist_train = strategy.unwrap( + strategy.experimental_run(train_step, train_iterator)) + dist_test = strategy.unwrap( + strategy.experimental_run(test_step, test_iterator)) + + training_loss_result = training_loss.result() + training_accuracy_result = training_accuracy.result() + test_loss_result = test_loss.result() + test_accuracy_result = test_accuracy.result() + + tf.contrib.distribute.initialize_tpu_system(cluster_resolver) + + train_iterator_init = train_iterator.initialize() + test_iterator_init = test_iterator.initialize() + + all_variables = ( + tf.global_variables() + + training_loss.variables + + training_accuracy.variables + + test_loss.variables + + test_accuracy.variables) + + with tf.Session(cluster_resolver.master()) as session: + session.run([v.initializer for v in all_variables]) + + for epoch in range(0, FLAGS.num_epochs): + # Train + print("Starting epoch {}".format(epoch)) + session.run(train_iterator_init) + while True: + try: + session.run(dist_train) + except tf.errors.OutOfRangeError: + break + print("Training loss: {:0.4f}, accuracy: {:0.2f}%".format( + session.run(training_loss_result), + session.run(training_accuracy_result) * 100)) + training_loss.reset_states() + training_accuracy.reset_states() + + # Test + session.run(test_iterator_init) + while True: + try: + session.run(dist_test) + except tf.errors.OutOfRangeError: + break + print("Test loss: {:0.4f}, accuracy: {:0.2f}%".format( + session.run(test_loss_result), + session.run(test_accuracy_result) * 100)) + test_loss.reset_states() + test_accuracy.reset_states() + + +if __name__ == "__main__": + flags.mark_flag_as_required("tpu") + app.run(main) diff --git a/tensorflow/contrib/distribute/python/keras_test.py b/tensorflow/contrib/distribute/python/keras_test.py index 77e241974f7c4c27382ab548a202891fdbbc6ba0..e733cc4e73d0e4b98ae581a7c85a1cf43e345b97 100644 --- a/tensorflow/contrib/distribute/python/keras_test.py +++ b/tensorflow/contrib/distribute/python/keras_test.py @@ -288,8 +288,8 @@ def all_strategy_combinations(): return strategy_minus_tpu_combinations() + tpu_strategy_combinations() -def all_strategy_combinations_minus_default(): - strategy_minus_default_combinations = combinations.combine( +def all_strategy_minus_default_and_tpu_combinations(): + return combinations.combine( distribution=[ combinations.one_device_strategy, combinations.one_device_strategy_gpu, @@ -298,7 +298,11 @@ def all_strategy_combinations_minus_default(): combinations.core_mirrored_strategy_with_gpu_and_cpu, combinations.core_mirrored_strategy_with_two_gpus], mode=['graph', 'eager']) - return strategy_minus_default_combinations + tpu_strategy_combinations() + + +def all_strategy_combinations_minus_default(): + return (all_strategy_minus_default_and_tpu_combinations() + + tpu_strategy_combinations()) def strategy_and_optimizer_combinations(): @@ -723,14 +727,22 @@ class TestDistributionStrategyWithNumpyArrays(test.TestCase, inputs = np.zeros((10, 3), dtype=np.float32) # As sample size is 10, we batch by 4 so that the last batch is - # a partial batch. Also `fit()` using numpy array as inputs without + # a partial batch. Also `predict()` using numpy array as inputs without # distribution strategy uses entire sample as a single batch. As so, # we remove parameters `batch_size` and `steps`. + predict_ground_truth = cpu_model.predict(inputs) cpu_model.set_weights(model_with_ds_strategy.get_weights()) self.assertAllClose( model_with_ds_strategy.predict(inputs, batch_size=4, steps=3), - cpu_model.predict(inputs), - atol=1e-5, rtol=1e-5) + predict_ground_truth, + atol=1e-5, + rtol=1e-5) + # Test that `steps` is inferred correctly when final partial batch exists. + self.assertAllClose( + model_with_ds_strategy.predict(inputs, batch_size=4), + predict_ground_truth, + atol=1e-5, + rtol=1e-5) @combinations.generate(tpu_strategy_combinations()) def test_predict_multi_output_model_with_partial_batch( @@ -1282,28 +1294,6 @@ class TestDistributionStrategyWithKerasModels(test.TestCase, model.predict(inputs, steps=1) model.evaluate(inputs, targets, steps=1) - # TODO(b/124377929): Remove error assertions once subclassed models - # are supported in DistributedStrategy. - @combinations.generate(all_strategy_combinations_minus_default()) - def test_distribution_strategy_on_subclassed_model(self, distribution): - with distribution.scope(): - model = simple_subclassed_model() - optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001) - loss = 'mse' - model.compile(optimizer, loss) - - inputs = np.zeros((64, 3), dtype=np.float32) - targets = np.zeros((64, 2), dtype=np.float32) - - with self.assertRaisesRegexp(AttributeError, 'has no attribute'): - model.fit(inputs, targets, epochs=1, steps_per_epoch=2) - - with self.assertRaisesRegexp(AttributeError, 'has no attribute'): - model.predict(inputs, steps=1) - - with self.assertRaisesRegexp(AttributeError, 'has no attribute'): - model.evaluate(inputs, targets, steps=1) - @combinations.generate(all_strategy_combinations_minus_default()) def test_distribution_strategy_one_dimensional(self, distribution): with distribution.scope(): diff --git a/tensorflow/contrib/distribute/python/keras_utils_test.py b/tensorflow/contrib/distribute/python/keras_utils_test.py index 36eaee77f21a9f6d62a7c3f616d0126b7a4a8902..fef40671369ec1c0897d37e0a78b42f3288d8399 100644 --- a/tensorflow/contrib/distribute/python/keras_utils_test.py +++ b/tensorflow/contrib/distribute/python/keras_utils_test.py @@ -34,6 +34,7 @@ from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.keras.engine import distributed_training_utils from tensorflow.python.keras.optimizer_v2 import gradient_descent as gradient_descent_keras +from tensorflow.python.ops import math_ops from tensorflow.python.training import gradient_descent @@ -318,6 +319,99 @@ class TestDistributionStrategyErrorCases(test.TestCase, parameterized.TestCase): verbose=0, callbacks=[keras.callbacks.ReduceLROnPlateau()]) + @combinations.generate( + combinations.combine( + distribution=[combinations.one_device_strategy], mode=['graph'])) + def test_distribution_strategy_with_add_metric_add_loss(self, distribution): + with distribution.scope(): + x = keras.layers.Input(shape=(1,)) + y = keras.layers.Dense(1, kernel_initializer='ones')(x) + + err_msg = ( + 'We currently do not support compiling the model with distribution ' + r'strategy if `model.add_loss\(tensor\)` or ' + r'`model.add_metric\(tensor\)` has been called.') + + # Test with add_metric. + model = keras.models.Model(x, y) + model.add_metric( + math_ops.reduce_sum(y), name='metric_1', aggregation='mean') + with self.assertRaisesRegex(ValueError, err_msg): + model.compile('sgd',) + + # Test with add_loss. + model = keras.models.Model(x, y) + model.add_loss(math_ops.reduce_mean(y)) + with self.assertRaisesRegex(ValueError, err_msg): + model.compile('sgd',) + + @combinations.generate( + combinations.combine( + distribution=[combinations.one_device_strategy], mode=['eager'])) + def test_distribution_strategy_with_run_eagerly(self, distribution): + with distribution.scope(): + x = keras.layers.Input(shape=(1,)) + y = keras.layers.Dense(1, kernel_initializer='ones')(x) + model = keras.models.Model(x, y) + + err_msg = ('We currently do not support enabling `run_eagerly` with ' + 'distribution strategy.') + with self.assertRaisesRegex(ValueError, err_msg): + model.compile('sgd', run_eagerly=True) + + # TODO(b/124377929): Remove error assertions once subclassed models + # are supported in DistributedStrategy. + @combinations.generate( + combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu + ], + mode=['graph', 'eager'])) + def test_distribution_strategy_on_subclassed_model(self, distribution): + with distribution.scope(): + + class _SimpleMLP(keras.Model): + + def __init__(self, num_labels): + super(_SimpleMLP, self).__init__() + self.dense = keras.layers.Dense(num_labels) + + def call(self, inputs): + return self.dense(inputs) + + model = _SimpleMLP(3) + + with self.assertRaisesRegexp( + ValueError, + 'We currently do not support distribution strategy with a ' + '`Sequential` model that is created without ' + '`input_shape`/`input_dim` set in its first layer or ' + 'a subclassed model.'): + model.compile('sgd') + + @combinations.generate( + combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu + ], + mode=['graph', 'eager'])) + def test_distribution_strategy_on_deferred_sequential_model( + self, distribution): + with distribution.scope(): + model = keras.models.Sequential() + model.add(keras.layers.Dense(16, activation='relu')) + model.add(keras.layers.Dense(3, activation='softmax')) + + with self.assertRaisesRegexp( + ValueError, + 'We currently do not support distribution strategy with a ' + '`Sequential` model that is created without ' + '`input_shape`/`input_dim` set in its first layer or ' + 'a subclassed model.'): + model.compile('sgd') + class TestDistributionStrategyWithLossMasking(test.TestCase, parameterized.TestCase): @@ -414,8 +508,7 @@ class TestDistributionStrategySaveLoadWeights(test.TestCase, @combinations.generate( keras_test_lib.all_strategy_combinations_minus_default()) - def test_save_load_trackable(self, distribution): - # TODO(sourabhbajaj): Test fails with optimizer v2 without h5 + def test_save_load_trackable_optimizer_v1(self, distribution): with self.cached_session(): dataset = keras_test_lib.get_dataset(distribution) with distribution.scope(): @@ -433,6 +526,27 @@ class TestDistributionStrategySaveLoadWeights(test.TestCase, keras_test_lib.get_predict_dataset(distribution), steps=2) model_2.fit(dataset, epochs=1, steps_per_epoch=1) + @combinations.generate( + keras_test_lib.all_strategy_minus_default_and_tpu_combinations()) + def test_save_load_trackable_optimizer_v2(self, distribution): + # TODO(b/123533246): Enable the test for TPU once bug is fixed + with self.cached_session(): + dataset = keras_test_lib.get_dataset(distribution) + with distribution.scope(): + model = keras_test_lib.get_model() + model.compile(gradient_descent_keras.SGD(0.01), 'mse') + model.fit(dataset, epochs=1, steps_per_epoch=1) + + weights_file = tempfile.mktemp() + model.save_weights(weights_file) + + model_2 = keras_test_lib.get_model() + model_2.compile(gradient_descent_keras.SGD(0.01), 'mse') + model_2.load_weights(weights_file) + model_2.predict( + keras_test_lib.get_predict_dataset(distribution), steps=2) + model_2.fit(dataset, epochs=1, steps_per_epoch=1) + class TestDistributionStrategyValidation(test.TestCase, parameterized.TestCase): diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py index 5ce731816ccefe36c1f876c79589e448f00b86f5..e533bdf27715d7657596b086800fb00e6d326360 100644 --- a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py +++ b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import json import sys from absl.testing import parameterized @@ -29,12 +30,14 @@ from tensorflow.contrib.distribute.python import multi_worker_test_base from tensorflow.contrib.distribute.python import strategy_test_lib from tensorflow.core.protobuf import config_pb2 from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib from tensorflow.python.distribute import device_util from tensorflow.python.distribute import distribution_strategy_context as ds_context from tensorflow.python.distribute import reduce_util from tensorflow.python.distribute import values from tensorflow.python.eager import backprop from tensorflow.python.eager import context +from tensorflow.python.eager import def_function from tensorflow.python.eager import function from tensorflow.python.eager import test from tensorflow.python.framework import constant_op @@ -269,6 +272,28 @@ class MirroredStrategyCallForEachReplicaTest(test.TestCase): self.assertEqual(in_scope, unwrapped[0]) self.assertEqual(in_scope, originally) + def testFunctionInCallForEachReplicaNoMergeCall(self, distribution): + @def_function.function + def model_fn(): + return 0. + + with distribution.scope(): + result = distribution.extended.call_for_each_replica(model_fn) + self.assertEqual((0., 0.), self.evaluate(result.values)) + + def testFunctionInCallForEachReplicaWithMergeCall(self, distribution): + def merge_fn(_): + pass + + @def_function.function + def model_fn(): + ds_context.get_replica_context().merge_call(merge_fn) + return 0. + + with distribution.scope(): + with self.assertRaisesRegexp( + RuntimeError, "`merge_call` called while defining a new graph."): + distribution.extended.call_for_each_replica(model_fn) @combinations.generate(combinations.combine( distribution=[ @@ -430,7 +455,7 @@ class MirroredStrategyVariableCreationTest(test.TestCase): self.assertEqual("var0:0", v0.name) self.assertIsInstance(v1, values.MirroredVariable) self.assertEqual("common/var1:0", v1.name) - self.assertIsInstance(v2, values.ReplicaLocalVariable) + self.assertIsInstance(v2, values.SyncOnReadVariable) self.assertEqual("common/var2:0", v2.name) self.assertEqual(variable_scope.VariableAggregation.SUM, v2.aggregation) self.assertIsInstance(v3, values.MirroredVariable) @@ -467,7 +492,7 @@ class MirroredStrategyVariableCreationTest(test.TestCase): self.assertEqual("main/var0:0", v0.name) self.assertIsInstance(v1, values.MirroredVariable) self.assertEqual("main/common/var1:0", v1.name) - self.assertIsInstance(v2, values.ReplicaLocalVariable) + self.assertIsInstance(v2, values.SyncOnReadVariable) self.assertEqual("main/common/var2:0", v2.name) self.assertEqual(variable_scope.VariableAggregation.SUM, v2.aggregation) @@ -619,7 +644,7 @@ class MirroredStrategyVariableCreationTest(test.TestCase): with self.assertRaises(RuntimeError): _ = distribution.extended.call_for_each_replica(model_fn, args=(names,)) - def testReplicaLocalVariable(self, distribution): + def testSyncOnReadVariable(self, distribution): all_v_sum = {} all_v_mean = {} components_sum = {} @@ -635,8 +660,8 @@ class MirroredStrategyVariableCreationTest(test.TestCase): 4.0, synchronization=variable_scope.VariableSynchronization.ON_READ, aggregation=variable_scope.VariableAggregation.MEAN) - self.assertTrue(isinstance(v_sum, values.ReplicaLocalVariable)) - self.assertTrue(isinstance(v_mean, values.ReplicaLocalVariable)) + self.assertIsInstance(v_sum, values.SyncOnReadVariable) + self.assertIsInstance(v_mean, values.SyncOnReadVariable) updates = [v_sum.assign_add(2.0 + replica_id), v_mean.assign(6.0 * replica_id)] all_v_sum[replica_id] = v_sum @@ -650,7 +675,7 @@ class MirroredStrategyVariableCreationTest(test.TestCase): return updates, v_sum, v_mean, c_sum, c_mean with distribution.scope(): - # Create "sum" and "mean" versions of ReplicaLocalVariables. + # Create "sum" and "mean" versions of SyncOnReadVariables. ret_ops, ret_v_sum, ret_v_mean, regrouped_sum, regrouped_mean = ( distribution.extended.call_for_each_replica(model_fn)) # Should see the same wrapping instance in all replicas. @@ -716,13 +741,13 @@ class MirroredStrategyVariableCreationTest(test.TestCase): _, v1 = distribution.unwrap(v) self.assertStartsWith(v1._op.name, "replica_1/") - def testReplicaLocalVariableUpdate(self, distribution): + def testSyncOnReadVariableUpdate(self, distribution): def model_fn(): v_sum = variable_scope.variable( 1.0, synchronization=variable_scope.VariableSynchronization.ON_READ, aggregation=variable_scope.VariableAggregation.SUM) - self.assertTrue(isinstance(v_sum, values.ReplicaLocalVariable)) + self.assertIsInstance(v_sum, values.SyncOnReadVariable) return v_sum def update(var, value): @@ -733,7 +758,7 @@ class MirroredStrategyVariableCreationTest(test.TestCase): # Initialize variables. self.evaluate(variables.global_variables_initializer()) - # Assert that the aggregated value of the replica local vars is the sum + # Assert that the aggregated value of the sync on read var is the sum # of the individual values before running the update ops. self.assertEqual(1.0, self.evaluate(ret_v_sum.get( distribution.extended.worker_devices[0]).read_value())) @@ -743,7 +768,7 @@ class MirroredStrategyVariableCreationTest(test.TestCase): update_ops = distribution.extended.update( ret_v_sum, update, args=(5.0,), group=False) self.evaluate(update_ops) - # Assert that the aggregated value of the replica local vars is the sum + # Assert that the aggregated value of the sync on read vars is the sum # of the individual values after running the update ops. self.assertEqual(5.0, self.evaluate(ret_v_sum.get( distribution.extended.worker_devices[0]).read_value())) @@ -752,11 +777,11 @@ class MirroredStrategyVariableCreationTest(test.TestCase): def testVarDistributeStrategy(self, distribution): with distribution.scope(): mirrored = variable_scope.variable(1.0) - replica_local = variable_scope.variable( + sync_on_read = variable_scope.variable( 1.0, synchronization=variable_scope.VariableSynchronization.ON_READ) self.assertIs(distribution, mirrored.distribute_strategy) - self.assertIs(distribution, replica_local.distribute_strategy) + self.assertIs(distribution, sync_on_read.distribute_strategy) @combinations.generate(combinations.combine( @@ -1123,7 +1148,7 @@ class MirroredVariableUpdateTest(test.TestCase): combinations.mirrored_strategy_with_gpu_and_cpu, combinations.core_mirrored_strategy_with_gpu_and_cpu], mode=["graph", "eager"])) -class MirroredAndReplicaLocalVariableInitializerTest(test.TestCase): +class MirroredAndSyncOnReadVariableInitializerTest(test.TestCase): def testAssignMirroredVarInitializer(self, distribution): # This test is not eager compatible since in eager variables are initialized @@ -1149,17 +1174,16 @@ class MirroredAndReplicaLocalVariableInitializerTest(test.TestCase): 1.0, synchronization=variable_scope.VariableSynchronization.ON_READ, aggregation=variable_scope.VariableAggregation.SUM) - self.assertTrue(isinstance(v_sum, values.ReplicaLocalVariable)) + self.assertIsInstance(v_sum, values.SyncOnReadVariable) return v_sum with distribution.scope(): - replica_local_var = distribution.extended.call_for_each_replica( + sync_on_read_var = distribution.extended.call_for_each_replica( model_fn) - self.assertTrue(isinstance(replica_local_var, - values.ReplicaLocalVariable)) - self.assertFalse(self.evaluate(replica_local_var.is_initialized())) - self.evaluate(replica_local_var.initializer) - self.assertTrue(self.evaluate(replica_local_var.is_initialized())) + self.assertIsInstance(sync_on_read_var, values.SyncOnReadVariable) + self.assertFalse(self.evaluate(sync_on_read_var.is_initialized())) + self.evaluate(sync_on_read_var.initializer) + self.assertTrue(self.evaluate(sync_on_read_var.is_initialized())) @combinations.generate(combinations.combine( @@ -1167,7 +1191,7 @@ class MirroredAndReplicaLocalVariableInitializerTest(test.TestCase): combinations.mirrored_strategy_with_gpu_and_cpu, combinations.core_mirrored_strategy_with_gpu_and_cpu], mode=["graph", "eager"])) -class ReplicaLocalVariableAssignTest(test.TestCase): +class SyncOnReadVariableAssignTest(test.TestCase): def testAssignReplicaLocalVarSumAggregation(self, distribution): def model_fn(): @@ -1178,24 +1202,23 @@ class ReplicaLocalVariableAssignTest(test.TestCase): return v_sum with distribution.scope(): - replica_local_var = distribution.extended.call_for_each_replica(model_fn) - self.assertTrue(isinstance(replica_local_var, - values.ReplicaLocalVariable)) + sync_on_read_var = distribution.extended.call_for_each_replica(model_fn) + self.assertIsInstance(sync_on_read_var, values.SyncOnReadVariable) self.evaluate(variables.global_variables_initializer()) # Each replica has a value of 1.0 assigned to it in replica context. # When we read the value using `read_var` we should see the SUM of each of # values on each of the replicas. self.assertEqual(2.0, self.evaluate( - distribution.extended.read_var(replica_local_var))) + distribution.extended.read_var(sync_on_read_var))) # Assigning 6.0 in cross replica context will assign a value of # 6.0/num_replicas to each replica. - tlv_ops = replica_local_var.assign(6.0) + tlv_ops = sync_on_read_var.assign(6.0) self.evaluate(tlv_ops) - # On reading the replica local var we should get the assigned value back. + # On reading the sync on read var we should get the assigned value back. # The value on all the replicas are added before being returned by # `read_var`. self.assertEqual(6.0, self.evaluate( - distribution.extended.read_var(replica_local_var))) + distribution.extended.read_var(sync_on_read_var))) def testAssignReplicaLocalVarMeanAggregation(self, distribution): def model_fn(): @@ -1206,21 +1229,20 @@ class ReplicaLocalVariableAssignTest(test.TestCase): return v_sum with distribution.scope(): - replica_local_var = distribution.extended.call_for_each_replica(model_fn) - self.assertTrue(isinstance(replica_local_var, - values.ReplicaLocalVariable)) + sync_on_read_var = distribution.extended.call_for_each_replica(model_fn) + self.assertIsInstance(sync_on_read_var, values.SyncOnReadVariable) self.evaluate(variables.global_variables_initializer()) # Each replica has a value of 1.0 assigned to it in replica context. # When we read the value using `read_var` we should see the MEAN of values # on all replicas which is the value assigned in replica context. self.assertEqual(1.0, self.evaluate( - distribution.extended.read_var(replica_local_var))) - tlv_ops = replica_local_var.assign(6.0) + distribution.extended.read_var(sync_on_read_var))) + tlv_ops = sync_on_read_var.assign(6.0) self.evaluate(tlv_ops) - # On reading the replica local var we should get the MEAN of all values + # On reading the sync on read var we should get the MEAN of all values # which is equal to the value assigned. self.assertEqual(6.0, self.evaluate( - distribution.extended.read_var(replica_local_var))) + distribution.extended.read_var(sync_on_read_var))) class MockModel(object): @@ -1515,6 +1537,25 @@ class MultiWorkerMirroredStrategyTestWithChief( strategy.configure(cluster_spec=self._cluster_spec) self._test_minimize_loss_graph(strategy, learning_rate=0.05) + def testMinimizeLossGraphCoreMirroredStrategyWithOneNode(self): + cluster_spec = {} + cluster_spec["chief"] = self._cluster_spec["chief"] + tf_config = {"cluster": cluster_spec} + with test.mock.patch.dict("os.environ", + {"TF_CONFIG": json.dumps(tf_config)}): + strategy = mirrored_strategy.CoreMirroredStrategy() + self.assertIsInstance(strategy.extended._inferred_cross_device_ops, + cross_device_ops_lib.NcclAllReduce) + self._test_minimize_loss_graph(strategy, learning_rate=0.05) + + def testInitializeFromTFConfig(self): + tf_config = {"cluster": self._cluster_spec} + with test.mock.patch.dict("os.environ", + {"TF_CONFIG": json.dumps(tf_config)}): + strategy = mirrored_strategy.CoreMirroredStrategy() + self.assertEqual( + max(context.num_gpus(), 1) * 3, strategy.num_replicas_in_sync) + def _replica_id(): replica_id = ds_context.get_replica_context().replica_id_in_sync_group diff --git a/tensorflow/contrib/distribute/python/multi_worker_test_base.py b/tensorflow/contrib/distribute/python/multi_worker_test_base.py index b05aac431f65b4281d9ed9c2fa95c210d55f4008..ce448840f14e3816f1d40328239256fd5acd51bf 100644 --- a/tensorflow/contrib/distribute/python/multi_worker_test_base.py +++ b/tensorflow/contrib/distribute/python/multi_worker_test_base.py @@ -37,12 +37,16 @@ except ImportError as _error: # pylint: disable=invalid-name from tensorflow.core.protobuf import config_pb2 from tensorflow.core.protobuf import rewriter_config_pb2 from tensorflow.python.client import session +from tensorflow.python.distribute import distribute_coordinator as dc from tensorflow.python.estimator import run_config from tensorflow.python.platform import test from tensorflow.python.platform import tf_logging as logging from tensorflow.python.training import coordinator from tensorflow.python.training import server_lib + +original_run_std_server = dc._run_std_server # pylint: disable=protected-access + ASSIGNED_PORTS = set() lock = threading.Lock() @@ -343,9 +347,9 @@ class MockOsEnv(collections.Mapping): def __iter__(self): if not hasattr(self._thread_local, 'dict'): self._thread_local.dict = dict() - for x in self._thread_local.dict.items(): + for x in self._thread_local.dict: yield x - for x in self._dict.items(): + for x in self._dict: yield x def __len__(self): @@ -357,6 +361,22 @@ class MockOsEnv(collections.Mapping): class IndependentWorkerTestBase(test.TestCase): """Testing infra for independent workers.""" + def _make_mock_run_std_server(self): + thread_local = threading.local() + + def _mock_run_std_server(*args, **kwargs): + ret = original_run_std_server(*args, **kwargs) + # Wait for all std servers to be brought up in order to reduce the chance + # of remote sessions taking local ports that have been assigned to std + # servers. Only call this barrier the first time this function is run for + # each thread. + if not getattr(thread_local, 'server_started', False): + self._barrier.wait() + thread_local.server_started = True + return ret + + return _mock_run_std_server + def setUp(self): self._mock_os_env = MockOsEnv() self._mock_context = test.mock.patch.object(os, 'environ', @@ -409,3 +429,25 @@ class IndependentWorkerTestBase(test.TestCase): def join_independent_workers(self, worker_threads): self._coord.join(worker_threads) + + +def get_tf_config_task(): + return json.loads(os.environ['TF_CONFIG'])['task'] + + +def get_tf_config_cluster_spec(): + return json.loads(os.environ['TF_CONFIG'])['cluster'] + + +def get_task_type(): + return get_tf_config_task()['type'] + + +def get_task_index(): + return get_tf_config_task()['index'] + + +def is_chief(): + return ('chief' not in get_tf_config_cluster_spec() + and get_task_type() == 'worker' + and get_task_index() == 0) diff --git a/tensorflow/contrib/distribute/python/parameter_server_strategy.py b/tensorflow/contrib/distribute/python/parameter_server_strategy.py index e42bc50fdc4e5e93c998708b0790fdea7768faf2..be863322256f7b5b93d91fa2e7ae1754b2494e3d 100644 --- a/tensorflow/contrib/distribute/python/parameter_server_strategy.py +++ b/tensorflow/contrib/distribute/python/parameter_server_strategy.py @@ -158,7 +158,7 @@ class ParameterServerExtended(CoreParameterServerExtended): cluster_spec=tfconfig.cluster_spec(), task_type=tfconfig.task_type, task_id=tfconfig.task_id, - num_accelerators=num_gpus_per_worker) + num_accelerators={'GPU': num_gpus_per_worker}) super(ParameterServerExtended, self).__init__( container_strategy, cluster_resolver=cluster_resolver) diff --git a/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py b/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py index 3de2041ae35775de6df5bca02c0f1d04a9c2f24e..18f6904959dd2afe3cf8e95b999b1b47dca03073 100644 --- a/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py +++ b/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py @@ -91,11 +91,11 @@ def create_test_objects(cluster_spec=None, cluster_spec=multi_worker_util.normalize_cluster_spec(cluster_spec), task_type=task_type, task_id=task_id, - num_accelerators=num_gpus) + num_accelerators={'GPU': num_gpus}) target = 'grpc://' + cluster_spec[WORKER][task_id] else: cluster_resolver = SimpleClusterResolver( - ClusterSpec({}), num_accelerators=num_gpus) + ClusterSpec({}), num_accelerators={'GPU': num_gpus}) target = '' distribution = MockCoreParameterServerStrategy(cluster_resolver) diff --git a/tensorflow/contrib/distribute/python/tpu_strategy.py b/tensorflow/contrib/distribute/python/tpu_strategy.py index 2d9d221f427422f8bbeba55c5644658af9a9a620..04e0af767bfaf94ed6a53ba9f8ed71ae4f9cdc4a 100644 --- a/tensorflow/contrib/distribute/python/tpu_strategy.py +++ b/tensorflow/contrib/distribute/python/tpu_strategy.py @@ -21,778 +21,6 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections -import copy - -from tensorflow.contrib.tpu.python.ops import tpu_ops -from tensorflow.contrib.tpu.python.tpu import device_assignment as device_assignment_lib -from tensorflow.contrib.tpu.python.tpu import functional as tpu_functional_ops -from tensorflow.contrib.tpu.python.tpu import topology -from tensorflow.contrib.tpu.python.tpu import tpu -from tensorflow.contrib.tpu.python.tpu import tpu_system_metadata as tpu_system_metadata_lib -from tensorflow.contrib.tpu.python.tpu import training_loop -from tensorflow.core.protobuf import config_pb2 -from tensorflow.python.client import session as session_lib -from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib -from tensorflow.python.distribute import device_util -from tensorflow.python.distribute import distribute_lib -from tensorflow.python.distribute import input_lib -from tensorflow.python.distribute import numpy_dataset -from tensorflow.python.distribute import reduce_util -from tensorflow.python.distribute import values -from tensorflow.python.distribute.cluster_resolver import TPUClusterResolver -from tensorflow.python.eager import context -from tensorflow.python.eager import function -from tensorflow.python.eager import tape -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import device as tf_device -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_util -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import control_flow_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import variable_scope as vs -from tensorflow.python.platform import tf_logging as logging -from tensorflow.python.util import compat -from tensorflow.python.util import nest - - -def initialize_tpu_system(cluster_resolver=None): - """Initialize the TPU devices in a separate session and graph. - - Args: - cluster_resolver: A tf.contrib.cluster_resolver.TPUClusterResolver, - which provides information about the TPU cluster. - Returns: - The tf.contrib.tpu.Topology object for the topology of the TPU cluster. - """ - if cluster_resolver is None: - cluster_resolver = TPUClusterResolver("") - master = cluster_resolver.master() - - logging.info("Initializing the TPU system.") - - if context.executing_eagerly(): - # This function looks as it is for the following non-intuitive reasons. - # tpu.initialize_system creates a dummy op whose sole purpose is to trigger - # DistributedTPURewritePass. This pass actually adds real ops that - # initialize the TPU system. Thus, we can't simply run tpu.initialize_system - # eagerly. We need to wrap it in defun and trigger the rewrite passes on it. - # The easiest way to trigger a rewrite is to run the function with - # TPUPartitionedCallOp. - @function.defun - def _tpu_init_fn(): - return tpu.initialize_system() - - # We can't call _tpu_init_fn normally (because it contains just a dummy op, - # see above) but need to define it to get it added to eager context - # and get its assigned name. - # pylint: disable=protected-access - graph_func = _tpu_init_fn._get_concrete_function_internal() - func_name = compat.as_str(graph_func._inference_function.name) - # pylint: enable=protected-access - - output = tpu_functional_ops.TPUPartitionedCall( - args=[], device_ordinal=0, Tout=[dtypes.string], f=func_name) - serialized_topology = output[0].numpy() - else: - session_config = config_pb2.ConfigProto(allow_soft_placement=True) - with ops.Graph().as_default(): - with session_lib.Session(config=session_config, target=master) as sess: - serialized_topology = sess.run(tpu.initialize_system()) - - logging.info("Finished initializing TPU system.") - return topology.Topology(serialized=serialized_topology) - - -def get_tpu_system_metadata(tpu_cluster_resolver): - """Retrieves TPU system metadata given a TPUClusterResolver.""" - master = tpu_cluster_resolver.master() - - # pylint: disable=protected-access - cluster_spec = tpu_cluster_resolver.cluster_spec() - cluster_def = cluster_spec.as_cluster_def() if cluster_spec else None - tpu_system_metadata = ( - tpu_system_metadata_lib._query_tpu_system_metadata( - master, - cluster_def=cluster_def, - query_topology=False)) - - return tpu_system_metadata - - -# TODO(jhseu): Deduplicate with MirroredStrategy? -def _create_tpu_mirrored_variable( # pylint: disable=missing-docstring - strategy, device_map, logical_device, real_mirrored_creator, - *args, **kwargs): - # Figure out what collections this variable should be added to. - # We'll add the TPUMirroredVariable to those collections instead. - var_collections = kwargs.pop("collections", None) - if var_collections is None: - var_collections = [ops.GraphKeys.GLOBAL_VARIABLES] - kwargs["collections"] = [] - - # TODO(jhseu): Should we have different behavior for different - # synchronization settings? - - # Get aggregation value - # TODO(jhseu): Support aggregation in a replica context. - aggregation = kwargs.pop("aggregation", vs.VariableAggregation.NONE) - if aggregation not in [ - vs.VariableAggregation.NONE, - vs.VariableAggregation.SUM, - vs.VariableAggregation.MEAN, - vs.VariableAggregation.ONLY_FIRST_REPLICA, - ]: - raise ValueError("Invalid variable aggregation mode: {} for variable: {}" - .format(aggregation, kwargs["name"])) - - # Ignore user-specified caching device, not needed for mirrored variables. - kwargs.pop("caching_device", None) - - # TODO(josh11b,apassos): It would be better if variable initialization - # was never recorded on the tape instead of having to do this manually - # here. - with tape.stop_recording(): - devices = device_map.logical_to_actual_devices(logical_device) - value_list = real_mirrored_creator(devices, *args, **kwargs) - result = values.TPUMirroredVariable( - strategy, device_map, value_list, aggregation, - logical_device=logical_device) - - if not (context.executing_eagerly() or ops.inside_function()): - g = ops.get_default_graph() - # If "trainable" is True, next_creator() will add the member variables - # to the TRAINABLE_VARIABLES collection, so we manually remove - # them and replace with the MirroredVariable. We can't set - # "trainable" to False for next_creator() since that causes functions - # like implicit_gradients to skip those variables. - if kwargs.get("trainable", True): - var_collections.append(ops.GraphKeys.TRAINABLE_VARIABLES) - l = g.get_collection_ref(ops.GraphKeys.TRAINABLE_VARIABLES) - for v in value_list: - l.remove(v) - g.add_to_collections(var_collections, result) - return result - - -class TPUStrategy(distribute_lib.DistributionStrategy): - """TPU distribution strategy implementation.""" - - def __init__(self, - tpu_cluster_resolver=None, - steps_per_run=None, - device_assignment=None): - """Initializes the TPUStrategy object. - - Args: - tpu_cluster_resolver: A tf.contrib.cluster_resolver.TPUClusterResolver, - which provides information about the TPU cluster. - steps_per_run: Number of steps to run on device before returning to the - host. Note that this can have side-effects on performance, hooks, - metrics, summaries etc. - This parameter is only used when Distribution Strategy is used with - estimator or keras. - device_assignment: Optional `tf.contrib.tpu.DeviceAssignment` to specify - the placement of replicas on the TPU cluster. Currently only supports - the usecase of using a single core within a TPU cluster. - """ - super(TPUStrategy, self).__init__(TPUExtended( - self, tpu_cluster_resolver, steps_per_run, device_assignment)) - - @property - def steps_per_run(self): - """DEPRECATED: use .extended.steps_per_run instead.""" - return self._extended.steps_per_run - - # TODO(cjfj): Modify `_call_for_each_replica` in `TPUExtended` such that this - # can use the default implementation. - # This implementation runs a single step. It does not use infeed or outfeed. - def experimental_run(self, fn, input_iterator=None): - """See base class.""" - if context.executing_eagerly() and not ops.inside_function(): - raise NotImplementedError( - "Eager mode not supported in TPUStrategy outside TF functions.") - - if input_iterator is None: - inputs = [] - else: - inputs = input_iterator.get_next() - - result = [None] - def replicated_fn(replica_id, replica_input): - """Wraps user function to provide replica ID and `Tensor` inputs.""" - with _TPUReplicaContext(self, replica_id_in_sync_group=replica_id): - if input_iterator is None: - result[0] = fn() - else: - result[0] = fn(replica_input) - return result[0] - - replicate_inputs = [] # By replica. - for i in range(self.num_replicas_in_sync): - replicate_inputs.append( - [constant_op.constant(i, dtype=dtypes.int32), - values.select_replica(i, inputs)]) - - with self.scope(): - replicate_outputs = tpu.replicate(replicated_fn, replicate_inputs) - - # Workaround for `tpu.replicate` behaviour when single `Tensor` returned. - replicate_outputs = [ - nest.pack_sequence_as(result[0], nest.flatten(replica_outputs)) - for replica_outputs in replicate_outputs] - - device_map = self.extended._device_map # pylint: disable=protected-access - return values.regroup(device_map, replicate_outputs) - - -class TPUExtended(distribute_lib.DistributionStrategyExtended): - """Implementation of TPUStrategy.""" - - def __init__(self, - container_strategy, - tpu_cluster_resolver=None, - steps_per_run=None, - device_assignment=None): - super(TPUExtended, self).__init__(container_strategy) - - if tpu_cluster_resolver is None: - tpu_cluster_resolver = TPUClusterResolver("") - - if steps_per_run is None: - # TODO(frankchn): Warn when we are being used by DS/Keras and this is - # not specified. - steps_per_run = 1 - - self._tpu_cluster_resolver = tpu_cluster_resolver - self._tpu_metadata = get_tpu_system_metadata(self._tpu_cluster_resolver) - self._device_assignment = device_assignment - - # Device assignment is currently only supported for 1 core case. - if self._device_assignment: - assert isinstance(self._device_assignment, - device_assignment_lib.DeviceAssignment) - if self._device_assignment.num_replicas != 1: - raise ValueError("Device assignment is only supported for a single " - "core single replica case currently.") - if self._device_assignment.num_cores_per_replica != 1: - raise ValueError("Device assignment is only supported for a single " - "core single replica case currently.") - if not all(self._device_assignment.core_assignment[0][0] == [0, 0, 0]): - raise ValueError("Device assignment is only supported for a single " - "core single replica case currently.") - - # TODO(jhseu): Switch to DeviceAssignment to support pods and model - # parallelism. - self._device_index = { - d.name: i for i, d in enumerate(self._tpu_metadata.devices) - if "device:TPU:" in d.name - } - self._host_device = self.get_host_cpu_device(0) - self._tpu_devices = tuple(sorted(self._device_index.keys())) - # Only create variables for the number of replicas we're running. - self._tpu_devices = self._tpu_devices[:self._num_replicas_in_sync] - self._device_map = values.ReplicaDeviceMap(self._tpu_devices) - - # Preload the data onto the TPUs. - input_worker_devices = collections.OrderedDict() - for tpu_device in self._tpu_devices: - host_device = _get_host_for_device(tpu_device) - input_worker_devices.setdefault(host_device, []) - input_worker_devices[host_device].append(tpu_device) - self._input_workers = input_lib.InputWorkers( - self._device_map, tuple(input_worker_devices.items())) - - # TODO(sourabhbajaj): Remove this once performance of running one step - # at a time is comparable to multiple steps. - self.steps_per_run = steps_per_run - self._require_static_shapes = True - - def _validate_colocate_with_variable(self, colocate_with_variable): - values.validate_colocate_tpu_variable(colocate_with_variable, self) - - def _get_enqueue_op_per_host(self, host_id, multi_worker_iterator, - input_shapes, iterations): - """Create an enqueue op for a single host identified using host_id. - - The while_loop op returned will run `iterations` times and in each run - enqueue batches for each shard. - - Args: - host_id: integer, id of the host to run the enqueue ops on. - multi_worker_iterator: MultiWorkerDataIterator to read the input data. - input_shapes: shape of inputs to be enqueue on the queue. This is same as - the value of `nest.flatten(iterator.output_shapes)`. - iterations: integer, number of iterations to be run; determines the - number of batches to be enqueued. - - Returns: - while_loop_op running `iterations` times; in each run we enqueue a batch - on the infeed queue from the host with id `host_id` for each device shard. - """ - host = self.get_host_cpu_device(host_id) - # TODO(sourabhbajaj): Possibly make changes to MultiWorkerDataset - # to work with TPU Prefetch so clean up this code. - iterator = ( - multi_worker_iterator.get_iterator(self.get_host(host_id))._iterator) # pylint: disable=protected-access - - def _infeed_enqueue_ops_fn(): - """Enqueue ops for one iteration.""" - control_deps = [] - sharded_inputs = [] - enqueue_ops = [] - - with ops.device(host): - for _ in range(self.num_replicas_per_host): - # Use control dependencies to ensure a deterministic ordering. - with ops.control_dependencies(control_deps): - inputs = nest.flatten(iterator.get_next()) - control_deps.extend(inputs) - sharded_inputs.append(inputs) - - for core_id, shard_input in enumerate(sharded_inputs): - enqueue_ops.append( - tpu_ops.infeed_enqueue_tuple( - inputs=shard_input, - shapes=input_shapes, - device_ordinal=core_id)) - return enqueue_ops - - def enqueue_ops_loop_body(i): - """Callable for the loop body of the while_loop instantiated below.""" - with ops.control_dependencies(_infeed_enqueue_ops_fn()): - return i + 1 - - with ops.device(host): - enqueue_op_per_host = control_flow_ops.while_loop( - lambda i: i < iterations, - enqueue_ops_loop_body, - [constant_op.constant(0)], - parallel_iterations=1) - - return enqueue_op_per_host - - def _make_dataset_iterator(self, dataset): - """Make iterators for each of the TPU hosts.""" - return input_lib.DatasetIterator(dataset, self._input_workers, - self._num_replicas_in_sync) - - def _make_input_fn_iterator( - self, - input_fn, - replication_mode=distribute_lib.InputReplicationMode.PER_WORKER): - input_contexts = [] - num_workers = self._input_workers.num_workers - for i in range(num_workers): - input_contexts.append(distribute_lib.InputContext( - num_input_pipelines=num_workers, - input_pipeline_id=i, - num_replicas_in_sync=self._num_replicas_in_sync)) - return input_lib.InputFunctionIterator( - input_fn, self._input_workers, input_contexts) - - def _experimental_make_numpy_dataset(self, numpy_input, session): - return numpy_dataset.one_host_numpy_dataset( - numpy_input, numpy_dataset.SingleDevice(self.get_host_cpu_device(0)), - session) - - # TODO(priyag): Deal with OutOfRange errors once b/111349762 is fixed. - # TODO(sourabhbajaj): Remove the initial_loop_values parameter when we have - # a mechanism to infer the outputs of `fn`. Pending b/110550782. - def _experimental_run_steps_on_iterator( - self, fn, multi_worker_iterator, iterations, initial_loop_values=None): - output_shapes = multi_worker_iterator.output_shapes - shapes = nest.flatten(output_shapes) - if any(not s.is_fully_defined() for s in shapes): - raise ValueError( - "TPU currently requires fully defined shapes. Either use " - "set_shape() on the input tensors or use " - "dataset.batch(..., drop_remainder=True).") - - # Wrap `fn` for repeat. - if initial_loop_values is None: - initial_loop_values = {} - initial_loop_values = nest.flatten(initial_loop_values) - ctx = input_lib.MultiStepContext() - - def run_fn(inputs): - """Single step on the TPU device.""" - fn_result = fn(ctx, inputs) - flat_last_step_outputs = nest.flatten(ctx.last_step_outputs) - if flat_last_step_outputs: - with ops.control_dependencies([fn_result]): - return [array_ops.identity(f) for f in flat_last_step_outputs] - else: - return fn_result - - # We capture the control_flow_context at this point, before we run `fn` - # inside a while_loop and TPU replicate context. This is useful in cases - # where we might need to exit these contexts and get back to the outer - # context to do some things, for e.g. create an op which should be - # evaluated only once at the end of the loop on the host. One such usage - # is in creating metrics' value op. - self._outer_control_flow_context = ( - ops.get_default_graph()._get_control_flow_context()) # pylint: disable=protected-access - - def rewrite_fn(*args): - """The rewritten step fn running on TPU.""" - del args - - per_replica_inputs = multi_worker_iterator.get_next() - replicate_inputs = [] - for replica_id in range(self._num_replicas_in_sync): - select_replica = lambda x: values.select_replica(replica_id, x) # pylint: disable=cell-var-from-loop - replicate_inputs.append((nest.map_structure( - select_replica, per_replica_inputs),)) - - replicate_outputs = tpu.replicate(run_fn, replicate_inputs) - - # If run_fn has tensor outputs, tpu.replicate returns a list of list. We - # will flatten it in this case. If run_fn has no tensor outputs, - # tpu.replicate returns a list of no_ops, we will keep the output as it - # is. - if isinstance(replicate_outputs[0], list): - replicate_outputs = nest.flatten(replicate_outputs) - - return replicate_outputs - - # TODO(sourabhbajaj): The input to while loop should be based on the - # output type of the step_fn - assert isinstance(initial_loop_values, list) - initial_loop_values = initial_loop_values * self._num_replicas_in_sync - - # Put the while loop op on host 0. - with ops.device(self.get_host_cpu_device(0)): - replicate_outputs = training_loop.repeat(iterations, rewrite_fn, - initial_loop_values) - - del self._outer_control_flow_context - ctx.run_op = control_flow_ops.group(replicate_outputs) - - if isinstance(replicate_outputs, list): - # Filter out any ops from the outputs, typically this would be the case - # when there were no tensor outputs. - last_step_tensor_outputs = [ - x for x in replicate_outputs if not isinstance(x, ops.Operation) - ] - - # Outputs are currently of the structure (flattened) - # [output0_device0, output1_device0, output2_device0, - # output0_device1, output1_device1, output2_device1, - # ...] - # Convert this to the following structure instead: (grouped by output) - # [[output0_device0, output0_device1], - # [output1_device0, output1_device1], - # [output2_device0, output2_device1]] - output_num = len(last_step_tensor_outputs) // self._num_replicas_in_sync - last_step_tensor_outputs = [ - last_step_tensor_outputs[i::output_num] for i in range(output_num) - ] - else: - # no tensors returned. - last_step_tensor_outputs = [] - - _set_last_step_outputs(ctx, last_step_tensor_outputs) - return ctx - - def _call_for_each_replica(self, fn, args, kwargs): - # TODO(jhseu): Consider making it so call_for_each_replica implies that - # we're in a tpu.rewrite(), and update TPUMirroredVariable accordingly. - with _TPUReplicaContext(self._container_strategy()): - return fn(*args, **kwargs) - - def _experimental_initialize_system(self): - """Experimental method added to be used by Estimator. - - This is a private method only to be used by Estimator. Other frameworks - should directly be calling `tf.contrib.distribute.initialize_tpu_system` - """ - initialize_tpu_system(self._tpu_cluster_resolver) - - def _create_variable(self, next_creator, *args, **kwargs): - """Create a TPUMirroredVariable. See `DistributionStrategy.scope`.""" - colocate_with = kwargs.pop("colocate_with", None) - if colocate_with is None: - device_map = self._device_map - logical_device = 0 # TODO(josh11b): Get logical device from scope here. - elif isinstance(colocate_with, numpy_dataset.SingleDevice): - with ops.device(colocate_with.device): - return next_creator(*args, **kwargs) - else: - device_map = colocate_with.device_map - logical_device = colocate_with.logical_device - - def _real_mirrored_creator(devices, *args, **kwargs): # pylint: disable=g-missing-docstring - value_list = [] - for i, d in enumerate(devices): - with ops.device(d): - if i > 0: - # Give replicas meaningful distinct names: - var0name = value_list[0].name.split(":")[0] - # We append a / to variable names created on replicas with id > 0 to - # ensure that we ignore the name scope and instead use the given - # name as the absolute name of the variable. - kwargs["name"] = "%s/replica_%d/" % (var0name, i) - # Initialize replicas with the same value: - if context.executing_eagerly() or ops.inside_function(): - with ops.init_scope(): - kwargs["initial_value"] = array_ops.identity( - value_list[0].value()) - else: - def initial_value_fn(device=d): - with ops.device(device): - return array_ops.identity(value_list[0].initial_value) - kwargs["initial_value"] = initial_value_fn - with context.context().device_policy(context.DEVICE_PLACEMENT_SILENT): - v = next_creator(*args, **kwargs) - assert not isinstance(v, values.TPUMirroredVariable) - value_list.append(v) - return value_list - - return _create_tpu_mirrored_variable( - self._container_strategy(), device_map, logical_device, - _real_mirrored_creator, *args, **kwargs) - - def _reduce_to(self, reduce_op, value, destinations): - if values._enclosing_tpu_context() is not None: # pylint: disable=protected-access - if reduce_op == reduce_util.ReduceOp.MEAN: - # TODO(jhseu): Revisit once we support model-parallelism. - value *= (1. / self._num_replicas_in_sync) - elif reduce_op != reduce_util.ReduceOp.SUM: - raise NotImplementedError( - "Currently only support sum & mean in TPUStrategy.") - return tpu_ops.cross_replica_sum(value) - - if not isinstance(value, values.DistributedValues): - # This function handles reducing values that are not PerReplica or - # Mirrored values. For example, the same value could be present on all - # replicas in which case `value` would be a single value or value could - # be 0. - return cross_device_ops_lib.reduce_non_distributed_value( - reduce_op, self._device_map, value, destinations) - - devices = cross_device_ops_lib.get_devices_from(destinations) - if len(devices) != 1: - raise ValueError("Multiple devices are not supported for TPUStrategy") - - # Always performs the reduction on the TPU host. - with ops.device(self._host_device): - output = math_ops.add_n(value.values) - if reduce_op == reduce_util.ReduceOp.MEAN: - output *= (1. / len(value.values)) - - # If necessary, copy to requested destination. - dest_canonical = device_util.canonicalize(devices[0]) - host_canonical = device_util.canonicalize(self._host_device) - - if dest_canonical != host_canonical: - with ops.device(devices[0]): - output = array_ops.identity(output) - - return output - - def _update(self, var, fn, args, kwargs, group): - assert isinstance(var, values.TPUMirroredVariable) - if values._enclosing_tpu_context() is not None: # pylint: disable=protected-access - if group: - return fn(var, *args, **kwargs) - else: - return (fn(var, *args, **kwargs),) - - # Otherwise, we revert to MirroredStrategy behavior and update each variable - # directly. - updates = [] - for i, (d, v) in enumerate(zip(var.devices, var.values)): - name = "update_%d" % i - with ops.device(d), distribute_lib.UpdateContext(d), ops.name_scope(name): - # If args and kwargs are not mirrored, the value is returned as is. - updates.append(fn(v, - *values.select_device_mirrored(d, args), - **values.select_device_mirrored(d, kwargs))) - return values.update_regroup(self, self._device_map, updates, group) - - def read_var(self, var): - assert isinstance(var, values.TPUMirroredVariable) - return var.read_value() - - def _unwrap(self, val): - if isinstance(val, values.DistributedValues): - # Return in a deterministic order. - return tuple(val.get(device=d) for d in sorted(val.devices)) - elif isinstance(val, list): - # TODO(josh11b): We need to remove this case; per device values should - # be represented using a PerReplica wrapper instead of a list with - # one entry per device. - return tuple(val) - elif isinstance(val, values.TPUMirroredVariable): - # pylint: disable=protected-access - if values._enclosing_tpu_context() is not None: - return (val,) - return val.values - return (val,) - - def value_container(self, value): - return value - - def _broadcast_to(self, tensor, destinations): - del destinations - return tensor - - @property - def num_hosts(self): - if self._device_assignment is None: - return self._tpu_metadata.num_hosts - - return len(set([self._device_assignment.host_device(r) - for r in range(self._device_assignment.num_replicas)])) - - @property - def num_replicas_per_host(self): - if self._device_assignment is None: - return self._tpu_metadata.num_of_cores_per_host - - # TODO(sourabhbajaj): Remove this method we use inputs and remove infeed - # as the computation of num_replicas_per_host is not a constant - # when using device_assignment. This is a temporary workaround to support - # StatefulRNN as everything is 1 in that case. - # This method needs to take host_id as input for correct computation. - max_models_per_host = (self._tpu_metadata.num_of_cores_per_host // - self._device_assignment.num_cores_per_replica) - models_per_host = min(self._device_assignment.num_replicas, - max_models_per_host) - return models_per_host * self._device_assignment.num_cores_per_replica - - @property - def _num_replicas_in_sync(self): - if self._device_assignment is None: - return self._tpu_metadata.num_cores - return (self._device_assignment.num_replicas * - self._device_assignment.num_cores_per_replica) - - @property - def experimental_between_graph(self): - return False - - @property - def experimental_should_init(self): - return True - - @property - def should_checkpoint(self): - return True - - @property - def should_save_summary(self): - return True - - @property - def worker_devices(self): - return self._tpu_devices - - @property - def parameter_devices(self): - return self._tpu_devices - - def non_slot_devices(self, var_list): - return self._host_device - - def _update_non_slot(self, colocate_with, fn, args, kwargs, group): - del colocate_with - with ops.device(self._host_device), distribute_lib.UpdateContext( - self._host_device): - result = fn(*args, **kwargs) - if group: - return result - else: - return nest.map_structure(self._unwrap, result) - - def get_host(self, host_id): - if self._tpu_cluster_resolver.get_master() in ("", "local"): - return "/replica:0/task:0" - job_name = self._tpu_cluster_resolver.get_job_name() or "tpu_worker" - return "/job:%s/task:%d" % (job_name, host_id) - - def get_host_cpu_device(self, host_id): - return self.get_host(host_id) + "/device:CPU:0" - - def _configure(self, - session_config=None, - cluster_spec=None, - task_type=None, - task_id=None): - del cluster_spec, task_type, task_id - if session_config: - session_config.CopyFrom(self._update_config_proto(session_config)) - - def _update_config_proto(self, config_proto): - updated_config = copy.deepcopy(config_proto) - updated_config.isolate_session_state = True - cluster_spec = self._tpu_cluster_resolver.cluster_spec() - if cluster_spec: - updated_config.cluster_def.CopyFrom(cluster_spec.as_cluster_def()) - return updated_config - - # TODO(priyag): Delete this once all strategies use global batch size. - @property - def _global_batch_size(self): - """`make_dataset_iterator` and `make_numpy_iterator` use global batch size. - - `make_input_fn_iterator` assumes per-replica batching. - - Returns: - Boolean. - """ - return True - - -class _TPUReplicaContext(distribute_lib.ReplicaContext): - """Replication Context class for TPU Strategy.""" - - # TODO(sourabhbajaj): Call for each replica should be updating this. - # TODO(b/118385803): Always properly initialize replica_id. - def __init__(self, strategy, replica_id_in_sync_group=None): - if replica_id_in_sync_group is None: - replica_id_in_sync_group = constant_op.constant(0, dtypes.int32) - distribute_lib.ReplicaContext.__init__( - self, strategy, replica_id_in_sync_group=replica_id_in_sync_group) - - @property - def devices(self): - distribute_lib.require_replica_context(self) - ds = self._strategy - replica_id = tensor_util.constant_value(self._replica_id_in_sync_group) - - if replica_id is None: # Non-constant `Tensor` inside `tpu.replicate`. - # TODO(cjfj): Return other devices when model parallelism is supported. - return (tpu.core(0),) - else: - return (ds.extended.worker_devices[replica_id],) - - -def _get_host_for_device(device): - spec = tf_device.DeviceSpec.from_string(device) - return tf_device.DeviceSpec( - job=spec.job, replica=spec.replica, task=spec.task, - device_type="CPU", device_index=0).to_string() - - -def _set_last_step_outputs(ctx, last_step_tensor_outputs): - """Sets the last step outputs on the given context.""" - # Convert replicate_outputs to the original dict structure of - # last_step_outputs. - last_step_tensor_outputs_dict = nest.pack_sequence_as( - ctx.last_step_outputs, last_step_tensor_outputs) - - for name, reduce_op in ctx._last_step_outputs_reduce_ops.items(): # pylint: disable=protected-access - output = last_step_tensor_outputs_dict[name] - # For outputs that have already been reduced, take the first value - # from the list as each value should be the same. Else return the full - # list of values. - # TODO(josh11b): If reduce_op is NONE, we should return a PerReplica - # value. - if reduce_op is not None: - # TODO(priyag): Should this return the element or a list with 1 element - last_step_tensor_outputs_dict[name] = output[0] - ctx._set_last_step_outputs(last_step_tensor_outputs_dict) # pylint: disable=protected-access +# pylint: disable=unused-import +from tensorflow.python.distribute.tpu_strategy import TPUStrategy +from tensorflow.python.tpu.tpu_strategy_util import initialize_tpu_system diff --git a/tensorflow/contrib/distribute/python/values_test.py b/tensorflow/contrib/distribute/python/values_test.py index 51c58b0b2f3dc2ab63e22718825a471b8657f892..f9bd3078ef69eaf8c989c500161cce8d751ed23e 100644 --- a/tensorflow/contrib/distribute/python/values_test.py +++ b/tensorflow/contrib/distribute/python/values_test.py @@ -522,11 +522,11 @@ def _make_replica_local(method, strategy=None): with ops.device(d): v.append(variable_scope.get_variable( name=n, initializer=init, use_resource=True)) - replica_local = values.ReplicaLocalVariable(strategy, device_map, v, method) + replica_local = values.SyncOnReadVariable(strategy, device_map, v, method) return v, replica_local -class ReplicaLocalVariablePropertiesTest(test.TestCase): +class SyncOnReadVariablePropertiesTest(test.TestCase): config = config_pb2.ConfigProto() config.allow_soft_placement = True @@ -549,7 +549,7 @@ class ReplicaLocalVariablePropertiesTest(test.TestCase): v = variable_scope.get_variable( name="v", initializer=[1.], use_resource=True) device_map = values.ReplicaDeviceMap(("/job:foo/device:CPU:0",)) - replica_local = values.ReplicaLocalVariable( + replica_local = values.SyncOnReadVariable( None, device_map, (v,), variable_scope.VariableAggregation.MEAN) self.assertEqual(v.name, replica_local.name) @@ -577,7 +577,7 @@ class ReplicaLocalVariablePropertiesTest(test.TestCase): combinations.mirrored_strategy_with_gpu_and_cpu, combinations.core_mirrored_strategy_with_gpu_and_cpu], mode=["graph", "eager"])) -class ReplicaLocalVariableTest(test.TestCase, parameterized.TestCase): +class SyncOnReadVariableTest(test.TestCase, parameterized.TestCase): def _assign_replica_local(self, devices, v, new): for d, var, n in zip(devices, v, new): @@ -656,7 +656,8 @@ class ReplicaLocalVariableTest(test.TestCase, parameterized.TestCase): def _save_replica_local_sum(self, distribution): """Save variables with mirroring, returns save_path.""" with self.session(graph=ops.Graph()) as sess: - v, replica_local = _make_replica_local("sum", distribution) + v, replica_local = _make_replica_local( + variable_scope.VariableAggregation.SUM, distribution) # Overwrite the initial values. self._assign_replica_local(_devices, v, [1.5, 2.]) diff --git a/tensorflow/contrib/eager/python/examples/nmt_with_attention/nmt_with_attention.ipynb b/tensorflow/contrib/eager/python/examples/nmt_with_attention/nmt_with_attention.ipynb index 436e887736158ec1ba8e46eac8de4ac7b8e6be01..512605a17eb77a85a5ec98197f4ed8fda6863932 100644 --- a/tensorflow/contrib/eager/python/examples/nmt_with_attention/nmt_with_attention.ipynb +++ b/tensorflow/contrib/eager/python/examples/nmt_with_attention/nmt_with_attention.ipynb @@ -688,7 +688,7 @@ " for t in range(max_length_targ):\n", " predictions, dec_hidden, attention_weights = decoder(dec_input, dec_hidden, enc_out)\n", " \n", - " # storing the attention weigths to plot later on\n", + " # storing the attention weights to plot later on\n", " attention_weights = tf.reshape(attention_weights, (-1, ))\n", " attention_plot[t] = attention_weights.numpy()\n", "\n", @@ -842,4 +842,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/tensorflow/contrib/estimator/python/estimator/early_stopping.py b/tensorflow/contrib/estimator/python/estimator/early_stopping.py index 11856ece38bf08dfdf16e8b0d9890bbfb0033216..47f568ed3d3e1b94e74c1423f774352df5c30f45 100644 --- a/tensorflow/contrib/estimator/python/estimator/early_stopping.py +++ b/tensorflow/contrib/estimator/python/estimator/early_stopping.py @@ -23,7 +23,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow_estimator.contrib.estimator.python.estimator import early_stopping +from tensorflow_estimator.python.estimator import early_stopping # Include attrs that start with single underscore. _HAS_DYNAMIC_ATTRIBUTES = True @@ -31,4 +31,4 @@ early_stopping.__all__ = [ s for s in dir(early_stopping) if not s.startswith('__') ] -from tensorflow_estimator.contrib.estimator.python.estimator.early_stopping import * +from tensorflow_estimator.python.estimator.early_stopping import * diff --git a/tensorflow/contrib/gan/README.md b/tensorflow/contrib/gan/README.md index 4eac4e80cdacd779fdbedef19e4a654196f0caf1..3c1d814e70f7fdad4083583c9d89450a60bc2e20 100644 --- a/tensorflow/contrib/gan/README.md +++ b/tensorflow/contrib/gan/README.md @@ -9,8 +9,9 @@ explicitly model the distribution and without writing an explicit loss. For example, the generator could learn to draw samples from the distribution of natural images. For more details on this technique, see ['Generative Adversarial Networks'](https://arxiv.org/abs/1406.2661) by -Goodfellow et al. See [tensorflow/models](https://github.com/tensorflow/models/tree/master/research/gan/) for examples, and [this tutorial](http://https://github.com/tensorflow/models/tree/master/research/gan/tutorial.ipynb) for an -introduction. +Goodfellow et al. See +[tensorflow/models](https://github.com/tensorflow/models/tree/master/research/gan/) +for examples, and [this tutorial](https://github.com/tensorflow/models/tree/master/research/gan/tutorial.ipynb) for an introduction. #### Usage ```python diff --git a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py index 5b9c54e43a16adf457d5ed0e7e73dcd168ab0d67..66af79d1e81bbc450141673dd54d865e5c7932d5 100644 --- a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py +++ b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py @@ -23,7 +23,6 @@ import tempfile from absl.testing import parameterized import numpy as np -import six from tensorflow.contrib import layers from tensorflow.contrib.gan.python import namedtuples as tfgan_tuples @@ -238,10 +237,10 @@ class GANEstimatorIntegrationTest(test.TestCase): # Evaluate. scores = est.evaluate(eval_input_fn) self.assertEqual(num_steps, scores[ops.GraphKeys.GLOBAL_STEP]) - self.assertIn('loss', six.iterkeys(scores)) + self.assertIn('loss', scores) self.assertEqual(scores['discriminator_loss'] + scores['generator_loss'], scores['loss']) - self.assertIn('mse_custom_metric', six.iterkeys(scores)) + self.assertIn('mse_custom_metric', scores) # Predict. predictions = np.array([x for x in est.predict(predict_input_fn)]) diff --git a/tensorflow/contrib/gan/python/estimator/python/stargan_estimator_test.py b/tensorflow/contrib/gan/python/estimator/python/stargan_estimator_test.py index c00ff4399748a77f88d9753df7592bf3859d754e..0fcd1b7924eb02f5d617b45af16852baf2e2bb48 100644 --- a/tensorflow/contrib/gan/python/estimator/python/stargan_estimator_test.py +++ b/tensorflow/contrib/gan/python/estimator/python/stargan_estimator_test.py @@ -23,7 +23,6 @@ import tempfile from absl.testing import parameterized import numpy as np -import six from tensorflow.contrib import layers from tensorflow.contrib.gan.python import namedtuples as tfgan_tuples @@ -235,10 +234,10 @@ class StarGANEstimatorIntegrationTest(test.TestCase): # EVALUTE scores = est.evaluate(eval_input_fn) self.assertEqual(num_steps, scores[ops.GraphKeys.GLOBAL_STEP]) - self.assertIn('loss', six.iterkeys(scores)) + self.assertIn('loss', scores) self.assertEqual(scores['discriminator_loss'] + scores['generator_loss'], scores['loss']) - self.assertIn('mse_custom_metric', six.iterkeys(scores)) + self.assertIn('mse_custom_metric', scores) # PREDICT predictions = np.array([x for x in est.predict(predict_input_fn)]) diff --git a/tensorflow/contrib/gan/python/estimator/python/tpu_gan_estimator_test.py b/tensorflow/contrib/gan/python/estimator/python/tpu_gan_estimator_test.py index 0d9e6489bdd1d89cc49bfedc2eed784999c31d2b..baf2c28df4b63cff525dcf3ff880730768ad000a 100644 --- a/tensorflow/contrib/gan/python/estimator/python/tpu_gan_estimator_test.py +++ b/tensorflow/contrib/gan/python/estimator/python/tpu_gan_estimator_test.py @@ -23,7 +23,6 @@ import tempfile from absl.testing import parameterized import numpy as np -import six from tensorflow.contrib import layers from tensorflow.contrib.gan.python import namedtuples as tfgan_tuples @@ -184,11 +183,11 @@ class TPUGANEstimatorIntegrationTest(test.TestCase, parameterized.TestCase): # Evaluate. num_steps_eval = 2 scores = est.evaluate(eval_input_fn, steps=num_steps_eval) - self.assertIn(ops.GraphKeys.GLOBAL_STEP, six.iterkeys(scores)) - self.assertIn('loss', six.iterkeys(scores)) + self.assertIn(ops.GraphKeys.GLOBAL_STEP, scores) + self.assertIn('loss', scores) self.assertEqual(scores['discriminator_loss'] + scores['generator_loss'], scores['loss']) - self.assertIn('mse_custom_metric', six.iterkeys(scores)) + self.assertIn('mse_custom_metric', scores) # Predict. predictions = np.array([x['generated_data'] for x in diff --git a/tensorflow/contrib/gan/python/features/python/spectral_normalization_impl.py b/tensorflow/contrib/gan/python/features/python/spectral_normalization_impl.py index 0cc653f0a7907f407e66add5537d1e0a5adb6d8b..3764c43cdfc8f6515e0376cd6aa1d244b21e2e89 100644 --- a/tensorflow/contrib/gan/python/features/python/spectral_normalization_impl.py +++ b/tensorflow/contrib/gan/python/features/python/spectral_normalization_impl.py @@ -53,7 +53,7 @@ def compute_spectral_norm(w_tensor, power_iteration_rounds=1, name=None): Args: w_tensor: The weight matrix whose spectral norm should be computed. power_iteration_rounds: The number of iterations of the power method to - perform. A higher number yeilds a better approximation. + perform. A higher number yields a better approximation. name: An optional scope name. Returns: @@ -105,7 +105,7 @@ def spectral_normalize(w, power_iteration_rounds=1, name=None): Args: w: The weight matrix to be normalized. power_iteration_rounds: The number of iterations of the power method to - perform. A higher number yeilds a better approximation. + perform. A higher number yields a better approximation. name: An optional scope name. Returns: @@ -126,7 +126,7 @@ def spectral_norm_regularizer(scale, power_iteration_rounds=1, scope=None): Args: scale: A scalar multiplier. 0.0 disables the regularizer. power_iteration_rounds: The number of iterations of the power method to - perform. A higher number yeilds a better approximation. + perform. A higher number yields a better approximation. scope: An optional scope name. Returns: @@ -221,7 +221,7 @@ def spectral_normalization_custom_getter(name_filter=_default_name_filter, name_filter: Optionally, a method that takes a Variable name as input and returns whether this Variable should be normalized. power_iteration_rounds: The number of iterations of the power method to - perform per step. A higher number yeilds a better approximation of the + perform per step. A higher number yields a better approximation of the true spectral norm. Returns: @@ -294,7 +294,7 @@ def keras_spectral_normalization(name_filter=_default_name_filter, name_filter: Optionally, a method that takes a Variable name as input and returns whether this Variable should be normalized. power_iteration_rounds: The number of iterations of the power method to - perform per step. A higher number yeilds a better approximation of the + perform per step. A higher number yields a better approximation of the true spectral norm. Yields: diff --git a/tensorflow/contrib/graph_editor/transform.py b/tensorflow/contrib/graph_editor/transform.py index e79ccd8da1f8952758ae322d3a92dec34910a9db..5b37239665d46db38fc249e9004d2200abb3d610 100644 --- a/tensorflow/contrib/graph_editor/transform.py +++ b/tensorflow/contrib/graph_editor/transform.py @@ -22,7 +22,6 @@ from __future__ import print_function from copy import deepcopy from functools import partial from six import iteritems -from six import iterkeys from six import string_types from six import StringIO from tensorflow.contrib.graph_editor import reroute @@ -735,9 +734,8 @@ def graph_replace(target_ts, replacement_ts, dst_scope="", # control dependencies. graph = util.get_unique_graph(flatten_target_ts, check_types=(tf_ops.Tensor)) control_ios = util.ControlOutputs(graph) - ops = select.get_walks_intersection_ops(list(iterkeys(replacement_ts)), - flatten_target_ts, - control_ios=control_ios) + ops = select.get_walks_intersection_ops( + list(replacement_ts), flatten_target_ts, control_ios=control_ios) if not ops: raise ValueError("Targets and replacements are not connected!") diff --git a/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op_gpu.cu.cc b/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op_gpu.cu.cc index bbb3a3b18fd7bfdc68e8b8532568985245154794..f97e790b56c511ffb7859b4120b7a4220b75c506 100644 --- a/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op_gpu.cu.cc +++ b/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op_gpu.cu.cc @@ -55,9 +55,10 @@ void AdjustHsvInYiqGPU::operator()(OpKernelContext* ctx, int channel_count, &tranformation_matrix)); // TODO(huangyp): It takes about 3.5 us to compute tranformation_matrix // with one thread. Improve its performance if necessary. - internal::compute_tranformation_matrix_cuda<<<1, 1, 0, cu_stream>>>( - delta_h, scale_s, scale_v, tranformation_matrix.flat().data(), - tranformation_matrix.flat().size()); + TF_CHECK_OK(CudaLaunchKernel(internal::compute_tranformation_matrix_cuda, 1, + 1, 0, cu_stream, delta_h, scale_s, scale_v, + tranformation_matrix.flat().data(), + tranformation_matrix.flat().size())); // Call cuBlas C = A * B directly. auto no_transpose = se::blas::Transpose::kNoTranspose; auto a_ptr = diff --git a/tensorflow/contrib/kafka/ops/kafka_ops.cc b/tensorflow/contrib/kafka/ops/kafka_ops.cc deleted file mode 100644 index 8cdf16103bab2b22d51c144d21a589e1e39f2f0b..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/kafka/ops/kafka_ops.cc +++ /dev/null @@ -1,44 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/core/framework/common_shape_fns.h" -#include "tensorflow/core/framework/op.h" -#include "tensorflow/core/framework/shape_inference.h" - -namespace tensorflow { - -REGISTER_OP("KafkaDataset") - .Input("topics: string") - .Input("servers: string") - .Input("group: string") - .Input("eof: bool") - .Input("timeout: int64") - .Output("handle: variant") - .SetIsStateful() - .SetShapeFn(shape_inference::ScalarShape) - .Doc(R"doc( -Creates a dataset that emits the messages of one or more Kafka topics. - -topics: A `tf.string` tensor containing one or more subscriptions, - in the format of [topic:partition:offset:length], - by default length is -1 for unlimited. -servers: A list of bootstrap servers. -group: The consumer group id. -eof: If True, the kafka reader will stop on EOF. -timeout: The timeout value for the Kafka Consumer to wait - (in millisecond). -)doc"); - -} // namespace tensorflow diff --git a/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.py b/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.py index 08ebcdb544645d3585a1af25c86c6182a1589dcb..3651275f935b50ac9d21bb831fd257eb22a6b793 100644 --- a/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.py +++ b/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function from tensorflow.contrib.kafka.python.ops import kafka_dataset_ops +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import iterator_ops from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors @@ -49,7 +50,8 @@ class KafkaDatasetTest(test.TestCase): topics, group="test", eof=True).repeat(num_epochs) batch_dataset = repeat_dataset.batch(batch_size) - iterator = iterator_ops.Iterator.from_structure(batch_dataset.output_types) + iterator = iterator_ops.Iterator.from_structure( + dataset_ops.get_legacy_output_types(batch_dataset)) init_op = iterator.make_initializer(repeat_dataset) init_batch_op = iterator.make_initializer(batch_dataset) get_next = iterator.get_next() diff --git a/tensorflow/contrib/keras/api/keras/losses/__init__.py b/tensorflow/contrib/keras/api/keras/losses/__init__.py index c4476a7bbd5056fa898468a46031bf3d8b1e44cf..b12832d2e2a3cccb4948d9e3bf3d226030121ac2 100644 --- a/tensorflow/contrib/keras/api/keras/losses/__init__.py +++ b/tensorflow/contrib/keras/api/keras/losses/__init__.py @@ -22,7 +22,7 @@ from __future__ import print_function from tensorflow.python.keras.losses import binary_crossentropy from tensorflow.python.keras.losses import categorical_crossentropy from tensorflow.python.keras.losses import categorical_hinge -from tensorflow.python.keras.losses import cosine_proximity +from tensorflow.python.keras.losses import cosine_similarity from tensorflow.python.keras.losses import hinge from tensorflow.python.keras.losses import kullback_leibler_divergence from tensorflow.python.keras.losses import logcosh diff --git a/tensorflow/contrib/keras/api/keras/metrics/__init__.py b/tensorflow/contrib/keras/api/keras/metrics/__init__.py index 7317fdb52c5b79e787a49d71be49f5261d6b1fff..095b5d798df9ac9038fa1088cdd402dff304e87e 100644 --- a/tensorflow/contrib/keras/api/keras/metrics/__init__.py +++ b/tensorflow/contrib/keras/api/keras/metrics/__init__.py @@ -23,7 +23,7 @@ from tensorflow.python.keras.metrics import binary_accuracy from tensorflow.python.keras.metrics import binary_crossentropy from tensorflow.python.keras.metrics import categorical_accuracy from tensorflow.python.keras.metrics import categorical_crossentropy -from tensorflow.python.keras.metrics import cosine_proximity +from tensorflow.python.keras.metrics import cosine_similarity from tensorflow.python.keras.metrics import hinge from tensorflow.python.keras.metrics import kullback_leibler_divergence from tensorflow.python.keras.metrics import mean_absolute_error diff --git a/tensorflow/contrib/kinesis/python/kernel_tests/kinesis_test.py b/tensorflow/contrib/kinesis/python/kernel_tests/kinesis_test.py index bf89922318b9b9a569e4bd1d71fe6283810cadda..af7018f8368116172511b3f78c42caf3fc215632 100644 --- a/tensorflow/contrib/kinesis/python/kernel_tests/kinesis_test.py +++ b/tensorflow/contrib/kinesis/python/kernel_tests/kinesis_test.py @@ -29,6 +29,7 @@ from __future__ import print_function import boto3 from tensorflow.contrib.kinesis.python.ops import kinesis_dataset_ops +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import iterator_ops from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors @@ -59,7 +60,8 @@ class KinesisDatasetTest(test.TestCase): stream, read_indefinitely=False).repeat(num_epochs) batch_dataset = repeat_dataset.batch(batch_size) - iterator = iterator_ops.Iterator.from_structure(batch_dataset.output_types) + iterator = iterator_ops.Iterator.from_structure( + dataset_ops.get_legacy_output_types(batch_dataset)) init_op = iterator.make_initializer(repeat_dataset) init_batch_op = iterator.make_initializer(batch_dataset) get_next = iterator.get_next() @@ -102,7 +104,8 @@ class KinesisDatasetTest(test.TestCase): stream, shard, read_indefinitely=False).repeat(num_epochs) batch_dataset = repeat_dataset.batch(batch_size) - iterator = iterator_ops.Iterator.from_structure(batch_dataset.output_types) + iterator = iterator_ops.Iterator.from_structure( + dataset_ops.get_legacy_output_types(batch_dataset)) init_op = iterator.make_initializer(repeat_dataset) init_batch_op = iterator.make_initializer(batch_dataset) get_next = iterator.get_next() diff --git a/tensorflow/contrib/labeled_tensor/python/ops/core.py b/tensorflow/contrib/labeled_tensor/python/ops/core.py index 0c6bba758b429a8c4112bc6abb2fae542b5dfc14..8ee554ffa7ab6bbcc2d36c525ad68e03bacb594b 100644 --- a/tensorflow/contrib/labeled_tensor/python/ops/core.py +++ b/tensorflow/contrib/labeled_tensor/python/ops/core.py @@ -321,8 +321,8 @@ class LabeledTensor(object): for (d, axis) in zip(shape, unvalidated_axes.values()): if d != axis.size: raise ValueError( - 'Provided axis size %d does not match tensor dimension size %d' % - (axis.size, d)) + 'Provided axis size %d does not match tensor dimension size %d' + 'in tensor %r' % (axis.size, d, tensor)) self._axes = unvalidated_axes diff --git a/tensorflow/contrib/layers/python/layers/layers.py b/tensorflow/contrib/layers/python/layers/layers.py index 9d9524e4e4b995d795b7c71b5bd083d11c60d5ce..1d959b3c78445977b4fe74ee6c20c86aaf7f86da 100644 --- a/tensorflow/contrib/layers/python/layers/layers.py +++ b/tensorflow/contrib/layers/python/layers/layers.py @@ -2312,7 +2312,9 @@ def layer_norm(inputs, norm_axes = list(range(begin_norm_axis, inputs_rank)) mean, variance = nn.moments(inputs, norm_axes, keep_dims=True) # Compute layer normalization using the batch_normalization function. - variance_epsilon = 1e-12 + # Note that epsilon must be increased for float16 due to the limited + # representable range. + variance_epsilon = 1e-12 if dtype != dtypes.float16 else 1e-3 outputs = nn.batch_normalization( inputs, mean, diff --git a/tensorflow/contrib/layers/python/layers/layers_test.py b/tensorflow/contrib/layers/python/layers/layers_test.py index 1c0088186c030437454c0f764decab9e5a276adc..2cd72410d3d08587b92f38d7dec9546445dc7562 100644 --- a/tensorflow/contrib/layers/python/layers/layers_test.py +++ b/tensorflow/contrib/layers/python/layers/layers_test.py @@ -2869,10 +2869,19 @@ class LayerNormTest(test.TestCase): tol=1e-5, begin_norm_axis=1, dtype=dtypes.float64): + eps = 1e-12 if dtype != dtypes.float16 else 1e-3 expected_mean = np.zeros(input_shape[:begin_norm_axis]) - expected_var = np.ones(input_shape[:begin_norm_axis]) - for mu in [0.0, 1e2]: - for sigma in [1.0, 0.1]: + expected_var_uncorrected = np.ones(input_shape[:begin_norm_axis]) + sigma_list = [1.0, 0.1] + if dtype == dtypes.float16: + # This causes the variance to underflow in float16, and requires that + # variance_epsilon be set appropriately to avoid NaNs in the output. + sigma_list.append(1e-4) + # Note that the mean:variance ratio must be limited to the representable + # range for float16. + for mu in [0.0, 1e2 if dtype != dtypes.float16 else 1e1]: + for sigma in sigma_list: + expected_var = expected_var_uncorrected / (1.0 + eps / sigma**2) input_values = np.random.randn(*input_shape) * sigma + mu with ops.Graph().as_default() as g: with self.session(graph=g) as sess: @@ -2893,10 +2902,13 @@ class LayerNormTest(test.TestCase): outputs, beta, gamma = sess.run((output_t, beta_var, gamma_var)) # Make sure that there are no NaNs self.assertFalse(np.isnan(outputs).any()) + if outputs.dtype != np.float64: + # Cast to float64 before computing mean/variance to avoid + # overflow and precision issues. + outputs = outputs.astype(np.float64) mean = np.mean(outputs, axis=moments_axis) var = np.var(outputs, axis=moments_axis) # Layer-norm implemented in numpy - eps = 1e-12 expected_out = ( (gamma * (input_values - np.mean( input_values, axis=moments_axis, keepdims=True)) / @@ -2933,6 +2945,12 @@ class LayerNormTest(test.TestCase): def testOutputBigInput(self): self.doOutputTest((1, 100, 100, 1)) + def testOutputBigInputFloat32(self): + self.doOutputTest((1, 100, 1000, 1), tol=1e-4, dtype=dtypes.float32) + + def testOutputBigInputFloat16(self): + self.doOutputTest((1, 100, 1000, 1), tol=5e-2, dtype=dtypes.float16) + class GDNTest(test.TestCase): diff --git a/tensorflow/contrib/layers/python/layers/optimizers.py b/tensorflow/contrib/layers/python/layers/optimizers.py index 2fdcd849b026d52ed4aff724838f6c71e3a315d0..3b0758750354e2843577be9142751011c55d587e 100644 --- a/tensorflow/contrib/layers/python/layers/optimizers.py +++ b/tensorflow/contrib/layers/python/layers/optimizers.py @@ -109,11 +109,12 @@ def optimize_loss(loss, gradient_multipliers: dict of variables or variable names to floats. If present, gradients for specified variables will be multiplied by given constant. - clip_gradients: float, callable or `None`. If float, is provided, a global - clipping is applied to prevent the norm of the gradient to exceed this - value. Alternatively, a callable can be provided e.g.: adaptive_clipping. - This callable takes a `list` of `(gradients, variables)` `tuple`s and - returns the same thing with the gradients modified. + clip_gradients: float, callable or `None`. If a float is provided, a global + clipping is applied to prevent the norm of the gradient from exceeding + this value. Alternatively, a callable can be provided, e.g., + `adaptive_clipping_fn()`. This callable takes a list of + `(gradients, variables)` tuples and returns the same thing with the + gradients modified. learning_rate_decay_fn: function, takes `learning_rate` and `global_step` `Tensor`s, returns `Tensor`. Can be used to implement any learning rate decay diff --git a/tensorflow/contrib/learn/python/learn/estimators/estimator_test.py b/tensorflow/contrib/learn/python/learn/estimators/estimator_test.py index 8a461a0bd7ba457fcf830769f23c6ca2860a2732..cbcae338a0a195da2aca1eea2e1b4c7eb8b0e35e 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/estimator_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/estimator_test.py @@ -1181,14 +1181,14 @@ class EstimatorTest(test.TestCase): ] self.assertItemsEqual([expected_vocab_file], assets) graph_ops = [x.name for x in graph.get_operations()] - self.assertTrue('input_example_tensor' in graph_ops) - self.assertTrue('ParseExample/ParseExample' in graph_ops) - self.assertTrue('linear/linear/feature/matmul' in graph_ops) + self.assertIn('input_example_tensor', graph_ops) + self.assertIn('ParseExample/ParseExample', graph_ops) + self.assertIn('linear/linear/feature/matmul', graph_ops) # Since there were no transforms, both save ops are still present. - self.assertTrue('save/SaveV2/tensor_names' in graph_ops) - self.assertTrue('save_1/SaveV2/tensor_names' in graph_ops) + self.assertIn('save/SaveV2/tensor_names', graph_ops) + self.assertIn('save_1/SaveV2/tensor_names', graph_ops) # Since there were no transforms, the hash table lookup is still there. - self.assertTrue('hash_table_Lookup' in graph_ops) + self.assertIn('hash_table_Lookup/LookupTableFindV2', graph_ops) # Restore, to validate that the export was well-formed. # tag_2, tag_3 was subjected to strip_unused_nodes. diff --git a/tensorflow/contrib/learn/python/learn/estimators/head_test.py b/tensorflow/contrib/learn/python/learn/estimators/head_test.py index 7c2d9bb0767cb979dae9c84b5342d129225677ed..a52d25acf402bdda46771e9146a40cfb71e99d53 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/head_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/head_test.py @@ -62,8 +62,8 @@ def _assert_no_variables(test_case): def _assert_metrics(test_case, expected_loss, expected_eval_metrics, model_fn_ops): test_case.assertAlmostEqual(expected_loss, model_fn_ops.loss.eval(), places=4) - for k in six.iterkeys(expected_eval_metrics): - test_case.assertIn(k, six.iterkeys(model_fn_ops.eval_metric_ops)) + for k in expected_eval_metrics: + test_case.assertIn(k, model_fn_ops.eval_metric_ops) variables.initialize_local_variables().run() for key, expected_value in six.iteritems(expected_eval_metrics): value_tensor, update_tensor = model_fn_ops.eval_metric_ops[key] @@ -545,19 +545,19 @@ class MultiLabelHeadTest(test.TestCase): with session.Session(): self.assertListEqual( [1, 0, 0], model_fn_ops.predictions["classes"].eval().tolist()[0]) - self.assertItemsEqual( - ["head_name"], six.iterkeys(model_fn_ops.output_alternatives)) + self.assertItemsEqual(["head_name"], + list(model_fn_ops.output_alternatives)) self.assertEqual( constants.ProblemType.CLASSIFICATION, model_fn_ops.output_alternatives["head_name"][0]) predictions_for_serving = ( model_fn_ops.output_alternatives["head_name"][1]) - self.assertIn("classes", six.iterkeys(predictions_for_serving)) + self.assertIn("classes", predictions_for_serving) self.assertAllEqual( [[b"0", b"1", b"2"], [b"0", b"1", b"2"]], predictions_for_serving["classes"].eval()) - self.assertIn("probabilities", six.iterkeys(predictions_for_serving)) + self.assertIn("probabilities", predictions_for_serving) self.assertAllClose( [[0.731059, 0.5, 0.5], [0.5, 0.5, 0.731059,]], @@ -850,18 +850,18 @@ class BinaryClassificationHeadTest(test.TestCase): with session.Session(): self.assertListEqual( [1, 1], list(model_fn_ops.predictions["classes"].eval())) - self.assertItemsEqual( - ["head_name"], six.iterkeys(model_fn_ops.output_alternatives)) + self.assertItemsEqual(["head_name"], + list(model_fn_ops.output_alternatives)) self.assertEqual( constants.ProblemType.LOGISTIC_REGRESSION, model_fn_ops.output_alternatives["head_name"][0]) predictions_for_serving = ( model_fn_ops.output_alternatives["head_name"][1]) - self.assertIn("classes", six.iterkeys(predictions_for_serving)) + self.assertIn("classes", predictions_for_serving) predicted_classes = predictions_for_serving["classes"].eval().tolist() self.assertListEqual( [b"0", b"1"], predicted_classes[0]) - self.assertIn("probabilities", six.iterkeys(predictions_for_serving)) + self.assertIn("probabilities", predictions_for_serving) def testBinaryClassificationInferMode_withWeightColumn(self): n_classes = 2 @@ -1349,18 +1349,18 @@ class MultiClassHeadTest(test.TestCase): self.assertAllEqual( [0, 2], model_fn_ops.predictions["classes"].eval()) - self.assertItemsEqual( - ["head_name"], six.iterkeys(model_fn_ops.output_alternatives)) + self.assertItemsEqual(["head_name"], + list(model_fn_ops.output_alternatives)) self.assertEqual( constants.ProblemType.CLASSIFICATION, model_fn_ops.output_alternatives["head_name"][0]) predictions_for_serving = ( model_fn_ops.output_alternatives["head_name"][1]) - self.assertIn("classes", six.iterkeys(predictions_for_serving)) + self.assertIn("classes", predictions_for_serving) self.assertAllEqual( [[b"0", b"1", b"2"], [b"0", b"1", b"2"]], predictions_for_serving["classes"].eval()) - self.assertIn("probabilities", six.iterkeys(predictions_for_serving)) + self.assertIn("probabilities", predictions_for_serving) self.assertAllClose( [[0.576117, 0.2119416, 0.2119416], [0.2119416, 0.2119416, 0.576117]], @@ -1401,18 +1401,18 @@ class MultiClassHeadTest(test.TestCase): self.assertAllEqual( [b"key0", b"key2"], model_fn_ops.predictions["classes"].eval()) - self.assertItemsEqual( - ["head_name"], six.iterkeys(model_fn_ops.output_alternatives)) + self.assertItemsEqual(["head_name"], + list(model_fn_ops.output_alternatives)) self.assertEqual( constants.ProblemType.CLASSIFICATION, model_fn_ops.output_alternatives["head_name"][0]) predictions_for_serving = ( model_fn_ops.output_alternatives["head_name"][1]) - self.assertIn("classes", six.iterkeys(predictions_for_serving)) + self.assertIn("classes", predictions_for_serving) self.assertAllEqual( [[b"key0", b"key1", b"key2"], [b"key0", b"key1", b"key2"]], predictions_for_serving["classes"].eval()) - self.assertIn("probabilities", six.iterkeys(predictions_for_serving)) + self.assertIn("probabilities", predictions_for_serving) self.assertAllClose( [[0.576117, 0.2119416, 0.2119416], [0.2119416, 0.2119416, 0.576117]], diff --git a/tensorflow/contrib/learn/python/learn/estimators/model_fn.py b/tensorflow/contrib/learn/python/learn/estimators/model_fn.py index dcb161180c99ce71195c820217e8bdaf79d70901..96adc8b83b5bec912460dbb54899ce5f168b8f25 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/model_fn.py +++ b/tensorflow/contrib/learn/python/learn/estimators/model_fn.py @@ -219,7 +219,7 @@ class ModelFnOps( used if a Servo request does not explicitly mention which head to infer on. Pass the key of the output alternative here that you want to designate as default. A separate ExportOutpout for this default head - wil be added to the export_outputs dict with the special key + will be added to the export_outputs dict with the special key signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, unless there is already an enry in output_alternatives with this special key. diff --git a/tensorflow/contrib/lite/python/BUILD b/tensorflow/contrib/lite/python/BUILD deleted file mode 100644 index 893ddd78231c8a0d819cbe5776e6873bdab57355..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/python/BUILD +++ /dev/null @@ -1,12 +0,0 @@ -licenses(["notice"]) - -# DO NOT USE THIS TARGET. TensorFlow Lite has moved to tensorflow/lite. -py_library( - name = "lite", - srcs = ["__init__.py"], - srcs_version = "PY2AND3", - visibility = ["//visibility:public"], - deps = [ - "//tensorflow/lite/python:lite", - ], -) diff --git a/tensorflow/contrib/lookup/lookup_ops.py b/tensorflow/contrib/lookup/lookup_ops.py index 3d21fb68a1452c97f7eb85491fc850d9e846266a..20e86e56bbe911eca2bba661aff7165e53fa159e 100644 --- a/tensorflow/contrib/lookup/lookup_ops.py +++ b/tensorflow/contrib/lookup/lookup_ops.py @@ -30,6 +30,7 @@ from tensorflow.python.ops.lookup_ops import IdTableWithHashBuckets from tensorflow.python.ops.lookup_ops import index_table_from_file from tensorflow.python.ops.lookup_ops import index_to_string_table_from_file from tensorflow.python.ops.lookup_ops import InitializableLookupTableBase +from tensorflow.python.ops.lookup_ops import InitializableLookupTableBaseV1 from tensorflow.python.ops.lookup_ops import KeyValueTensorInitializer from tensorflow.python.ops.lookup_ops import LookupInterface from tensorflow.python.ops.lookup_ops import StrongHashSpec @@ -284,7 +285,7 @@ def index_to_string(tensor, mapping, default_value="UNK", name=None): return table.lookup(tensor) -class HashTable(InitializableLookupTableBase): +class HashTable(InitializableLookupTableBaseV1): """A generic hash table implementation. Example usage: @@ -325,7 +326,7 @@ class HashTable(InitializableLookupTableBase): super(HashTable, self).__init__(default_value, initializer) self._value_shape = self._default_value.get_shape() - def create_resource(self): + def _create_resource(self): table_ref = gen_lookup_ops.hash_table_v2( shared_name=self._shared_name, key_dtype=self._initializer.key_dtype, @@ -337,6 +338,10 @@ class HashTable(InitializableLookupTableBase): self._table_name = table_ref.op.name.split("/")[-1] return table_ref + @property + def init(self): + return self.initializer + @property def name(self): return self._table_name @@ -362,4 +367,4 @@ class HashTable(InitializableLookupTableBase): MutableHashTable = lookup_ops.MutableHashTable -MutableDenseHashTable = lookup_ops.MutableDenseHashTable +MutableDenseHashTable = lookup_ops.DenseHashTable diff --git a/tensorflow/contrib/makefile/Makefile b/tensorflow/contrib/makefile/Makefile index 7ea6e34cf50ed8e292f11314550d992c3dde34c0..d22548d5007997491370b002f84ae93eb041f75a 100644 --- a/tensorflow/contrib/makefile/Makefile +++ b/tensorflow/contrib/makefile/Makefile @@ -439,7 +439,6 @@ $(MARCH_OPTION) \ -Itensorflow/core/kernels \ -I$(MAKEFILE_DIR)/downloads/cub \ -I$(MAKEFILE_DIR)/downloads/cub/cub_archive/cub/device \ --Ithird_party/toolchains/gpus/cuda \ -I$(JETPACK)/cuda/include \ -I$(JETPACK) \ -I$(JETPACK)/cuDNN/aarch64 \ diff --git a/tensorflow/contrib/makefile/proto_text_pb_cc_files.txt b/tensorflow/contrib/makefile/proto_text_pb_cc_files.txt index 8330c45cc16ffa536107e25699379bb5d9e8993b..1c1460ce77c99d29785c7e8b8a8e9f770a45b59f 100644 --- a/tensorflow/contrib/makefile/proto_text_pb_cc_files.txt +++ b/tensorflow/contrib/makefile/proto_text_pb_cc_files.txt @@ -25,6 +25,7 @@ tensorflow/core/framework/variable.pb.cc tensorflow/core/framework/versions.pb.cc tensorflow/core/grappler/costs/op_performance_data.pb.cc tensorflow/core/lib/core/error_codes.pb.cc +tensorflow/core/protobuf/trackable_object_graph.pb.cc tensorflow/core/protobuf/cluster.pb.cc tensorflow/core/protobuf/config.pb.cc tensorflow/core/protobuf/eager_service.pb.cc @@ -34,7 +35,9 @@ tensorflow/core/protobuf/meta_graph.pb.cc tensorflow/core/protobuf/named_tensor.pb.cc tensorflow/core/protobuf/queue_runner.pb.cc tensorflow/core/protobuf/rewriter_config.pb.cc +tensorflow/core/protobuf/saved_object_graph.pb.cc tensorflow/core/protobuf/saver.pb.cc +tensorflow/core/protobuf/struct.pb.cc tensorflow/core/protobuf/tensorflow_server.pb.cc tensorflow/core/protobuf/verifier_config.pb.cc tensorflow/core/util/event.pb.cc diff --git a/tensorflow/contrib/makefile/proto_text_pb_h_files.txt b/tensorflow/contrib/makefile/proto_text_pb_h_files.txt index 7257ac8feedfb8ed18c4d691cd85766e70a48ae8..5def632e8a7b65272a1339bdacd92c1fa23012d2 100644 --- a/tensorflow/contrib/makefile/proto_text_pb_h_files.txt +++ b/tensorflow/contrib/makefile/proto_text_pb_h_files.txt @@ -25,6 +25,7 @@ tensorflow/core/framework/variable.pb.h tensorflow/core/framework/versions.pb.h tensorflow/core/grappler/costs/op_performance_data.pb.h tensorflow/core/lib/core/error_codes.pb.h +tensorflow/core/protobuf/trackable_object_graph.pb.h tensorflow/core/protobuf/cluster.pb.h tensorflow/core/protobuf/config.pb.h tensorflow/core/protobuf/debug.pb.h @@ -34,7 +35,9 @@ tensorflow/core/protobuf/meta_graph.pb.h tensorflow/core/protobuf/named_tensor.pb.h tensorflow/core/protobuf/queue_runner.pb.h tensorflow/core/protobuf/rewriter_config.pb.h +tensorflow/core/protobuf/saved_object_graph.pb.h tensorflow/core/protobuf/saver.pb.h +tensorflow/core/protobuf/struct.pb.h tensorflow/core/protobuf/tensor_bundle.pb.h tensorflow/core/protobuf/tensorflow_server.pb.h tensorflow/core/protobuf/verifier_config.pb.h diff --git a/tensorflow/contrib/makefile/tf_op_files.txt b/tensorflow/contrib/makefile/tf_op_files.txt index 2cd7d6d519a55423a96526b541845392d9ec6bc2..0ed87544ce3c43b25e631757dbd66412d9cd42c6 100644 --- a/tensorflow/contrib/makefile/tf_op_files.txt +++ b/tensorflow/contrib/makefile/tf_op_files.txt @@ -43,7 +43,9 @@ tensorflow/core/kernels/conv_grad_input_ops.cc tensorflow/core/kernels/conv_grad_ops.cc tensorflow/core/kernels/conv_ops.cc tensorflow/core/kernels/conv_ops_3d.cc -tensorflow/core/kernels/conv_ops_fused.cc +tensorflow/core/kernels/conv_ops_fused_double.cc +tensorflow/core/kernels/conv_ops_fused_float.cc +tensorflow/core/kernels/conv_ops_fused_half.cc tensorflow/core/kernels/conv_ops_using_gemm.cc tensorflow/core/kernels/crop_and_resize_op.cc tensorflow/core/kernels/ctc_decoder_ops.cc diff --git a/tensorflow/contrib/makefile/tf_proto_files.txt b/tensorflow/contrib/makefile/tf_proto_files.txt index 24d86d313b76343ed9450a33cf185d9c426696bb..deb6a5b94020a02b878bdd68a33b3737a97fcf2b 100644 --- a/tensorflow/contrib/makefile/tf_proto_files.txt +++ b/tensorflow/contrib/makefile/tf_proto_files.txt @@ -31,6 +31,7 @@ tensorflow/core/framework/versions.proto tensorflow/core/grappler/costs/op_performance_data.proto tensorflow/core/kernels/boosted_trees/boosted_trees.proto tensorflow/core/lib/core/error_codes.proto +tensorflow/core/protobuf/trackable_object_graph.proto tensorflow/core/protobuf/cluster.proto tensorflow/core/protobuf/config.proto tensorflow/core/protobuf/debug.proto @@ -40,7 +41,9 @@ tensorflow/core/protobuf/meta_graph.proto tensorflow/core/protobuf/named_tensor.proto tensorflow/core/protobuf/queue_runner.proto tensorflow/core/protobuf/rewriter_config.proto +tensorflow/core/protobuf/saved_object_graph.proto tensorflow/core/protobuf/saver.proto +tensorflow/core/protobuf/struct.proto tensorflow/core/protobuf/tensor_bundle.proto tensorflow/core/protobuf/tensorflow_server.proto tensorflow/core/protobuf/verifier_config.proto diff --git a/tensorflow/contrib/mpi_collectives/kernels/ring.cu.cc b/tensorflow/contrib/mpi_collectives/kernels/ring.cu.cc index b04abde4694199d827a1738850bded9bf696d56c..ca3ddfa721d45a2de3ea51c80d6adfa2371c3c94 100644 --- a/tensorflow/contrib/mpi_collectives/kernels/ring.cu.cc +++ b/tensorflow/contrib/mpi_collectives/kernels/ring.cu.cc @@ -96,13 +96,14 @@ __global__ void elemwise_accum(T* out, const T* in, const size_t N) { // Synchronously accumulate tensors on the GPU, using a different stream than // the default and than TensorFlow to avoid synchronizing on operations // unrelated to the allreduce. -#define GENERATE_ACCUMULATE(type) \ - template <> \ - void AccumulateTensorData(type * dst, type * src, \ - size_t size) { \ - auto stream = CudaStreamForMPI(); \ - elemwise_accum<<<32, 256, 0, stream>>>(dst, src, size); \ - cudaStreamSynchronize(stream); \ +#define GENERATE_ACCUMULATE(type) \ + template <> \ + void AccumulateTensorData(type * dst, type * src, \ + size_t size) { \ + auto stream = CudaStreamForMPI(); \ + TF_CHECK_OK(CudaLaunchKernel(elemwise_accum, 32, 256, 0, stream, \ + dst, src, size)); \ + cudaStreamSynchronize(stream); \ }; GENERATE_ACCUMULATE(int); GENERATE_ACCUMULATE(long long); diff --git a/tensorflow/contrib/mpi_collectives/ring.cu.cc b/tensorflow/contrib/mpi_collectives/ring.cu.cc index 2f3eef366a9a3c10e59cd5298fc1626e1094dff8..c73156d230820e8f89d88d8d4c8599fd1a5f68d8 100644 --- a/tensorflow/contrib/mpi_collectives/ring.cu.cc +++ b/tensorflow/contrib/mpi_collectives/ring.cu.cc @@ -96,13 +96,14 @@ __global__ void elemwise_accum(T* out, const T* in, const size_t N) { // Synchronously accumulate tensors on the GPU, using a different stream than // the default and than TensorFlow to avoid synchronizing on operations // unrelated to the allreduce. -#define GENERATE_ACCUMULATE(type) \ - template <> \ - void AccumulateTensorData(type * dst, type * src, \ - size_t size) { \ - auto stream = CudaStreamForMPI(); \ - elemwise_accum<<<32, 256, 0, stream>>>(dst, src, size); \ - cudaStreamSynchronize(stream); \ +#define GENERATE_ACCUMULATE(type) \ + template <> \ + void AccumulateTensorData(type * dst, type * src, \ + size_t size) { \ + auto stream = CudaStreamForMPI(); \ + TF_CHECK_OK(CudaLaunchKernel(elemwise_accum, 32, 256, 0, stream, \ + dst, src, size)); \ + cudaStreamSynchronize(stream); \ }; GENERATE_ACCUMULATE(int); GENERATE_ACCUMULATE(long long); diff --git a/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py b/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py index b2ea3daf82ed8daa6e0b9acd8e3cf258b8181615..1e210ec666c072b8fc132966c3936729737ba869 100644 --- a/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py +++ b/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py @@ -704,8 +704,7 @@ class CheckpointCompatibilityTests(test.TestCase): if context.executing_eagerly(): self._check_sentinels(root) if context.executing_eagerly(): - with self.assertRaisesRegexp(AssertionError, "OBJECT_CONFIG_JSON"): - status.assert_consumed() + status.assert_consumed() else: # When graph building, we haven't read any keys, so we don't know # whether the restore will be complete. diff --git a/tensorflow/contrib/optimizer_v2/optimizer_v2.py b/tensorflow/contrib/optimizer_v2/optimizer_v2.py index a7f978634ed45012144b2cc49ed069f6fca44f66..436ece79a79810d4688e259523a4f86a1ca7f5a5 100644 --- a/tensorflow/contrib/optimizer_v2/optimizer_v2.py +++ b/tensorflow/contrib/optimizer_v2/optimizer_v2.py @@ -965,7 +965,7 @@ class OptimizerV2(optimizer_v1.Optimizer): # `update_op`. # TODO(josh11b): Make different state objects for each device to # avoid needing to set the device_policy. - device_policy = context.context().device_policy( + device_policy = context.device_policy( context.DEVICE_PLACEMENT_SILENT) with ops.name_scope("update_" + scope_name), device_policy: return processor.update_op(self, g, state) @@ -981,7 +981,7 @@ class OptimizerV2(optimizer_v1.Optimizer): def finish(): # TODO(josh11b): Make different state objects for each device to # avoid needing to set the device_policy. - with context.context().device_policy(context.DEVICE_PLACEMENT_SILENT): + with context.device_policy(context.DEVICE_PLACEMENT_SILENT): return self._finish(state) update_ops = control_flow_ops.group(update_ops) diff --git a/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops_gpu.cu.cc b/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops_gpu.cu.cc index 204b83f7f5f118f418815edb6c482b1c06673845..13fbd974e9ce6a680a31507f7f49df17d121535f 100644 --- a/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops_gpu.cu.cc +++ b/tensorflow/contrib/reduce_slice_ops/kernels/reduce_slice_ops_gpu.cu.cc @@ -77,10 +77,10 @@ namespace functor { sizex, sizey, sizez, d, ReduceSliceDeviceKernel##reduceop, \ 0, 0); \ \ - ReduceSliceDeviceKernel##reduceop \ - <<>>( \ - config, indices_width, bound, beginning(), indices.data(), \ - data.data(), output.data()); \ + TF_CHECK_OK(CudaLaunchKernel( \ + ReduceSliceDeviceKernel##reduceop, config.block_count, \ + config.thread_per_block, 0, d.stream(), config, indices_width, \ + bound, beginning(), indices.data(), data.data(), output.data())); \ } \ }; diff --git a/tensorflow/contrib/resampler/kernels/resampler_ops_gpu.cu.cc b/tensorflow/contrib/resampler/kernels/resampler_ops_gpu.cu.cc index 3c07051f685c74b6e45fb782c80871f38dffbbf4..3b2ee098b3e24287298273a04f80e41f6d9dcd86 100644 --- a/tensorflow/contrib/resampler/kernels/resampler_ops_gpu.cu.cc +++ b/tensorflow/contrib/resampler/kernels/resampler_ops_gpu.cu.cc @@ -119,10 +119,10 @@ struct Resampler2DFunctor { batch_size * num_sampling_points * data_channels; ::tensorflow::CudaLaunchConfig config = ::tensorflow::GetCudaLaunchConfig(output_data_size, d); - Resampler2DKernel - <<>>( - data, warp, output, batch_size, data_height, data_width, - data_channels, num_sampling_points); + TF_CHECK_OK(CudaLaunchKernel( + Resampler2DKernel, config.block_count, config.thread_per_block, 0, + d.stream(), data, warp, output, batch_size, data_height, data_width, + data_channels, num_sampling_points)); } }; @@ -254,22 +254,23 @@ struct ResamplerGrad2DFunctor { ::tensorflow::CudaLaunchConfig config = ::tensorflow::GetCudaLaunchConfig(grad_warp_size, d); - ::tensorflow:: - SetZero<<>>( - grad_warp_size, grad_warp); + TF_CHECK_OK(::tensorflow::CudaLaunchKernel( + SetZero, config.block_count, config.thread_per_block, 0, d.stream(), + grad_warp_size, grad_warp)); config = ::tensorflow::GetCudaLaunchConfig(grad_data_size, d); - ::tensorflow:: - SetZero<<>>( - grad_data_size, grad_data); + TF_CHECK_OK(::tensorflow::CudaLaunchKernel( + SetZero, config.block_count, config.thread_per_block, 0, d.stream(), + grad_data_size, grad_data)); const int resampler_output_size = batch_size * num_sampling_points * data_channels; config = ::tensorflow::GetCudaLaunchConfig(resampler_output_size, d); - ResamplerGrad2DKernel - <<>>( - data, warp, grad_output, grad_data, grad_warp, batch_size, - data_height, data_width, data_channels, num_sampling_points); + TF_CHECK_OK(CudaLaunchKernel(ResamplerGrad2DKernel, config.block_count, + config.thread_per_block, 0, d.stream(), data, + warp, grad_output, grad_data, grad_warp, + batch_size, data_height, data_width, + data_channels, num_sampling_points)); } }; diff --git a/tensorflow/contrib/rnn/kernels/lstm_ops_gpu.cu.cc b/tensorflow/contrib/rnn/kernels/lstm_ops_gpu.cu.cc index 15ae95f13cffa5d1469d737b23f2a83b9e5a694f..81beb2942c183e6a831b64e946fea89c050b88db 100644 --- a/tensorflow/contrib/rnn/kernels/lstm_ops_gpu.cu.cc +++ b/tensorflow/contrib/rnn/kernels/lstm_ops_gpu.cu.cc @@ -242,8 +242,9 @@ void LSTMBlockCellFpropWithCUDA( const int block_dim = 128; const int grid_dim = Eigen::divup(batch_size * (cell_size + input_size), block_dim); - concat_xh<<>>( - xh.data(), x.data(), h_prev.data(), batch_size, cell_size, input_size); + TF_CHECK_OK(CudaLaunchKernel(concat_xh, grid_dim, block_dim, 0, cu_stream, + xh.data(), x.data(), h_prev.data(), batch_size, + cell_size, input_size)); // states1 = xh * w typename TTypes::ConstMatrix const_xh(xh.data(), xh.dimensions()); @@ -261,15 +262,17 @@ void LSTMBlockCellFpropWithCUDA( Eigen::divup(cell_size, static_cast(block_dim_2d.y))); if (use_peephole) { - lstm_gates<<>>( + TF_CHECK_OK(CudaLaunchKernel( + lstm_gates, grid_dim_2d, block_dim_2d, 0, cu_stream, icfo.data(), b.data(), cs_prev.data(), wci.data(), wcf.data(), wco.data(), o.data(), h.data(), ci.data(), cs.data(), co.data(), - i.data(), f.data(), forget_bias, cell_clip, batch_size, cell_size); + i.data(), f.data(), forget_bias, cell_clip, batch_size, cell_size)); } else { - lstm_gates<<>>( + TF_CHECK_OK(CudaLaunchKernel( + lstm_gates, grid_dim_2d, block_dim_2d, 0, cu_stream, icfo.data(), b.data(), cs_prev.data(), wci.data(), wcf.data(), wco.data(), o.data(), h.data(), ci.data(), cs.data(), co.data(), - i.data(), f.data(), forget_bias, cell_clip, batch_size, cell_size); + i.data(), f.data(), forget_bias, cell_clip, batch_size, cell_size)); } } @@ -374,12 +377,13 @@ void LSTMBlockCellBpropWithCUDA( dim3 grid_dim_2d(Eigen::divup(batch_size, static_cast(block_dim_2d.x)), Eigen::divup(cell_size, static_cast(block_dim_2d.y))); - lstm_gates_bprop<<>>( + TF_CHECK_OK(CudaLaunchKernel( + lstm_gates_bprop, grid_dim_2d, block_dim_2d, 0, cu_stream, cs_prev.data(), h_prev.data(), w.data(), wci.data(), wcf.data(), wco.data(), b.data(), i.data(), cs.data(), f.data(), o.data(), ci.data(), co.data(), cs_grad.data(), h_grad.data(), do_.data(), dcs.data(), dci.data(), df.data(), di.data(), dicfo.data(), cs_prev_grad.data(), - batch_size, cell_size, use_peephole); + batch_size, cell_size, use_peephole)); if (use_peephole) { Eigen::array p_shape({1, cell_size}); diff --git a/tensorflow/contrib/rnn/python/ops/rnn_cell.py b/tensorflow/contrib/rnn/python/ops/rnn_cell.py index 482e547a16be85804beec88a91fa03b053d09b27..9111044e5b3bf93feffb19a6fdf0d634d2499083 100644 --- a/tensorflow/contrib/rnn/python/ops/rnn_cell.py +++ b/tensorflow/contrib/rnn/python/ops/rnn_cell.py @@ -3153,7 +3153,7 @@ class IndyGRUCell(rnn_cell_impl.LayerRNNCell): r"""Independently Gated Recurrent Unit cell. Based on IndRNNs (https://arxiv.org/abs/1803.04831) and similar to GRUCell, - yet with the \(U_r\), \(U_z\), and \(U\) matrices in equations 5, 6, and + yet with the \\(U_r\\), \\(U_z\\), and \\(U\\) matrices in equations 5, 6, and 8 of http://arxiv.org/abs/1406.1078 respectively replaced by diagonal matrices, i.e. a Hadamard product with a single vector: @@ -3164,12 +3164,10 @@ class IndyGRUCell(rnn_cell_impl.LayerRNNCell): $$\tilde{h}^{(t)}_j = \phi\left([\mathbf W \mathbf x]_j + [\mathbf u \circ \mathbf r \circ \mathbf h_{(t-1)}]_j\right)$$ - where \(\circ\) denotes the Hadamard operator. This means that each IndyGRU + where \\(\circ\\) denotes the Hadamard operator. This means that each IndyGRU node sees only its own state, as opposed to seeing all states in the same layer. - TODO(gonnet): Write a paper describing this and add a reference here. - Args: num_units: int, The number of units in the GRU cell. activation: Nonlinearity to use. Default: `tanh`. @@ -3254,7 +3252,7 @@ class IndyGRUCell(rnn_cell_impl.LayerRNNCell): self.built = True def call(self, inputs, state): - """Gated recurrent unit (GRU) with nunits cells.""" + """Recurrently independent Gated Recurrent Unit (GRU) with nunits cells.""" gate_inputs = math_ops.matmul(inputs, self._gate_kernel_w) + ( gen_array_ops.tile(state, [1, 2]) * self._gate_kernel_u) @@ -3278,10 +3276,9 @@ class IndyLSTMCell(rnn_cell_impl.LayerRNNCell): r"""Basic IndyLSTM recurrent network cell. Based on IndRNNs (https://arxiv.org/abs/1803.04831) and similar to - BasicLSTMCell, yet with the \(U_f\), \(U_i\), \(U_o\) and \(U_c\) - matrices in - https://en.wikipedia.org/wiki/Long_short-term_memory#LSTM_with_a_forget_gate - replaced by diagonal matrices, i.e. a Hadamard product with a single vector: + BasicLSTMCell, yet with the \\(U_f\\), \\(U_i\\), \\(U_o\\) and \\(U_c\\) + matrices in the regular LSTM equations replaced by diagonal matrices, i.e. a + Hadamard product with a single vector: $$f_t = \sigma_g\left(W_f x_t + u_f \circ h_{t-1} + b_f\right)$$ $$i_t = \sigma_g\left(W_i x_t + u_i \circ h_{t-1} + b_i\right)$$ @@ -3289,8 +3286,8 @@ class IndyLSTMCell(rnn_cell_impl.LayerRNNCell): $$c_t = f_t \circ c_{t-1} + i_t \circ \sigma_c\left(W_c x_t + u_c \circ h_{t-1} + b_c\right)$$ - where \(\circ\) denotes the Hadamard operator. This means that each IndyLSTM - node sees only its own state \(h\) and \(c\), as opposed to seeing all + where \\(\circ\\) denotes the Hadamard operator. This means that each IndyLSTM + node sees only its own state \\(h\\) and \\(c\\), as opposed to seeing all states in the same layer. We add forget_bias (default: 1) to the biases of the forget gate in order to @@ -3298,11 +3295,6 @@ class IndyLSTMCell(rnn_cell_impl.LayerRNNCell): It does not allow cell clipping, a projection layer, and does not use peep-hole connections: it is the basic baseline. - - For advanced models, please use the full `tf.nn.rnn_cell.LSTMCell` - that follows. - - TODO(gonnet): Write a paper describing this and add a reference here. """ def __init__(self, @@ -3429,7 +3421,7 @@ class MinimalRNNCell(rnn_cell_impl.LayerRNNCell): Propagation in Recurrent Neural Networks." ICML, 2018. A MinimalRNN cell first projects the input to the hidden space. The new - hidden state is then calcuated as a weighted sum of the projected input and + hidden state is then calculated as a weighted sum of the projected input and the previous hidden state, using a single update gate. """ @@ -3543,7 +3535,7 @@ class CFNCell(rnn_cell_impl.LayerRNNCell): "A recurrent neural network without chaos." ICLR, 2017. A CFN cell first projects the input to the hidden space. The hidden state - goes through a contractive mapping. The new hidden state is then calcuated + goes through a contractive mapping. The new hidden state is then calculated as a linear combination of the projected input and the contracted previous hidden state, using decoupled input and forget gates. """ diff --git a/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model.py b/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model.py index 0392ed9eee79391c60318faf68d8dfd6eb64a994..a61e9579b84a60d74b73e45a6100a2c772d9cff8 100644 --- a/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model.py +++ b/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model.py @@ -22,5 +22,5 @@ from tensorflow.python.keras import saving # TODO(kathywu): Remove all contrib callers, switch to tf.keras. -save_keras_model = saving.export +save_keras_model = saving.export_saved_model load_keras_model = saving.load_from_saved_model diff --git a/tensorflow/contrib/seq2seq/BUILD b/tensorflow/contrib/seq2seq/BUILD index 2a70b08f5c46e11e7fd83fe134741b9a241153f5..8e2ce82294287dda07d2067c5b9f012f510dbd08 100644 --- a/tensorflow/contrib/seq2seq/BUILD +++ b/tensorflow/contrib/seq2seq/BUILD @@ -269,4 +269,5 @@ cuda_py_test( "//tensorflow/python:platform_test", "//tensorflow/python:variables", ], + shard_count = 4, ) diff --git a/tensorflow/contrib/seq2seq/kernels/beam_search_ops_gpu.cu.cc b/tensorflow/contrib/seq2seq/kernels/beam_search_ops_gpu.cu.cc index bc28d492fe1a25afe0d0783539aa9e759e7b703f..be2aa4782c3cbc2ecce23b57d332e9bf0cec18bc 100644 --- a/tensorflow/contrib/seq2seq/kernels/beam_search_ops_gpu.cu.cc +++ b/tensorflow/contrib/seq2seq/kernels/beam_search_ops_gpu.cu.cc @@ -91,16 +91,11 @@ struct GatherTree { beams.device(d) = beams.constant(end_token); CudaLaunchConfig config = GetCudaLaunchConfig(batch_size * beam_width, d); - // clang-format off - GatherTreeOpKernel - <<>>( - batch_size, max_time, beam_width, - step_ids.data(), - parent_ids.data(), - max_sequence_length.data(), - end_token, - beams.data()); - // clang-format on + TF_CHECK_OK(CudaLaunchKernel( + GatherTreeOpKernel, config.block_count, config.thread_per_block, 0, + d.stream(), batch_size, max_time, beam_width, step_ids.data(), + parent_ids.data(), max_sequence_length.data(), end_token, + beams.data())); } }; diff --git a/tensorflow/contrib/slim/README.md b/tensorflow/contrib/slim/README.md index f2bb458848fab5603128903868b52f29785efc92..7b54aafeb2cfb5f2a99a93b97d14fbc5bf6e8f9c 100644 --- a/tensorflow/contrib/slim/README.md +++ b/tensorflow/contrib/slim/README.md @@ -11,7 +11,7 @@ import tensorflow.contrib.slim as slim ## Why TF-Slim? -TF-Slim is a library that makes building, training and evaluation neural +TF-Slim is a library that makes defining, training and evaluating neural networks simple: * Allows the user to define models much more compactly by eliminating @@ -78,7 +78,7 @@ provides convenience wrappers for variable creation and manipulation. ## Defining Models Models can be succinctly defined using TF-Slim by combining its variables, -layers and scopes. Each of these elements are defined below. +layers and scopes. Each of these elements is defined below. ### Variables @@ -160,15 +160,15 @@ slim.add_model_variable(my_model_variable) ### Layers -While the set of TensorFlow operations is quite extensive, developers of -neural networks typically think of models in terms of higher level concepts -like "layers", "losses", "metrics", and "networks". A layer, -such as a Convolutional Layer, a Fully Connected Layer or a BatchNorm Layer -are more abstract than a single TensorFlow operation and typically involve -several operations. Furthermore, a layer usually (but not always) has -variables (tunable parameters) associated with it, unlike more primitive -operations. For example, a Convolutional Layer in a neural network -is composed of several low level operations: +While the set of TensorFlow operations is quite extensive, developers of neural +networks typically think of models in terms of higher level concepts like +"layers", "losses", "metrics", and "networks". A layer, such as a Convolutional +Layer, a Fully Connected Layer or a BatchNorm Layer is more abstract than a +single TensorFlow operation and typically involve several operations. +Furthermore, a layer usually (but not always) has variables (tunable parameters) +associated with it, unlike more primitive operations. For example, a +Convolutional Layer in a neural network is composed of several low level +operations: 1. Creating the weight and bias variables 2. Convolving the weights with the input from the previous layer @@ -455,9 +455,8 @@ loss = slim.losses.softmax_cross_entropy(predictions, labels) ``` In this example, we start by creating the model (using TF-Slim's VGG -implementation), and add the standard classification loss. Now, lets turn -to the case where we have a multi-task model that produces multiple outputs: - +implementation), and add the standard classification loss. Now, let's turn to +the case where we have a multi-task model that produces multiple outputs: ```python # Load the images and labels. @@ -555,8 +554,8 @@ that we'll save a model checkpoint every 10 minutes. ### Working Example: Training the VGG16 Model -To illustrate this, lets -examine the following sample of training the VGG network: +To illustrate this, let's examine the following sample of training the VGG +network: ```python import tensorflow as tf @@ -738,7 +737,7 @@ slim.learning.train(train_op, log_dir, init_fn=init_fn) Once we've trained a model (or even while the model is busy training) we'd like to see how well the model performs in practice. This is accomplished by picking -a set of evaluation metrics, which will grade the models performance, and the +a set of evaluation metrics, which will grade the model's performance, and the evaluation code which actually loads the data, performs inference, compares the results to the ground truth and records the evaluation scores. This step may be performed once or repeated periodically. diff --git a/tensorflow/contrib/tensor_forest/python/ops/model_ops.py b/tensorflow/contrib/tensor_forest/python/ops/model_ops.py index 40bf7081a3f22dfd68fd46f0f61695ee9ca7863b..d36d0eb0c46b0d68bea4b6fc29a20dc8876ac539 100644 --- a/tensorflow/contrib/tensor_forest/python/ops/model_ops.py +++ b/tensorflow/contrib/tensor_forest/python/ops/model_ops.py @@ -103,9 +103,9 @@ class TreeVariable(tracking.TrackableResource): self._container = container self._init_op = None super(TreeVariable, self).__init__() - self._resource_handle = self.create_resource() + self._resource_handle = self._create_resource() - def create_resource(self): + def _create_resource(self): if context.executing_eagerly(): # TODO(allenl): This will leak memory due to kernel caching by the # shared_name attribute value (but is better than the alternative of @@ -117,7 +117,7 @@ class TreeVariable(tracking.TrackableResource): return gen_model_ops.decision_tree_resource_handle_op( self._container, shared_name=shared_name, name=self._name) - def initialize(self): + def _initialize(self): return gen_model_ops.create_tree_variable( self.resource_handle, self._tree_config, @@ -126,7 +126,7 @@ class TreeVariable(tracking.TrackableResource): @property def initializer(self): if self._init_op is None: - self._init_op = self.initialize() + self._init_op = self._initialize() return self._init_op def is_initialized(self): diff --git a/tensorflow/contrib/tensor_forest/python/ops/stats_ops.py b/tensorflow/contrib/tensor_forest/python/ops/stats_ops.py index 80afcfb251f4d6455a9eb8ba5df4a6e43d2feb1c..7ac68fed20c3c9dfeaff05013e3fc686eea8cc2e 100644 --- a/tensorflow/contrib/tensor_forest/python/ops/stats_ops.py +++ b/tensorflow/contrib/tensor_forest/python/ops/stats_ops.py @@ -98,9 +98,9 @@ class FertileStatsVariable(tracking.TrackableResource): self._container = container self._init_op = None super(FertileStatsVariable, self).__init__() - self._resource_handle = self.create_resource() + self._resource_handle = self._create_resource() - def create_resource(self): + def _create_resource(self): if context.executing_eagerly(): # TODO(allenl): This will leak memory due to kernel caching by the # shared_name attribute value (but is better than the alternative of @@ -112,7 +112,7 @@ class FertileStatsVariable(tracking.TrackableResource): return gen_stats_ops.fertile_stats_resource_handle_op( self._container, shared_name=shared_name, name=self._name) - def initialize(self): + def _initialize(self): return gen_stats_ops.create_fertile_stats_variable( self.resource_handle, self._stats_config, @@ -121,7 +121,7 @@ class FertileStatsVariable(tracking.TrackableResource): @property def initializer(self): if self._init_op is None: - self._init_op = self.initialize() + self._init_op = self._initialize() return self._init_op def is_initialized(self): diff --git a/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc b/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc index 11335d7da637c813b301b4d4657462f4aae0c190..b683c14c0d77ebac74ad4d9b479c5ed493a3900a 100644 --- a/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc +++ b/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc @@ -21,9 +21,10 @@ limitations under the License. #include #define EIGEN_USE_GPU -#include "tensorflow/core/framework/op_kernel.h" #include "cuda/include/cuda_runtime_api.h" +#include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/platform/stream_executor.h" +#include "tensorflow/core/util/cuda_launch_config.h" namespace tensorflow { namespace tensorrt { @@ -38,8 +39,8 @@ void IncrementKernel(const float* d_input, float inc, float* d_output, int threads_per_block = 256; int blocks_per_grid = (count + threads_per_block - 1) / threads_per_block; - VecInc<<>>(d_input, inc, - d_output, count); + TF_CHECK_OK(CudaLaunchKernel(VecInc, threads_per_block, blocks_per_grid, 0, + stream, d_input, inc, d_output, count)); } // Note: this kernel definition is not needed in the plugin_test rule, but it is diff --git a/tensorflow/contrib/tpu/BUILD b/tensorflow/contrib/tpu/BUILD index e274ff081aa37b88b07a29c9499fa8355ad4798d..ee1cd3213efb0fff3a99536bdf1abd93c0c32a6e 100644 --- a/tensorflow/contrib/tpu/BUILD +++ b/tensorflow/contrib/tpu/BUILD @@ -18,6 +18,7 @@ package( "//learning/brain:__subpackages__", "//learning/deepmind:__subpackages__", "//medical/pathology:__subpackages__", + "//smartass/brain:__subpackages__", "//tensorflow:__subpackages__", "//vr/perception:__subpackages__", ], @@ -28,8 +29,7 @@ py_library( srcs = ["python/ops/tpu_ops.py"], srcs_version = "PY2AND3", deps = [ - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:tpu_ops_gen", + "//tensorflow/python/tpu:tpu_py", ], ) @@ -38,19 +38,7 @@ py_library( srcs = ["python/tpu/async_checkpoint.py"], srcs_version = "PY2AND3", deps = [ - "//tensorflow/python:array_ops", - "//tensorflow/python:control_flow_ops", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:init_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python:platform", - "//tensorflow/python:state_ops", - "//tensorflow/python:summary", - "//tensorflow/python:summary_ops_v2", - "//tensorflow/python:training", - "//tensorflow/python:variable_scope", - "//tensorflow/python:variables", - "//tensorflow/python/estimator:estimator_py", + "//tensorflow/python/tpu:async_checkpoint", ], ) @@ -72,24 +60,7 @@ py_library( ":tpu_embedding", ":tpu_lib", "//tensorflow/contrib/training:training_py", - "//tensorflow/core:protos_all_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:control_flow_ops", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:function", - "//tensorflow/python:init_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python:platform", - "//tensorflow/python:session", - "//tensorflow/python:state_ops", - "//tensorflow/python:summary", - "//tensorflow/python:summary_ops_v2", - "//tensorflow/python:training", - "//tensorflow/python:variable_scope", - "//tensorflow/python:variables", - "//tensorflow/python/estimator:estimator_py", - "//tensorflow/python/estimator:util", - "@six_archive//:six", + "//tensorflow/python/tpu:tpu_estimator", ], ) @@ -101,7 +72,7 @@ py_library( "//visibility:public", ], deps = [ - "//tensorflow/python:tpu_ops_gen", + "//tensorflow/python/tpu:functional", ], ) @@ -110,10 +81,7 @@ py_library( srcs = ["python/profiler/__init__.py"], srcs_version = "PY2AND3", deps = [ - "//tensorflow/contrib/tpu/profiler:tpu_profiler_analysis_pb2_grpc", - "//tensorflow/core/profiler:profiler_analysis_proto_py", - "//tensorflow/core/profiler:protos_all_py", - "//tensorflow/python:util", + "//tensorflow/python/tpu/profiler", ], ) @@ -130,6 +98,7 @@ py_library( ":tpu_embedding", ":tpu_estimator", ":tpu_lib", + "//tensorflow/python/tpu", ], ) @@ -196,29 +165,9 @@ py_library( ":functional", ":profiler", ":tpu_py", - "//tensorflow/compiler/xla/experimental/xla_sharding", - "//tensorflow/compiler/xla/python_api:xla_shape", "//tensorflow/contrib/cluster_resolver:cluster_resolver_py", "//tensorflow/contrib/compiler:xla", - "//tensorflow/core:protos_all_py", - "//tensorflow/core/protobuf/tpu:compilation_result_proto_py", - "//tensorflow/core/protobuf/tpu:dynamic_padding_proto_py", - "//tensorflow/core/protobuf/tpu:optimization_parameters_proto_py", - "//tensorflow/core/protobuf/tpu:topology_proto_py", - "//tensorflow/core/protobuf/tpu:tpu_embedding_configuration_proto_py", - "//tensorflow/core/protobuf/tpu:tpu_embedding_output_layout_proto_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:control_flow_ops", - "//tensorflow/python:control_flow_util", - "//tensorflow/python:dtypes", - "//tensorflow/python:framework", - "//tensorflow/python:framework_ops", - "//tensorflow/python:tensor_shape", - "//tensorflow/python:tpu_ops_gen", - "//tensorflow/python:training", - "//tensorflow/python:util", - "//tensorflow/python:variable_scope", - "//tensorflow/python/ops/losses", + "//tensorflow/python/tpu:tpu_lib", ], ) @@ -229,106 +178,7 @@ py_library( ], srcs_version = "PY2AND3", deps = [ - "//tensorflow/contrib/data/python/ops:batching", - "//tensorflow/contrib/data/python/ops:interleave_ops", - "//tensorflow/python:dtypes", - "//tensorflow/python:function", - "//tensorflow/python:functional_ops", - "//tensorflow/python/data/ops:dataset_ops", - "//tensorflow/python/data/ops:iterator_ops", - "//tensorflow/python/data/ops:readers", - ], -) - -tf_py_test( - name = "datasets_test", - size = "medium", - srcs = ["python/tpu/datasets_test.py"], - additional_deps = [ - "//tensorflow/python:client_testlib", - ":datasets", - ], - grpc_enabled = True, - shard_count = 4, - tags = ["no_oss"], -) - -tf_py_test( - name = "tpu_test", - size = "small", - srcs = ["python/tpu/tpu_test.py"], - additional_deps = [ - ":tpu", - "//tensorflow/python:client_testlib", - "//tensorflow/python:dtypes", - "//tensorflow/python:framework", - "//tensorflow/python:layers", - ], - tags = ["no_windows"], # TODO: needs investigation on Windows -) - -tf_py_test( - name = "tpu_sharding_test", - size = "small", - srcs = ["python/tpu/tpu_sharding_test.py"], - additional_deps = [ - ":tpu", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework", - ], -) - -tf_py_test( - name = "bfloat16_test", - size = "small", - srcs = ["python/tpu/bfloat16_test.py"], - additional_deps = [ - ":tpu", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework", - ], -) - -tf_py_test( - name = "tpu_infeed_test", - size = "small", - srcs = ["python/tpu/tpu_infeed_test.py"], - additional_deps = [ - ":tpu", - "//tensorflow/python:framework", - "//tensorflow/python:framework_test_lib", - ], -) - -tf_py_test( - name = "tpu_config_test", - size = "small", - srcs = ["python/tpu/tpu_config_test.py"], - additional_deps = [ - ":tpu_estimator", - "//tensorflow/python:framework", - "//tensorflow/python:framework_test_lib", - ], -) - -tf_py_test( - name = "tpu_estimator_signals_test", - size = "small", - srcs = ["python/tpu/tpu_estimator_signals_test.py"], - additional_deps = [ - ":tpu_estimator", - "//tensorflow/python:framework", - "//tensorflow/python:framework_test_lib", - ], -) - -tf_py_test( - name = "topology_test", - size = "medium", - srcs = ["python/tpu/topology_test.py"], - additional_deps = [ - ":tpu", - "//tensorflow/python:framework_test_lib", + "//tensorflow/python/tpu:datasets", ], ) @@ -341,16 +191,7 @@ py_library( srcs_version = "PY2AND3", deps = [ ":tpu_lib", - "//tensorflow/core/protobuf/tpu:tpu_embedding_configuration_proto_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:init_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python:partitioned_variables", - "//tensorflow/python:tpu_ops_gen", - "//tensorflow/python:variable_scope", - "//tensorflow/python:variables", - "@six_archive//:six", + "//tensorflow/python/tpu:tpu_embedding", ], ) @@ -359,31 +200,6 @@ py_library( srcs = ["python/tpu/feature_column.py"], deps = [ ":tpu_lib", - "//tensorflow/python:framework_ops", - "//tensorflow/python:init_ops", - "//tensorflow/python:variable_scope", - "//tensorflow/python/feature_column", - "//tensorflow/python/feature_column:feature_column_py", - ], -) - -tf_py_test( - name = "feature_column_test", - srcs = [ - "python/tpu/feature_column_test.py", - ], - additional_deps = [ - ":feature_column", - "//third_party/py/numpy", - "//tensorflow/python:dtypes", - "//tensorflow/python:framework_ops", - "//tensorflow/python:lookup_ops", - "//tensorflow/python:parsing_ops", - "//tensorflow/python:session", - "//tensorflow/python:sparse_tensor", - "//tensorflow/python:variables", - "//tensorflow/python/feature_column", - "//tensorflow/python/feature_column:feature_column_py", + "//tensorflow/python/tpu:feature_column", ], - main = "python/tpu/feature_column_test.py", ) diff --git a/tensorflow/contrib/tpu/profiler/BUILD b/tensorflow/contrib/tpu/profiler/BUILD index 0c3743cd118110e8c77445a78ac89de7184263dd..e2ce77e118182bb07193cbac82e176d3b2057e17 100644 --- a/tensorflow/contrib/tpu/profiler/BUILD +++ b/tensorflow/contrib/tpu/profiler/BUILD @@ -2,23 +2,6 @@ licenses(["notice"]) # Apache 2.0 load("//tensorflow:tensorflow.bzl", "tf_cc_binary") load("//tensorflow:tensorflow.bzl", "tf_cc_test") -load("//tensorflow/core:platform/default/build_config.bzl", "tf_proto_library") -load("//tensorflow/core:platform/default/build_config.bzl", "tf_additional_all_protos") - -cc_library( - name = "dump_tpu_profile", - srcs = ["dump_tpu_profile.cc"], - hdrs = ["dump_tpu_profile.h"], - visibility = ["//visibility:public"], - deps = [ - ":trace_events_to_json", - "//tensorflow/core:framework", - "//tensorflow/core:grpc_services", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core/profiler:protos_all_cc", - ], -) cc_library( name = "version", @@ -31,43 +14,13 @@ tf_cc_binary( srcs = [ "capture_tpu_profile.cc", ], + tags = ["no_windows"], visibility = ["//visibility:public"], deps = [ ":version", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", + "//tensorflow/core/platform/cloud:gcs_file_system", "//tensorflow/core/profiler/rpc/client:capture_profile", ], ) - -cc_library( - name = "trace_events_to_json", - srcs = ["trace_events_to_json.cc"], - hdrs = ["trace_events_to_json.h"], - deps = [ - "//tensorflow/core:lib", - "//tensorflow/core/profiler:protos_all_cc", - "@jsoncpp_git//:jsoncpp", - ], -) - -tf_cc_test( - name = "trace_events_to_json_test", - srcs = ["trace_events_to_json_test.cc"], - deps = [ - ":trace_events_to_json", - "//tensorflow/core:lib", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - "//tensorflow/core/profiler:protos_all_cc", - "@jsoncpp_git//:jsoncpp", - ], -) - -py_library( - name = "tpu_profiler_analysis_pb2_grpc", - srcs = ["tpu_profiler_analysis_pb2_grpc.py"], - srcs_version = "PY2AND3", - visibility = ["//visibility:public"], - deps = ["//tensorflow/core/profiler:profiler_analysis_proto_py"], -) diff --git a/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc b/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc index 508b929658dd18de9eda13d87c46a22891d1d36e..f11d1a9f37eeb19b95a876bd68575022e6b91521 100644 --- a/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc +++ b/tensorflow/contrib/tpu/profiler/capture_tpu_profile.cc @@ -76,7 +76,8 @@ int main(int argc, char** argv) { std::cout << usage.c_str() << std::endl; return 2; } - tensorflow::Status status = + tensorflow::Status status; + status = tensorflow::profiler::client::ValidateHostPortPair(FLAGS_service_addr); if (!status.ok()) { std::cout << status.error_message() << std::endl; @@ -103,9 +104,14 @@ int main(int argc, char** argv) { tensorflow::profiler::client::StartMonitoring( FLAGS_service_addr, duration_ms, FLAGS_monitoring_level, num_queries); } else { - tensorflow::profiler::client::StartTracing( + status = tensorflow::profiler::client::StartTracing( FLAGS_service_addr, FLAGS_logdir, FLAGS_workers_list, FLAGS_include_dataset_ops, duration_ms, num_tracing_attempts); + if (!status.ok()) { + std::cout << status.error_message() << std::endl; + std::cout << usage.c_str() << std::endl; + return 2; + } } return 0; } diff --git a/tensorflow/contrib/tpu/python/ops/tpu_ops.py b/tensorflow/contrib/tpu/python/ops/tpu_ops.py index ec0d5fec44e1687c20946c700769efe5b818af68..8605bae5c128513186d8c03835dcf49d3e4b6fd9 100644 --- a/tensorflow/contrib/tpu/python/ops/tpu_ops.py +++ b/tensorflow/contrib/tpu/python/ops/tpu_ops.py @@ -1,418 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ============================================================================= - -"""Operations for TPUs.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import platform - -from tensorflow.contrib.tpu.python.tpu import tpu_function -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.platform import tf_logging as logging - -if platform.system() != "Windows": - # pylint: disable=wildcard-import,unused-import,g-import-not-at-top - from tensorflow.python.ops import gen_tpu_ops - from tensorflow.python.ops.gen_tpu_ops import * - # pylint: enable=wildcard-import,unused-import,g-import-not-at-top - - def _create_default_group_assignment(): - num_shards = tpu_function.get_tpu_context().number_of_shards - if num_shards is None: - logging.warning( - "cross_replica_sum should be used within a tpu_shard_context, but " - "got unset number_of_shards. Assuming 1.") - num_shards = 1 - group_assignment = [list(range(num_shards))] - return group_assignment - - def all_to_all(x, - concat_dimension, - split_dimension, - split_count, - group_assignment=None, - name=None): - """Exchange data across TPU replicas. - - Args: - x: The local tensor. - concat_dimension: The dimension number to concatenate. - split_dimension: The dimension number to split. - split_count: The number of splits, this number must equal to the sub-group - size(group_assignment.get_shape()[1]) - group_assignment: Optional 2d int32 lists with shape [num_groups, - num_replicas_per_group]. `group_assignment[i]` represents the replica - ids in the ith subgroup. - name: Optional op name. - - Returns: - A `Tensor` which is concatenated by data from different replicas. - """ - if group_assignment is None: - group_assignment = _create_default_group_assignment() - return gen_tpu_ops.all_to_all( - x, - group_assignment, - concat_dimension=concat_dimension, - split_dimension=split_dimension, - split_count=split_count, - name=name) - - @ops.RegisterGradient("AllToAll") - def _all_to_all_grad(op, grad): - # The gradient of a all-to-all is also a all-to-all but the - # split_dimension and concat_dimension is swapped. - # The graident with respect to group_assignment is None. - return [ - gen_tpu_ops.all_to_all( - grad, - op.inputs[1], - concat_dimension=op.get_attr("split_dimension"), - split_dimension=op.get_attr("concat_dimension"), - split_count=op.get_attr("split_count")), None - ] - - def cross_replica_sum(x, group_assignment=None, name=None): - """Sum the input tensor across replicas according to group_assignment. - - Args: - x: The local tensor to the sum. - group_assignment: Optional 2d int32 lists with shape [num_groups, - num_replicas_per_group]. `group_assignment[i]` represents the replica - ids in the ith subgroup. - name: Optional op name. - - Returns: - A `Tensor` which is summed across replicas. - """ - if group_assignment is None: - group_assignment = _create_default_group_assignment() - - return gen_tpu_ops.cross_replica_sum(x, group_assignment, name=name) - - def collective_permute(x, source_target_pairs, name=None): - """Permute the input tensor across replicas given source_target_pairs. - - For each source_target_pair , we send replica a's input to replica b. - Each replica id must only appear once in the source column. Also it must - only appear once in the target column. - For the replica id not in the target column, this op returns a zero tensor - with the same shape and dtype of the input x. - - For example, suppose there are 4 TPU instances: `[A, B, C, D]`. Passing - source_target_pairs=`[[0,1],[1,2],[2,3]]` gets the outputs: - `[0, A, B, C]`. - - Args: - x: The local tensor to be permuted. - source_target_pairs: 2d int lists with shape [num_pairs, 2]. - source_target_pairs[i][0] represents the source replica id and - source_target_pairs[i][1] represents the target replica id. - name: Optional op name. - - Returns: - A `Tensor` which is permuted. - """ - return gen_tpu_ops.collective_permute(x, source_target_pairs, name=name) - - @ops.RegisterGradient("CollectivePermute") - def _collective_permute_grad(op, grad): - # The gradient of a collective permute operation is also a collective - # permute, but with source/target pairs reversed. The gradient with respect - # to input argument `source_target_pairs` is `None`. - source_target_pairs = op.inputs[1][:, ::-1] - return [gen_tpu_ops.collective_permute(grad, source_target_pairs), None] - - @ops.RegisterGradient("CrossReplicaSum") - def _cross_replica_sum_grad(op, grad): - # The gradient of a cross replica sum is also a cross-replica sum. - # The gradient with respect to group_assignment is None. - return [gen_tpu_ops.cross_replica_sum(grad, op.inputs[1]), None] - - # This extra type checking exists to give a more helpful error message in - # the common case that uint8 and int64 values are infed. Remove when both - # types are supported. - - _SUPPORTED_INFEED_DTYPES = set([ - dtypes.bool, dtypes.int32, dtypes.int64, dtypes.bfloat16, dtypes.float32, - dtypes.complex64, dtypes.uint32 - ]) - - @ops.RegisterGradient("TPUEmbeddingActivations") - def _embedding_activations_grad(activations_op, grad_wrt_activations): - """Saves the gradient of embedding activations ops in a graph collection.""" - g = ops.get_default_graph() - table_id = activations_op.get_attr("table_id") - lookup_id = activations_op.get_attr("lookup_id") - table_gradients = g.get_collection_ref( - "tpu_embedding_gradients_table_%d" % table_id) - - if not table_gradients: - raise RuntimeError( - "Gradients for TPUEmbedding have been generated in non-training mode." - "This is not expected. Consider putting your Optimizer.minimize code " - "behind the training mode condition check. For Estimator, you can " - "do \n\n" - " if mode == tf.estimator.ModeKeys.TRAIN:\n" - " train_op = opt.minimize(loss)\n" - "\n") - - table_gradients[lookup_id] = array_ops.identity(grad_wrt_activations) - return [ - # RegisterGradient requires that value be returned for all inputs. Since - # the first argument (tpu_gradient_variable_{table_name}) has shape [1], - # we will return zeros(shape=[1]). The actual gradient w.r.t. the - # embedding activations (grad_wrt_activations) has the same shape as the - # activations returned by embedding_activations. - array_ops.zeros(arg.shape, dtype=dtypes.float32) - for arg in activations_op.inputs - ] - - def infeed_dequeue(dtype, shape, name=None): - """A placeholder op for a value that will be fed into the computation. - - Args: - dtype: A `tf.DType`. The type of elements in the tensor. - shape: A `tf.TensorShape` or list of `ints`. The shape of the tensor. - name: A name for the operation (optional). - - Returns: - A `Tensor` of type `dtype`. - A tensor that will be provided using the infeed mechanism. - - Raises: - TypeError: If 'dtype` is not a supported infeed type. - """ - if dtype not in _SUPPORTED_INFEED_DTYPES: - raise TypeError( - "{} is not a supported TPU infeed type. Supported types are: " - "{}".format(dtype, list(_SUPPORTED_INFEED_DTYPES))) - - return gen_tpu_ops.infeed_dequeue(dtype, shape, name=name) - - # pylint: disable=redefined-outer-name - def infeed_dequeue_tuple(dtypes, shapes, name=None): - """A placeholder op for values fed into the TPU simultaneously as a tuple. - - Args: - dtypes: A list of `tf.DType`s that has length `>= 1`. - The element types of each element in `outputs`. - shapes: A list of shapes (each a `tf.TensorShape` or list of `ints`). - The shapes of each tensor in `outputs`. - name: A name for the operation (optional). - - Returns: - A list of `Tensor` objects of type `dtypes`. - A list of tensors that will be provided using the infeed mechanism. - - Raises: - TypeError: If a type in 'dtypes` is not a supported infeed type. - """ - for dtype in dtypes: - if dtype not in _SUPPORTED_INFEED_DTYPES: - raise TypeError( - "{} is not a supported TPU infeed type. Supported types are: " - "{}".format(dtype, list(_SUPPORTED_INFEED_DTYPES))) - return gen_tpu_ops.infeed_dequeue_tuple(dtypes, shapes, name=name) - # pylint: enable=redefined-outer-name - - # pylint: disable=protected-access - def send_tpu_embedding_gradients(inputs, - config, - learning_rates=None, - name=None): - """A placeholder op for feeding per-sample gradients to the embedding layer. - - Args: - inputs: A TensorList of gradients with which to update embedding tables. - This argument has the same length and shapes as the return value of - RecvTPUEmbeddingActivations, but contains gradients of the model's - loss with respect to the embedding activations. The embedding tables - are updated from these gradients via the optimizers specified in the - TPU embedding configuration given to tpu.initialize_system. - config: Serialized TPUEmbeddingConfiguration proto. - learning_rates: A TensorList of float32 scalars, one for each dynamic - learning rate tag: see the comments in - //third_party/tensorflow/core/protobuf/tpu/ - optimization_parameters.proto. - Multiple tables can share the same dynamic learning rate tag as - specified in the configuration. If the learning rates for all tables - are constant, this list should be empty. - name: A name for the operation (optional). - - Returns: - A SendTPUEmbeddingGradients operation. - """ - if learning_rates is None: - learning_rates = [] - return gen_tpu_ops.send_tpu_embedding_gradients( - inputs=inputs, learning_rates=learning_rates, config=config, name=name) - - send_tpu_embedding_gradients.__doc__ = ( - gen_tpu_ops.send_tpu_embedding_gradients.__doc__) - - # pylint: disable=protected-access - def enqueue_tpu_embedding_integer_batch(batch, - device_ordinal, - mode_override=None, - name=None): - """A placeholder op for enqueueing embedding IDs to the TPU. - - Args: - batch: A list of 1D tensors, one for each embedding table, containing the - indices into the tables. - device_ordinal: The TPU device to use. Should be >= 0 and less than the - number of TPU cores in the task on which the node is placed. - mode_override: A string input that overrides the mode specified in the - TPUEmbeddingConfiguration. Supported values are {'unspecified', - 'inference', 'training', 'backward_pass_only'}. When set to - 'unspecified', the mode set in TPUEmbeddingConfiguration is used, - otherwise mode_override is used (optional). - name: A name for the operation (optional). - - Returns: - An EnqueueTPUEmbeddingIntegerBatch operation. - """ - if mode_override is None: - mode_override = "unspecified" - return gen_tpu_ops.enqueue_tpu_embedding_integer_batch( - batch=batch, - device_ordinal=device_ordinal, - mode_override=mode_override, - name=name) - - enqueue_tpu_embedding_integer_batch.__doc__ = ( - gen_tpu_ops.enqueue_tpu_embedding_integer_batch.__doc__) - - # pylint: disable=protected-access - def enqueue_tpu_embedding_sparse_batch(sample_indices, - embedding_indices, - aggregation_weights, - device_ordinal, - combiners=None, - mode_override=None, - name=None): - """A placeholder op for enqueueing embedding IDs to the TPU. - - Args: - sample_indices: A list of rank 1 Tensors specifying the training example - and feature to which the corresponding embedding_indices and - aggregation_weights values belong. sample_indices[i] must equal b * nf + - f, where nf is the number of features from the corresponding table, f is - in [0, nf), and b is in [0, batch size). - embedding_indices: A list of rank 1 Tensors, indices into the embedding - tables. - aggregation_weights: A list of rank 1 Tensors containing per sample -- - i.e. per (training example, feature) -- aggregation weights. - device_ordinal: The TPU device to use. Should be >= 0 and less than the - number of TPU cores in the task on which the node is placed. - combiners: A list of string scalars, one for each embedding table that - specify how to normalize the embedding activations after weighted - summation. Supported combiners are 'mean', 'sum', or 'sqrtn'. It is - invalid to have the sum of the weights be 0 for 'mean' or the sum of the - squared weights be 0 for 'sqrtn'. If combiners isn't passed, the default - is to use 'sum' for all tables (optional). - mode_override: A string input that overrides the mode specified in the - TPUEmbeddingConfiguration. Supported values are {'unspecified', - 'inference', 'training', 'backward_pass_only'}. When set to - 'unspecified', the mode set in TPUEmbeddingConfiguration is used, - otherwise mode_override is used (optional). - name: A name for the operation (optional). - - Returns: - An EnqueueTPUEmbeddingSparseBatch operation. - """ - if mode_override is None: - mode_override = "unspecified" - return gen_tpu_ops.enqueue_tpu_embedding_sparse_batch( - sample_indices=sample_indices, - embedding_indices=embedding_indices, - aggregation_weights=aggregation_weights, - device_ordinal=device_ordinal, - combiners=combiners, - mode_override=mode_override, - name=name) - - enqueue_tpu_embedding_sparse_batch.__doc__ = ( - gen_tpu_ops.enqueue_tpu_embedding_sparse_batch.__doc__) - - # pylint: disable=protected-access - def enqueue_tpu_embedding_sparse_tensor_batch(sample_indices, - embedding_indices, - aggregation_weights, - table_ids, - device_ordinal, - combiners=None, - mode_override=None, - name=None): - """A placeholder op for enqueueing embedding IDs to the TPU. - - Args: - sample_indices: A list of rank 1 Tensors specifying the training example - to which the corresponding embedding_indices and aggregation_weights - values belong. It corresponds to sp_ids.indices[:,0] in - embedding_lookup_sparse(). - embedding_indices: A list of rank 1 Tensors, indices into the embedding - tables. It corresponds to sp_ids.values in embedding_lookup_sparse(). - aggregation_weights: A list of rank 1 Tensors containing per training - example aggregation weights. It corresponds to sp_weights.values in - embedding_lookup_sparse(). - table_ids: A list of integers specifying the identifier of the embedding - table (offset of TableDescriptor in the TPUEmbeddingConfiguration) to - lookup the corresponding input. The ith input is looked up using - table_ids[i]. The size of the table_ids list must be equal to that of - sample_indices, embedding_indices and aggregation_weights. - device_ordinal: The TPU device to use. Should be >= 0 and less than the - number of TPU cores in the task on which the node is placed. - combiners: A list of string scalars, one for each embedding table that - specify how to normalize the embedding activations after weighted - summation. Supported combiners are 'mean', 'sum', or 'sqrtn'. It is - invalid to have the sum of the weights be 0 for 'mean' or the sum of the - squared weights be 0 for 'sqrtn'. If combiners isn't passed, the default - is to use 'sum' for all tables (optional). - mode_override: A string input that overrides the mode specified in the - TPUEmbeddingConfiguration. Supported values are {'unspecified', - 'inference', 'training', 'backward_pass_only'}. When set to - 'unspecified', the mode set in TPUEmbeddingConfiguration is used, - otherwise mode_override is used (optional). - name: A name for the operation (optional). - - Returns: - An EnqueueTPUEmbeddingSparseTensorBatch operation. - """ - if mode_override is None: - mode_override = "unspecified" - return gen_tpu_ops.enqueue_tpu_embedding_sparse_tensor_batch( - sample_indices=sample_indices, - embedding_indices=embedding_indices, - aggregation_weights=aggregation_weights, - table_ids=table_ids, - device_ordinal=device_ordinal, - combiners=combiners, - mode_override=mode_override, - name=name) - - enqueue_tpu_embedding_sparse_tensor_batch.__doc__ = ( - gen_tpu_ops.enqueue_tpu_embedding_sparse_tensor_batch.__doc__) - -else: - # We have already built the appropriate libraries into the binary via CMake - # if we have built contrib, so we don't need this - pass +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.ops.tpu_ops import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/ops/tpu_ordinal_selector_op.py b/tensorflow/contrib/tpu/python/ops/tpu_ordinal_selector_op.py index 6917ac2e1a769378c77dcdcd0d63da2028a3a34c..788e1fe0568cf2f406c379e4d928100ea51a37a3 100644 --- a/tensorflow/contrib/tpu/python/ops/tpu_ordinal_selector_op.py +++ b/tensorflow/contrib/tpu/python/ops/tpu_ordinal_selector_op.py @@ -1,35 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ============================================================================= - -"""Operations to select TPU core to run.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import platform - -if platform.system() != "Windows": - # pylint: disable=wildcard-import,unused-import,g-import-not-at-top - from tensorflow.python.ops.gen_tpu_ops import tpu_ordinal_selector - - from tensorflow.contrib.util import loader - from tensorflow.python.platform import resource_loader - # pylint: enable=wildcard-import,unused-import,g-import-not-at-top - -else: - # We have already built the appropriate libraries into the binary via CMake - # if we have built contrib, so we don't need this - pass +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.ops.tpu_ordinal_selector_op import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/profiler/__init__.py b/tensorflow/contrib/tpu/python/profiler/__init__.py index 0c183aaf53543f7bb38475525d1777048925ff62..aeb061dbe114bc287946b50d08a86778c78c7b38 100644 --- a/tensorflow/contrib/tpu/python/profiler/__init__.py +++ b/tensorflow/contrib/tpu/python/profiler/__init__.py @@ -1,31 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ============================================================================= - -"""Classes for TPU trace events.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function # pylint: disable=wildcard-import,unused-import -from tensorflow.core.profiler.trace_events_pb2 import * -from tensorflow.core.profiler.profiler_analysis_pb2 import * +from tensorflow.python.tpu.profiler import * # pylint: enable=wildcard-import,unused-import - -from tensorflow.python.util.all_util import remove_undocumented - -_allowed_symbols = ['Trace', 'Resource', 'Device', 'TraceEvent'] - -remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/tpu/python/tpu/__init__.py b/tensorflow/contrib/tpu/python/tpu/__init__.py index 0dffd7064b19f353aed6afa3ad383564643a4a90..82d4f68c0221013706f70bcf54ae4c97cc7db1d3 100644 --- a/tensorflow/contrib/tpu/python/tpu/__init__.py +++ b/tensorflow/contrib/tpu/python/tpu/__init__.py @@ -1,20 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ============================================================================= - -"""Ops related to Tensor Processing Units.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function + +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/_tpu_estimator_embedding.py b/tensorflow/contrib/tpu/python/tpu/_tpu_estimator_embedding.py index 98aa7827fcf38b10e97318067ffa99008e93c557..41aa4d267812cabe775459723df7e01efaa83c93 100644 --- a/tensorflow/contrib/tpu/python/tpu/_tpu_estimator_embedding.py +++ b/tensorflow/contrib/tpu/python/tpu/_tpu_estimator_embedding.py @@ -1,334 +1,23 @@ -# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# =================================================================== -"""Tooling for support TPU embedding in TPUEstimator.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections - -from tensorflow.contrib.tpu.python.tpu import feature_column as tpu_fc -from tensorflow.contrib.tpu.python.tpu import tpu_embedding -from tensorflow.python.estimator import model_fn as model_fn_lib -from tensorflow.python.feature_column import feature_column as core_fc -from tensorflow.python.feature_column import feature_column_lib as core_fc_lib - -# pylint: disable=protected-access -_TPU_EMBEDDING_COLUMN_CLASSES = (tpu_fc._TPUEmbeddingColumn, - tpu_fc._TPUSharedEmbeddingColumn) -_EMBEDDING_COLUMN_CLASSES = (core_fc._EmbeddingColumn, - core_fc_lib.EmbeddingColumn, - core_fc._SharedEmbeddingColumn) -_SUPPORTED_FEATURE_COLUMNS = (core_fc._NumericColumn, core_fc_lib.NumericColumn) - -# pylint: enable=protected-access - -_TABLE_NAME_PREFIX = 'tbl_' -_LEN_TABLE_NAME_PREFIX = len(_TABLE_NAME_PREFIX) - - -def _get_table_name_from_embedding_var_name(embedding_var_name): - return '{}{}'.format(_TABLE_NAME_PREFIX, embedding_var_name) - - -def _get_embedding_var_name_from_table_name(table_name): - return table_name[_LEN_TABLE_NAME_PREFIX:] - - -def _get_embedding_variable_name(scope_name, var_name): - return '{}/{}'.format(scope_name, var_name) - - -def _get_slot_variable_names(scope_name, var_name, optimization_parameters): - """Return embedding variable names which are consistent with CPU runs.""" - if isinstance(optimization_parameters, tpu_embedding.AdagradParameters): - return tpu_embedding.AdagradSlotVariableName( - '{}/{}/Adagrad'.format(scope_name, var_name) - ) - elif isinstance(optimization_parameters, tpu_embedding.AdamParameters): - return tpu_embedding.AdamSlotVariableNames( - '{}/{}/Adam/m'.format(scope_name, var_name), - '{}/{}/Adam/v'.format(scope_name, var_name) - ) - elif isinstance(optimization_parameters, - tpu_embedding.StochasticGradientDescentParameters): - return None - else: - raise ValueError('Support to infer full variable name ' - 'for optimization_parameter {} has not been added.' - .format(optimization_parameters)) - - -def get_full_variable_names( - graph, table_to_config_dict, optimization_parameters): - """Return embedding variable names and slot variables which are consistent with CPU runs.""" - collection = graph.get_collection_ref(tpu_fc._TPU_FC_TO_SCOPE) # pylint: disable=protected-access - if not collection: - raise RuntimeError( - 'Embedding feature column did not capture any thing. Make sure the ' - 'feature columns passed to TPUEstimator constructor is properly ' - 'used in model_fn.') - - embedding_variable_name_by_table = {} - slot_variable_names_by_table = {} - for table_name in table_to_config_dict: - embedding_var_name = _get_embedding_var_name_from_table_name(table_name) - (scope_name, var_name) = collection[0][embedding_var_name] - embedding_variable_name_by_table[table_name] = ( - _get_embedding_variable_name(scope_name, var_name)) - slot_variable_names_by_table[table_name] = _get_slot_variable_names( - scope_name, var_name, optimization_parameters) - - graph.clear_collection(tpu_fc._TPU_FC_TO_SCOPE) # pylint: disable=protected-access - return embedding_variable_name_by_table, slot_variable_names_by_table - - -def get_tpu_embedding_config_from_feature_columns(feature_columns): - """Create configs for TPUEmbedding from a list of feature columns. - - This function will place one embedding tensor per table and the return is - intended to be used as input to TPUEmbedding. - - Args: - feature_columns: a list of supported feature columns. - - Returns: - A pair of dicts, the first maps tables to their config, the second maps - features to tables. - """ - - allowed = (tpu_fc._TPUEmbeddingColumn, tpu_fc._TPUSharedEmbeddingColumn) # pylint: disable=protected-access - - for column in feature_columns: - if not isinstance(column, allowed): - raise TypeError( - 'Unsupported feature column {}. Supported types are {}.'.format( - type(column), allowed)) - - table_to_config = {} - feature_to_table = {} - for column in feature_columns: - feature_name = column.get_feature_key_name() - table_name = _get_table_name_from_embedding_var_name( - column.get_embedding_var_name()) - if feature_name in feature_to_table: - raise ValueError( - 'Feature column {} is used with multiple embeddings and this is ' - 'not supported.'.format(feature_name)) - feature_to_table[feature_name] = table_name - vocabulary_size, dimension = column.get_embedding_table_size() - table_to_config[table_name] = tpu_embedding.TableConfig( - vocabulary_size=vocabulary_size, - dimension=dimension, - initializer=column.get_initializer(), - combiner=column.get_combiner()) - - return table_to_config, feature_to_table - - -def _get_tpu_embedding_optimization_parameters(embedding_config_spec): - """Get tpu_embedding._OptimizationParameters from EmbeddingConfigSpec.""" - if embedding_config_spec.optimizer_type == 'adagrad': - return tpu_embedding.AdagradParameters( - embedding_config_spec.learning_rate, - embedding_config_spec.adagrad_initial_accumulator, - embedding_config_spec.use_gradient_accumulation) - elif embedding_config_spec.optimizer_type == 'sgd': - return tpu_embedding.StochasticGradientDescentParameters( - embedding_config_spec.learning_rate, - embedding_config_spec.use_gradient_accumulation) - elif embedding_config_spec.optimizer_type == 'adam': - return tpu_embedding.AdamParameters( - embedding_config_spec.learning_rate, - embedding_config_spec.adam_parameters.beta1, - embedding_config_spec.adam_parameters.beta2, - embedding_config_spec.adam_parameters.epsilon, - use_gradient_accumulation=embedding_config_spec - .use_gradient_accumulation) - else: - raise ValueError('optimizer_type must be adagrad or sgd or adam for now.') - - -AdamParameters = collections.namedtuple('AdamParameters', - ['beta1', 'beta2', 'epsilon']) - - -# TODO(shizhiw): Improve the API to support more optimizer parameters in API. -class EmbeddingConfigSpec( - collections.namedtuple('EmbeddingConfigSpec', [ - 'feature_columns', 'learning_rate', 'optimizer_type', - 'adagrad_initial_accumulator', 'clipping_limit', - 'use_gradient_accumulation', 'adam_parameters' - ])): - """Class to keep track of embedding config specification.""" - - def __new__(cls, - feature_columns, - learning_rate, - optimizer_type='adagrad', - adagrad_initial_accumulator=None, - clipping_limit=None, - use_gradient_accumulation=False, - adam_parameters=None): - """Creates an EmbeddingConfigSpec instance. - - Args: - feature_columns: All `FeatureColumn`s used by model. - learning_rate: embedding optimizer learning rate. - optimizer_type: (String) Name of the optimizer for embedding gradients - updates. Must be either 'adagrad' ( `tf.train.AdagradOptimizer`, default - value), 'sgd' (`tf.train.GradientDescentOptimizer`), or 'adam' - (`tf.contrib.opt.LazyAdamOptimizer`) for lazy Adam. This optimizer will - be applied to all embedding variables specified by `feature_columns`. - adagrad_initial_accumulator: Initial accumulator for Adagrad. Used when - optimizer_type is 'adagrad'. Default is `0.1`. - clipping_limit: (Optional) Clipping limit (absolute value). - use_gradient_accumulation: (Experimental) Whether to accumulate the - gradients across TPU embedding mini-batches. Gradient accumulation does - not affect SGD and therefore this is applicable only for Adagrad. - adam_parameters: AdamParameters. Used when optimizer_type is 'adam'. - Default is 0.9 for beta1, 0.999 for beta2 and 1e-8 for epsilon. - - Returns: - An EmbeddingConfigSpec instance. - - Raises: - ValueError: If the feature_columns are not specified. - TypeError: If the feature columns are not of ths correct type (one of - _SUPPORTED_FEATURE_COLUMNS, _TPU_EMBEDDING_COLUMN_CLASSES OR - _EMBEDDING_COLUMN_CLASSES). - ValueError: If use_gradient_accumulation is True for SGD. - ValueError: If `optimizer_type` is not one of "adagrad" or "sgd" or - "adam". - """ - if not feature_columns: - raise ValueError('`feature_columns` cannot be `None` or empty.') - - # It is unknown at this moment, whether the TPUEstimator is running in CPU - # or TPU mode. So allow non-TPU embedding columns also. - supported_classes = tuple( - list(_SUPPORTED_FEATURE_COLUMNS) + list(_TPU_EMBEDDING_COLUMN_CLASSES) + - list(_EMBEDDING_COLUMN_CLASSES)) - - for column in feature_columns: - if not isinstance(column, supported_classes): - raise TypeError( - 'All feature columns must be supported types in {}. Got {}'.format( - supported_classes, type(column))) - - if optimizer_type == 'adagrad': - if adagrad_initial_accumulator is None: - adagrad_initial_accumulator = 0.1 - if adagrad_initial_accumulator <= 0: - raise ValueError('Adagrad initial_accumulator must be positive') - elif optimizer_type == 'sgd': - if use_gradient_accumulation: - raise ValueError('Gradient accumulation makes sense for Adagrad only.') - elif optimizer_type == 'adam': - if adam_parameters is None: - adam_parameters = AdamParameters(0.9, 0.999, 1e-8) - if adam_parameters.beta1 < 0. or adam_parameters.beta1 >= 1.: - raise ValueError('beta1 must be between 0. and 1; got {}.'.format( - adam_parameters.beta1)) - if adam_parameters.beta2 < 0. or adam_parameters.beta2 >= 1.: - raise ValueError('beta2 must be between 0. and 1; got {}.'.format( - adam_parameters.beta2)) - if adam_parameters.epsilon <= 0.: - raise ValueError('epsilon must be positive; got {}.'.format( - adam_parameters.epsilon)) - else: - raise ValueError('optimizer_type must be adagrad or sgd or adam for now.') - - return super(EmbeddingConfigSpec, cls).__new__( - cls, - feature_columns=feature_columns, - learning_rate=learning_rate, - optimizer_type=optimizer_type, - adagrad_initial_accumulator=adagrad_initial_accumulator, - clipping_limit=clipping_limit, - use_gradient_accumulation=use_gradient_accumulation, - adam_parameters=adam_parameters) - - -class EmbeddingConfig(object): - """This is the internal immutable object for embedding config. - - `_EmbeddingConfig` is responsible to _translate_ user provided - `EmbeddingConfigSpec` to internal data structures, mostly constructor - arguments of `TPUEmbedding`. - """ - - def __init__(self, embedding_config_spec, train_batch_size, eval_batch_size, - num_hosts, num_cores, master): - self._embedding_config_spec = embedding_config_spec - self._train_batch_size = train_batch_size - self._eval_batch_size = eval_batch_size - self._num_hosts = num_hosts - self._num_cores = num_cores - self._master = master - - self._table_to_config_dict, self._feature_to_table_dict = ( - get_tpu_embedding_config_from_feature_columns( - embedding_config_spec.feature_columns)) - self._optimization_parameters = _get_tpu_embedding_optimization_parameters( - self._embedding_config_spec) - self._mode_to_tpu_embedding_dict = {} - self.dummy_table_variables = None - - def has_embedding_tables(self): - return bool(self._table_to_config_dict) - - def _create_tpu_embedding(self, mode): - """Create tpu_embedding.TPUEmbedding based on mode.""" - if mode == model_fn_lib.ModeKeys.TRAIN: - batch_size = self._train_batch_size - else: - batch_size = self._eval_batch_size - - if mode == model_fn_lib.ModeKeys.TRAIN: - tpu_embedding_mode = tpu_embedding.TRAINING - elif (mode == model_fn_lib.ModeKeys.EVAL or - mode == model_fn_lib.ModeKeys.PREDICT): - tpu_embedding_mode = tpu_embedding.INFERENCE - else: - raise ValueError('Mode {} is not supported.'.format(mode)) - - tpu_embedding_ = tpu_embedding.TPUEmbedding( - self._table_to_config_dict, - self._feature_to_table_dict, - batch_size, - tpu_embedding_mode, - self._master, - self._optimization_parameters, - ) - return tpu_embedding_ - - def get_tpu_embedding(self, mode): - if mode not in self._mode_to_tpu_embedding_dict: - self._mode_to_tpu_embedding_dict[mode] = ( - self._create_tpu_embedding(mode)) - return self._mode_to_tpu_embedding_dict[mode] - - -def split_inputs(ctx, features, labels): - """Splits the dense and sparse tensors inside the features and labels.""" - sparse_features = collections.OrderedDict() - if ctx.embedding_config: - tpu_embedding_ = ctx.embedding_config.tpu_embedding - for feature_key in tpu_embedding_.feature_to_table_dict: - sparse_features[feature_key] = features.pop(feature_key) - - return features, labels, sparse_features +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu._tpu_estimator_embedding import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/async_checkpoint.py b/tensorflow/contrib/tpu/python/tpu/async_checkpoint.py index 1b09ce173a64ba3f93ec019c8fd65dc4710f0fcf..5eb8034e47474873ccef0b6123f2becd0668738c 100644 --- a/tensorflow/contrib/tpu/python/tpu/async_checkpoint.py +++ b/tensorflow/contrib/tpu/python/tpu/async_checkpoint.py @@ -1,212 +1,23 @@ -# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # -# Licensed under the Apache License, Version 2.0 (the 'License'); +# Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an 'AS IS' BASIS, +# distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ====================================== -"""Hook for asynchronous checkpointing. - -This hook dispatches checkpoint writing operations in a separate thread to -allow execution to continue on the main thread. -""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import os -import threading -import time - -from tensorflow.core.util.event_pb2 import SessionLog -from tensorflow.python.framework import meta_graph -from tensorflow.python.framework import ops -from tensorflow.python.platform import tf_logging as logging -from tensorflow.python.training import basic_session_run_hooks -from tensorflow.python.training import training_util -from tensorflow.python.training.session_run_hook import SessionRunArgs -from tensorflow.python.training.summary_io import SummaryWriterCache - - -class AsyncCheckpointSaverHook(basic_session_run_hooks.CheckpointSaverHook): - """Saves checkpoints every N steps or seconds.""" - - def __init__(self, - checkpoint_dir, - save_secs=None, - save_steps=None, - saver=None, - checkpoint_basename="model.ckpt", - scaffold=None, - listeners=None): - """Initializes a `CheckpointSaverHook`. - - Args: - checkpoint_dir: `str`, base directory for the checkpoint files. - save_secs: `int`, save every N secs. - save_steps: `int`, save every N steps. - saver: `Saver` object, used for saving. - checkpoint_basename: `str`, base name for the checkpoint files. - scaffold: `Scaffold`, use to get saver object. - listeners: List of `CheckpointSaverListener` subclass instances. Used for - callbacks that run immediately before or after this hook saves the - checkpoint. - - Raises: - ValueError: One of `save_steps` or `save_secs` should be set. - ValueError: At most one of `saver` or `scaffold` should be set. - """ - logging.info("Create AsyncCheckpointSaverHook.") - if saver is not None and scaffold is not None: - raise ValueError("You cannot provide both saver and scaffold.") - self._saver = saver - self._save_thread = None - self._write_graph_thread = None - self._checkpoint_dir = checkpoint_dir - self._save_path = os.path.join(checkpoint_dir, checkpoint_basename) - self._scaffold = scaffold - self._timer = basic_session_run_hooks.SecondOrStepTimer( - every_secs=save_secs, every_steps=save_steps) - self._listeners = listeners or [] - self._steps_per_run = 1 - self._summary_writer = None - self._global_step_tensor = None - - self._last_checkpoint_step = None - - def _set_steps_per_run(self, steps_per_run): - self._steps_per_run = steps_per_run - - def begin(self): - self._summary_writer = SummaryWriterCache.get(self._checkpoint_dir) - self._global_step_tensor = training_util._get_or_create_global_step_read() # pylint: disable=protected-access - if self._global_step_tensor is None: - raise RuntimeError( - "Global step should be created to use CheckpointSaverHook.") - for l in self._listeners: - l.begin() - - def after_create_session(self, session, coord): - global_step = session.run(self._global_step_tensor) - - # We do write graph and saver_def at the first call of before_run. - # We cannot do this in begin, since we let other hooks to change graph and - # add variables in begin. Graph is finalized after all begin calls. - def _write_graph_fn(self): - training_util.write_graph( - ops.get_default_graph().as_graph_def(add_shapes=True), - self._checkpoint_dir, "graph.pbtxt") - self._write_graph_thread = threading.Thread(target=_write_graph_fn, - args=[self]) - self._write_graph_thread.start() - - saver_def = self._get_saver().saver_def if self._get_saver() else None - graph = ops.get_default_graph() - meta_graph_def = meta_graph.create_meta_graph_def( - graph_def=graph.as_graph_def(add_shapes=True), saver_def=saver_def) - self._summary_writer.add_graph(graph) - self._summary_writer.add_meta_graph(meta_graph_def) - # The checkpoint saved here is the state at step "global_step". - self._save(session, global_step) - self._timer.update_last_triggered_step(global_step) - - def before_run(self, run_context): # pylint: disable=unused-argument - return SessionRunArgs(self._global_step_tensor) - - def after_run(self, run_context, run_values): - global_step = run_context.session.run(self._global_step_tensor) - if self._timer.should_trigger_for_step(global_step): - self._timer.update_last_triggered_step(global_step) - logging.info("Triggering checkpoint. %s", global_step) - if self._save(run_context.session, global_step): - run_context.request_stop() - - def end(self, session): - if self._save_thread: - logging.info("Waiting for any pending checkpoints to finish.") - self._save_thread.join() - if self._write_graph_thread: - logging.info("Waiting for any pending write_graph to finish.") - self._write_graph_thread.join() - - last_step = session.run(self._global_step_tensor) - - if self._last_checkpoint_step != last_step: - self._save(session, last_step, asynchronous=False) - - for l in self._listeners: - l.end(session, last_step) - - def _save(self, session, step, asynchronous=True): - """Saves the latest checkpoint, returns should_stop.""" - - # Skip saving on step 0 - if step == 0: - return - - def _save_fn(): - """Run the saver process.""" - logging.info("Saving checkpoints for %d into %s.", step, self._save_path) - - start_time = time.time() - for l in self._listeners: - l.before_save(session, step) - - self._get_saver().save(session, self._save_path, global_step=step) - self._summary_writer.add_session_log( - SessionLog( - status=SessionLog.CHECKPOINT, checkpoint_path=self._save_path), - step) - - for l in self._listeners: - l.after_save(session, step) - - end_time = time.time() - logging.info("Checkpoint actual writing time: (%.3f sec)", - end_time - start_time) - logging.info("Checkpoint finished for %d into %s.", step, self._save_path) - - if not asynchronous: - self._last_checkpoint_step = step - _save_fn() - return - - if self._save_thread is not None: - self._save_thread.join(timeout=0.1) - if self._save_thread.is_alive(): - logging.info("Saver thread still in progress, skipping checkpoint.") - return - - self._last_checkpoint_step = step - self._save_thread = threading.Thread(target=_save_fn) - self._save_thread.start() - - def _get_saver(self): - if self._saver is not None: - return self._saver - elif self._scaffold is not None: - return self._scaffold.saver - - # Get saver from the SAVERS collection if present. - collection_key = ops.GraphKeys.SAVERS - savers = ops.get_collection(collection_key) - if not savers: - raise RuntimeError( - "No items in collection {}. Please add a saver to the collection " - "or provide a saver or scaffold.".format(collection_key)) - elif len(savers) > 1: - raise RuntimeError( - "More than one item in collection {}. " - "Please indicate which one to use by passing it to the constructor." - .format(collection_key)) - - self._saver = savers[0] - return savers[0] +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.async_checkpoint import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/bfloat16.py b/tensorflow/contrib/tpu/python/tpu/bfloat16.py index fa74f651aa63c72d14eb78c8af479263810e9b7d..f3d392a8daec2a80f974d90051324a02be002afd 100644 --- a/tensorflow/contrib/tpu/python/tpu/bfloat16.py +++ b/tensorflow/contrib/tpu/python/tpu/bfloat16.py @@ -1,77 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ============================================================================= - -"""Helper context for running models with bfloat16.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.python.framework import dtypes -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import variable_scope -from tensorflow.python.util import tf_contextlib - - -def _get_custom_getter(): - """Returns a custom getter that this class's methods must be called under. - - All methods of this class must be called under a variable scope that was - passed this custom getter. Example: - - ```python - network = ConvNetBuilder(...) - with tf.variable_scope('cg', custom_getter=network.get_custom_getter()): - network.conv(...) - # Call more methods of network here - ``` - - Currently, this custom getter only does anything if self.use_tf_layers is - True. In that case, it causes variables to be stored as dtype - self.variable_type, then casted to the requested dtype, instead of directly - storing the variable as the requested dtype. - """ - - def inner_custom_getter(getter, *args, **kwargs): - """Custom getter that forces variables to have type self.variable_type.""" - cast_to_bfloat16 = False - requested_dtype = kwargs['dtype'] - if requested_dtype == dtypes.bfloat16: - # Only change the variable dtype if doing so does not decrease variable - # precision. - kwargs['dtype'] = dtypes.float32 - cast_to_bfloat16 = True - var = getter(*args, **kwargs) - # This if statement is needed to guard the cast, because batch norm - # assigns directly to the return value of this custom getter. The cast - # makes the return value not a variable so it cannot be assigned. Batch - # norm variables are always in fp32 so this if statement is never - # triggered for them. - if cast_to_bfloat16: - var = math_ops.cast(var, dtypes.bfloat16) - return var - - return inner_custom_getter - - -@tf_contextlib.contextmanager -def bfloat16_scope(): - """Scope class for bfloat16 variables so that the model uses custom getter. - - This enables variables to be read as bfloat16 type when using get_variable. - """ - with variable_scope.variable_scope( - '', custom_getter=_get_custom_getter()) as varscope: - yield varscope +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.bfloat16 import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/datasets.py b/tensorflow/contrib/tpu/python/tpu/datasets.py index bc0cd41d210ac6f8de1b20ebf744ee1e1dd04137..c20aac7e36aa31c5a9d88ca6fe02a8703f9ed5a3 100644 --- a/tensorflow/contrib/tpu/python/tpu/datasets.py +++ b/tensorflow/contrib/tpu/python/tpu/datasets.py @@ -1,191 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ====================================== -"""Library of Cloud TPU helper functions for data loading.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.python.data.experimental.ops import batching -from tensorflow.python.data.experimental.ops import interleave_ops -from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.data.ops import iterator_ops -from tensorflow.python.data.ops import readers -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import function -from tensorflow.python.framework import ops -from tensorflow.python.ops import functional_ops - - -def _TextLineDataset(filename): - buffer_size = 8 * 1024 * 1024 # 8 MiB per file - dataset = readers.TextLineDataset(filename, buffer_size=buffer_size) - return dataset - - -def _TFRecordDataset(filename): - buffer_size = 8 * 1024 * 1024 # 8 MiB per file - dataset = readers.TFRecordDataset(filename, buffer_size=buffer_size) - return dataset - - -_FILETYPE_MAP = { - 'tfrecord': _TFRecordDataset, - 'textline': _TextLineDataset, - 'text': _TextLineDataset, -} - - -def StreamingFilesDataset(files, - filetype=None, - file_reader_job=None, - worker_job=None, - num_epochs=None, - filename_shuffle_buffer_size=None, - num_parallel_reads=None, - batch_transfer_size=None, - sloppy=None): - """StreamingFilesDataset constructs a dataset to stream from workers (GCE VM). - - Because Cloud TPUs are allocated over the network, a Cloud TPU cannot read - files local to your GCE VM. In order to train using files stored on your local - VM (e.g. on local SSD for extreme performance), use the StreamingFilesDataset - helper to generate a dataset to feed your Cloud TPU with files from your GCE - VM. - - The resulting dataset may return an OutOfRangeError if there are no files - found as a result of the fileglob expansion. - - Note: StreamingFilesDataset assumes that the session is using a - TPUClusterResolver and has therefore a worker and a coordinator job. File - loading will be done on the coordinator job. - - Args: - files: A string glob to match files, or a `tf.data.Dataset` generating file - names. - filetype: A string (one of 'tfrecord', or 'textline') or a single-argument - TensorFlow function that when given a filename returns a dataset. - file_reader_job: An optional string that corresponds to the job that should - perform the file reads. - worker_job: An optional string that corresponds to the job that should - process the tensors (i.e. your GPU or TPU worker). - num_epochs: The number of epochs through the training set that should be - generated. By default, it will repeat infinitely. - filename_shuffle_buffer_size: An optional integer whose value controls the - shuffling of the file names. If you would like to read from the files in - the same order, set to 0 or False. - num_parallel_reads: An optional integer controlling the number of files to - read from concurrently. (Set to 1 for no parallelism.) - batch_transfer_size: An optional integer controlling the batching used to - amortize the remote function invocation overhead. Set to a very large - number to increase throughput. Set to a very small number to reduce memory - consumption. Set to False to skip batching. - sloppy: (Optional.) If `False`, read input data while maintaining a - deterministic order. (This may have significant performance impacts.) - sloppy defaults to: True. - Returns: - A `tf.data.Dataset` with an infinite stream of elements generated by a - parallel interleaving of the set of files matched (or generated) by `files` - with a type is the output of the dataset specified by `filetype`. - - Raises: - ValueError: if any argument is not of the expected type. - """ - if filetype is None: - filetype = 'tfrecord' - - if isinstance(filetype, str): - if filetype not in _FILETYPE_MAP: - raise ValueError('Unexpected filetype: %s' % filetype) - reader_fn = _FILETYPE_MAP[filetype] - elif callable(filetype): - reader_fn = filetype - else: - raise ValueError('filetype should be a string or a callable') - - file_reader_job = file_reader_job or 'coordinator' - - worker_job = worker_job or 'worker' - - if filename_shuffle_buffer_size is None: - filename_shuffle_buffer_size = 4096 - - num_parallel_reads = num_parallel_reads or 8 - - if batch_transfer_size is None: - batch_transfer_size = 256 - - if sloppy is None: - sloppy = True - - with ops.device('/job:%s' % file_reader_job): - if isinstance(files, str): - source_dataset = dataset_ops.Dataset.list_files(files) - elif isinstance(files, dataset_ops.DatasetV2): - source_dataset = files - else: - raise ValueError('files was not a string or a dataset: %s' % files) - - if filename_shuffle_buffer_size: - source_dataset = source_dataset.shuffle( - buffer_size=filename_shuffle_buffer_size) - - source_dataset = source_dataset.apply( - interleave_ops.parallel_interleave( - reader_fn, cycle_length=num_parallel_reads, sloppy=sloppy)) - - source_dataset = source_dataset.repeat(num_epochs) - - if batch_transfer_size: - source_dataset = source_dataset.batch(batch_transfer_size) - - source_dataset = source_dataset.prefetch(1) - - source_iterator = dataset_ops.make_one_shot_iterator(source_dataset) - source_handle = source_iterator.string_handle() - - @function.Defun(dtypes.string) - def LoadingFunc(h): - remote_iterator = iterator_ops.Iterator.from_string_handle( - h, source_dataset.output_types, source_dataset.output_shapes) - return remote_iterator.get_next() - - def MapFn(unused_input): - if isinstance(source_dataset.output_types, dtypes.DType): - output_types = [source_dataset.output_types] - elif isinstance(source_dataset.output_types, (list, tuple)): - output_types = source_dataset.output_types - else: - raise ValueError('source dataset has invalid output types') - remote_calls = functional_ops.remote_call( - args=[source_handle], - Tout=output_types, - f=LoadingFunc, - target='/job:%s/replica:0/task:0/cpu:0' % file_reader_job) - if len(remote_calls) == 1: - return remote_calls[0] - else: - return remote_calls - - with ops.device('/job:%s' % worker_job): - output_dataset = dataset_ops.Dataset.range(2).repeat().map( - MapFn, num_parallel_calls=4 if sloppy else None) - output_dataset = output_dataset.prefetch(1) - - if batch_transfer_size: - # Undo the batching used during the transfer. - output_dataset = output_dataset.apply(batching.unbatch()).prefetch(1) - - return output_dataset +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.datasets import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/device_assignment.py b/tensorflow/contrib/tpu/python/tpu/device_assignment.py index 112c6bb42d28daa1bcad14387412dfb57a13fea2..05dffef3a1efdae2ad7306ca5ad3bc7a9eac04cf 100644 --- a/tensorflow/contrib/tpu/python/tpu/device_assignment.py +++ b/tensorflow/contrib/tpu/python/tpu/device_assignment.py @@ -1,313 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ====================================== -"""Library of TPU helper functions.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import math -import numpy as np -from six.moves import xrange # pylint: disable=redefined-builtin - -from tensorflow.contrib.tpu.python.tpu.topology import Topology - - -SINGLE_CORE_ASSIGNMENT = [[[0, 0, 0]]] - - -def _compute_task_and_cores_to_replicas(core_assignment, topology): - """Computes a nested dict which maps task and logical core to replicas.""" - task_and_cores_to_replicas = {} - for replica in xrange(core_assignment.shape[0]): - for logical_core in xrange(core_assignment.shape[1]): - coordinates = core_assignment[replica, logical_core, :] - task_id = topology.task_ordinal_at_coordinates(coordinates) - if task_id not in task_and_cores_to_replicas: - task_and_cores_to_replicas[task_id] = {} - if logical_core not in task_and_cores_to_replicas[task_id]: - task_and_cores_to_replicas[task_id][logical_core] = set() - - task_and_cores_to_replicas[task_id][logical_core].add(replica) - - task_to_sorted_replica_id = {} - - for task, core_to_replicas in task_and_cores_to_replicas.items(): - core_to_sorted_replicas = {} - for core, replicas in core_to_replicas.items(): - core_to_sorted_replicas[core] = sorted(replicas) - - task_to_sorted_replica_id[task] = core_to_sorted_replicas - return task_to_sorted_replica_id - - -class DeviceAssignment(object): - """Mapping from logical cores in a computation to the physical TPU topology. - - Prefer to use the `device_assignment()` helper to construct a - `DeviceAssignment`; it is easier if less flexible than constructing a - `DeviceAssignment` directly. - """ - - def __init__(self, topology, core_assignment): - """Constructs a `DeviceAssignment` object. - - Args: - topology: A `Topology` object that describes the physical TPU topology. - core_assignment: A logical to physical core mapping, represented as a - rank 3 numpy array. See the description of the `core_assignment` - property for more details. - - Raises: - ValueError: If `topology` is not `Topology` object. - ValueError: If `core_assignment` is not a rank 3 numpy array. - """ - if not isinstance(topology, Topology): - raise ValueError("topology must be a Topology object, got {}".format( - type(topology))) - core_assignment = np.asarray(core_assignment, dtype=np.int32) - - self._topology = topology - - if core_assignment.ndim != 3: - raise ValueError("core_assignment must be a rank 3 numpy array, " - "got shape {}".format(core_assignment.shape)) - - self._num_replicas = core_assignment.shape[0] - self._num_cores_per_replica = core_assignment.shape[1] - - if core_assignment.shape[-1] != topology.mesh_rank: - raise ValueError( - "minor dimension of core_assignment must have size equal to topology " - "rank ({}), got shape {}".format(topology.mesh_rank, - core_assignment.shape)) - - self._core_assignment = core_assignment - self._task_and_cores_to_replicas = _compute_task_and_cores_to_replicas( - self._core_assignment, topology) - - @property - def topology(self): - """A `Topology` that describes the TPU topology.""" - return self._topology - - @property - def num_cores_per_replica(self): - """The number of cores per replica.""" - return self._num_cores_per_replica - - @property - def num_replicas(self): - """The number of replicas of the computation.""" - return self._num_replicas - - @property - def core_assignment(self): - """The logical to physical core mapping. - - Returns: - An integer numpy array of rank 3, with shape - `[num_replicas, num_cores_per_replica, topology_rank]`. Maps - (replica, logical core) pairs to physical topology coordinates. - """ - return self._core_assignment - - def coordinates(self, replica, logical_core): - """Returns the physical topology coordinates of a logical core.""" - return tuple(self.core_assignment[replica, logical_core, :]) - - def lookup_replicas(self, task_id, logical_core): - """Lookup replica ids by task number and logical core. - - Args: - task_id: TensorFlow task number. - logical_core: An integer, identifying a logical core. - Returns: - A sorted list of the replicas that are attached to that task and - logical_core. - Raises: - ValueError: If no replica exists in the task which contains the logical - core. - """ - try: - return self._task_and_cores_to_replicas[task_id][logical_core] - except KeyError: - raise ValueError( - "Can not find any replica in task: {} contains logical_core: {} ". - format(task_id, logical_core)) - - def tpu_ordinal(self, replica=0, logical_core=0): - """Returns the ordinal of the TPU device assigned to a logical core.""" - coordinates = self.coordinates(replica, logical_core) - return self._topology.tpu_device_ordinal_at_coordinates(coordinates) - - def host_device(self, replica=0, logical_core=0, job=None): - """Returns the CPU device attached to a logical core.""" - coordinates = self.coordinates(replica, logical_core) - return self._topology.cpu_device_name_at_coordinates(coordinates, job=job) - - def tpu_device(self, replica=0, logical_core=0, job=None): - """Returns the name of the TPU device assigned to a logical core.""" - coordinates = self.coordinates(replica, logical_core) - return self._topology.tpu_device_name_at_coordinates(coordinates, job=job) - - -def device_assignment(topology, - computation_shape=None, - computation_stride=None, - num_replicas=1): - """Computes a device_assignment of a computation across a TPU topology. - - Attempts to choose a compact grid of cores for locality. - - Returns a `DeviceAssignment` that describes the cores in the topology assigned - to each core of each replica. - - `computation_shape` and `computation_stride` values should be powers of 2 for - optimal packing. - - Args: - topology: A `Topology` object that describes the TPU cluster topology. - To obtain a TPU topology, evaluate the `Tensor` returned by - `initialize_system` using `Session.run`. Either a serialized - `TopologyProto` or a `Topology` object may be passed. Note: you must - evaluate the `Tensor` first; you cannot pass an unevaluated `Tensor` here. - computation_shape: A rank 1 int32 numpy array with size equal to the - topology rank, describing the shape of the computation's block of cores. - If None, the `computation_shape` is `[1] * topology_rank`. - computation_stride: A rank 1 int32 numpy array of size `topology_rank`, - describing the inter-core spacing of the `computation_shape` cores in the - TPU topology. If None, the `computation_stride` is `[1] * topology_rank`. - num_replicas: The number of computation replicas to run. The replicas will - be packed into the free spaces of the topology. - - Returns: - A DeviceAssignment object, which describes the mapping between the logical - cores in each computation replica and the physical cores in the TPU - topology. - - Raises: - ValueError: If `topology` is not a valid `Topology` object. - ValueError: If `computation_shape` or `computation_stride` are not 1D int32 - numpy arrays with shape [3] where all values are positive. - ValueError: If computation's replicas cannot fit into the TPU topology. - """ - # Deserialize the Topology proto, if it is a string. - if isinstance(topology, bytes): - topology = Topology(serialized=topology) - - if not isinstance(topology, Topology): - raise ValueError("`topology` is not a Topology object; got {}".format( - type(topology))) - - topology_rank = len(topology.mesh_shape) - mesh_shape = topology.mesh_shape - if computation_shape is None: - computation_shape = np.array([1] * topology_rank, dtype=np.int32) - else: - computation_shape = np.asarray(computation_shape, dtype=np.int32) - - if computation_stride is None: - computation_stride = np.array([1] * topology_rank, dtype=np.int32) - else: - computation_stride = np.asarray(computation_stride, dtype=np.int32) - - if computation_shape.shape != (topology_rank,): - raise ValueError("computation_shape must have shape [{}]; got {}".format( - topology_rank, computation_shape.shape)) - if computation_stride.shape != (topology_rank,): - raise ValueError("computation_stride must have shape [{}]; got {}".format( - topology_rank, computation_stride.shape)) - - if any(computation_shape < 1): - raise ValueError( - "computation_shape must be positive; got computation_shape={}".format( - computation_shape)) - if any(computation_stride < 1): - raise ValueError( - "computation_stride must be positive; got computation_stride={}".format( - computation_stride)) - - # Computes the physical size of one computation instance. - computation_footprint = computation_shape * computation_stride - if any(computation_footprint > mesh_shape): - raise ValueError( - "computation footprint {} does not fit in TPU topology shape {}".format( - computation_footprint, mesh_shape)) - - # Computes how many copies of the computation footprint fit in the mesh. - block_counts = mesh_shape // computation_footprint - - replica_counts = block_counts * computation_stride - max_replicas = np.prod(replica_counts) - if num_replicas > max_replicas: - raise ValueError( - "requested {} replicas but only {} replicas with shape {} and " - "computation_stride {} fit in a TPU mesh of shape {}".format( - num_replicas, max_replicas, computation_shape, computation_stride, - mesh_shape)) - - def ceil_of_ratio(n, m): - return (n + m - 1) // m - - replica_shape = [0] * topology_rank - if num_replicas > 0: - remaining_replicas = num_replicas - remaining_dims = topology_rank - - # Choose dimensions as close to an equal cube as possible, in order of - # increasing dimension size. By visiting dimensions in increasing size, we - # assign the most constrained dimension first, so we won't make infeasible - # choices. - # - # As a secondary sort order, visit the dimensions in reverse order. This - # means we try to use both cores on the same chip in preference to two cores - # on different chips. - for x, ni in sorted(((x, -i) for (i, x) in enumerate(replica_counts))): - i = -ni - target_size = int(math.ceil(remaining_replicas**(1.0 / remaining_dims))) - replica_shape[i] = min(target_size, x) - remaining_replicas = ceil_of_ratio(remaining_replicas, replica_shape[i]) - remaining_dims -= 1 - - assert remaining_replicas == 1 and remaining_dims == 0 - - # Assigns an offset to each replica such that no two replicas overlap. - replica_offsets = np.full([num_replicas, topology_rank], -1, dtype=np.int32) - for replica in xrange(num_replicas): - # Chooses a replica number in each axis. - t = replica - pos = [] - for dim in replica_shape[::-1]: - pos.append(t % dim) - t //= dim - replica_pos = np.array(pos[::-1], dtype=np.int32) - - # Determines where that replica starts in each axis. - outer = replica_pos // computation_stride - inner = replica_pos % computation_stride - replica_offsets[replica, :] = outer * computation_footprint + inner - - # Computes a complete logical core -> physical core mapping for each replica. - indices = [ - np.arange(0, computation_shape[i] * computation_stride[i], - computation_stride[i]) for i in xrange(topology_rank) - ] - indices = np.concatenate( - [i[..., np.newaxis] for i in np.meshgrid(*indices, indexing="ij")], - axis=-1) - indices = indices.reshape((-1, topology_rank)) - assignment = indices + replica_offsets[:, np.newaxis, :] - return DeviceAssignment(topology, core_assignment=assignment) +# pylint: disable=wildcard-import,unused-import,redefined-builtin +from tensorflow.python.tpu.device_assignment import * +# pylint: enable=wildcard-import,unused-import,redefined-builtin diff --git a/tensorflow/contrib/tpu/python/tpu/error_handling.py b/tensorflow/contrib/tpu/python/tpu/error_handling.py index 52e1ea42370d653d1de7c12eee4b456ec7ce921c..1b1328b4075d9a737e40693c13e33e0b7c1fbedf 100644 --- a/tensorflow/contrib/tpu/python/tpu/error_handling.py +++ b/tensorflow/contrib/tpu/python/tpu/error_handling.py @@ -1,132 +1,23 @@ -# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# =================================================================== -"""ErrorRendezvous handler for collecting errors from multiple threads.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import contextlib -import sys -import threading -import time - -import six - -from tensorflow.python.framework import errors -from tensorflow.python.platform import tf_logging as logging - -_UNINTERESTING_ERRORS = (errors.CancelledError,) - - -class ErrorRendezvous(object): - """Resolve errors from multiple threads during TPU execution. - - TPU errors can occur on the infeed or outfeed threads as well as the main - training thread. - - Depending on which thread "wins" and receives the session error first, we may - end up showing users a confusing and non-actionable error message (session - cancelled) instead of a root cause (e.g. a bad filename). - - The rendezvous object provides a location to capture these errors until all - threads terminate. At that point we can choose the most informative error - to report. - """ - - def __init__(self, num_sources): - # string -> (message, traceback) - self._errors = {} - self._num_sources = num_sources - self._session_cancel_timer = None - - def record_error(self, source, exc_info, session=None): - """Report an exception from the given source. - - If a session is passed, a timer will be registered to close it after a few - seconds. This is necessary to ensure the main training loop does not hang - if an infeed/oufeed error occurs. We sleep a few seconds to allow a more - interesting error from another thread to propagate. - - Args: - source: string, source of the error - exc_info: Output from `sys.exc_info` (type, value, traceback) - session: Session to close after delay. - """ - _, value, _ = exc_info - self._errors[source] = exc_info - logging.info('Error recorded from %s: %s', source, value) - - if session is not None and self._session_cancel_timer is None: - - def _cancel_session(): - time.sleep(5) - try: - session.close() - except: # pylint: disable=bare-except - pass - - self._session_cancel_timer = threading.Thread(target=_cancel_session,) - self._session_cancel_timer.daemon = True - self._session_cancel_timer.start() - - def record_done(self, source): - """Mark execution source `source` as done. - - If an error was originally reported from `source` it is left intact. - - Args: - source: `str`, source being recorded - """ - logging.info('%s marked as finished', source) - if source not in self._errors: - self._errors[source] = None - - @contextlib.contextmanager - def catch_errors(self, source, session=None): - """Context manager to report any errors within a block.""" - try: - yield - except Exception: # pylint: disable=broad-except - self.record_error(source, sys.exc_info(), session) - - def raise_errors(self, timeout_sec=0): - """Wait for up to `timeout` seconds for all error sources to finish. - - Preferentially raise "interesting" errors (errors not in the - _UNINTERESTING_ERRORS) set. - - Args: - timeout_sec: Seconds to wait for other error sources. - """ - for _ in range(timeout_sec): - if len(self._errors) == self._num_sources: - break - time.sleep(1) - - kept_errors = [(k, v) for (k, v) in self._errors.items() if v is not None] - - # First check for any interesting errors, then fall back on the session - # cancelled errors etc. - for k, (typ, value, traceback) in kept_errors: - if isinstance(value, _UNINTERESTING_ERRORS): - continue - else: - logging.warn('Reraising captured error') - six.reraise(typ, value, traceback) - - for k, (typ, value, traceback) in kept_errors: - logging.warn('Reraising captured error') - six.reraise(typ, value, traceback) +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.error_handling import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/feature_column.py b/tensorflow/contrib/tpu/python/tpu/feature_column.py index e80d0de70548b4b628d8ac7ebd6c9c1e5ded26ed..ded75e975b10c4265370af260bf804687c9caebc 100644 --- a/tensorflow/contrib/tpu/python/tpu/feature_column.py +++ b/tensorflow/contrib/tpu/python/tpu/feature_column.py @@ -1,435 +1,30 @@ -# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# =================================================================== -"""TPU Feature Column Library.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" + from __future__ import absolute_import from __future__ import division from __future__ import print_function -import math - -from tensorflow.contrib.tpu.python.tpu import tpu -from tensorflow.contrib.tpu.python.tpu import tpu_function -from tensorflow.python.feature_column import feature_column as fc -from tensorflow.python.feature_column import feature_column_lib as fc_lib -from tensorflow.python.framework import ops -from tensorflow.python.ops import init_ops -from tensorflow.python.ops import variable_scope -# pylint: disable=protected-access - - -_TPU_FC_TO_SCOPE = '_tpu_feature_column_scope' -_SUPPORTED_CATEGORICAL_COLUMNS = (fc._IdentityCategoricalColumn, - fc._VocabularyFileCategoricalColumn, - fc._VocabularyListCategoricalColumn, - fc._WeightedCategoricalColumn, - fc_lib.IdentityCategoricalColumn, - fc_lib.VocabularyFileCategoricalColumn, - fc_lib.VocabularyListCategoricalColumn, - fc_lib.WeightedCategoricalColumn) - - -def embedding_column(categorical_column, - dimension, - combiner='mean', - initializer=None): - """TPU embedding_column for `tf.feature_column.embedding_column`. - - Note that the interface for TPU embedding_column is different from the non-TPU - version. The following args available for the non-TPU version are NOT - supported: ckpt_to_load_from, tensor_name_in_ckp, max_norm and trainable. - - Args: - categorical_column: A categorical_column returned from - categorical_column_with_identity, weighted_categorical_column, - categorical_column_with_vocabulary_list or - categorical_column_with_vocabulary_file. - dimension: An integer specifying dimension of the embedding, must be > 0. - combiner: A string specifying how to reduce if there are multiple entries - in a single row. For more information, see - `tf.feature_column.embedding_column`. - initializer: A variable initializer function to be used in embedding - variable initialization. If not specified, defaults to - `tf.truncated_normal_initializer` with mean `0.0` and standard deviation - `1/sqrt(dimension)`. - - Returns: - A _TPUEmbeddingColumn. - - Raises: - ValueError: if `dimension` not > 0. - ValueError: if `initializer` is specified but not callable. - """ - if not isinstance(categorical_column, _SUPPORTED_CATEGORICAL_COLUMNS): - raise TypeError( - 'categorical_column for tpu ' - ' embedding_column must be type %s, got %s.' % (' or '.join([ - cc.__name__ for cc in _SUPPORTED_CATEGORICAL_COLUMNS - ]), type(categorical_column))) - if (dimension is None) or (dimension < 1): - raise ValueError('Invalid dimension {}.'.format(dimension)) - - if (initializer is not None) and (not callable(initializer)): - raise ValueError('initializer must be callable if specified. ' - 'Embedding of column_name: {}'.format( - categorical_column.name)) - if initializer is None: - initializer = init_ops.truncated_normal_initializer( - mean=0.0, stddev=1 / math.sqrt(dimension)) - - embedding_shape = categorical_column._num_buckets, dimension # pylint: disable=protected-access - - def _creator(weight_collections, scope): - embedding_column_layer = fc._EmbeddingColumnLayer( - embedding_shape=embedding_shape, - initializer=initializer, - weight_collections=weight_collections, - trainable=True, - name='embedding_column_layer') - return embedding_column_layer(None, scope=scope) # pylint: disable=not-callable - - column = _TPUEmbeddingColumn( - categorical_column=categorical_column, - dimension=dimension, - combiner=combiner, - layer_creator=_creator, - ckpt_to_load_from=None, - tensor_name_in_ckpt=None, - max_norm=None, - trainable=True) - # For Embedding column, the initializer is hidden inside the creator Fn, which - # is not accessiable later. So, we attach it to a speicial field. Also note - # that non-TPU Embedding column and non-TPU shared Embedding column handle the - # initializer differently. See shared_embedding_columns for details. - column._tpu_initializer = initializer - return column - - -def shared_embedding_columns(categorical_columns, - dimension, - combiner='mean', - initializer=None, - shared_embedding_collection_name=None): - """List of dense columns that convert from sparse, categorical input.""" - for categorical_column in categorical_columns: - if not isinstance(categorical_column, _SUPPORTED_CATEGORICAL_COLUMNS): - raise TypeError( - 'categorical_column for tpu ' - ' shared_embedding_columns must be type %s, got %s.' % (' or '.join([ - cc.__name__ for cc in _SUPPORTED_CATEGORICAL_COLUMNS - ]), type(categorical_column))) - columns = fc_lib.shared_embedding_columns( - categorical_columns, - dimension, - combiner=combiner, - initializer=initializer, - shared_embedding_collection_name=shared_embedding_collection_name, - ckpt_to_load_from=None, - tensor_name_in_ckpt=None, - max_norm=None, - trainable=True) - - # Use the initializer and shared_embedding_collection_name to create TPU - # version - initializer = columns[0].initializer - shared_embedding_collection_name = columns[0].shared_embedding_collection_name - tpu_columns = [] - - # Create the state (_SharedEmbeddingColumnLayer) here. - for categorical_column in categorical_columns: - column = _TPUSharedEmbeddingColumn( - categorical_column=categorical_column, - dimension=dimension, - combiner=combiner, - initializer=initializer, - shared_embedding_collection_name=shared_embedding_collection_name, - ckpt_to_load_from=None, - tensor_name_in_ckpt=None, - max_norm=None, - trainable=True) - tpu_columns.append(column) - - return tpu_columns - - -class _TPUBaseEmbeddingColumn(object): - """Base class for TPU Embedding Column.""" - - def __init__(self, categorical_column): - self._tpu_categorical_column = categorical_column - - def get_combiner(self): - """Returns the embedding combiner.""" - raise NotImplementedError('not implemented') - - def get_embedding_table_size(self): - """Returns the embedding table size, tuple of vocab size and dimension.""" - raise NotImplementedError('not implemented') - - def get_feature_key_name(self): - """Returns the feature key name in the features dict.""" - raise NotImplementedError('not impl') - - def get_weight_key_name(self): - """Return the key name for weights.""" - raise NotImplementedError('not impl') - - def get_embedding_var_name(self): - """Returns the embedding variable name. - - Feature key name and embedding variable name are usually one-to-one mapping. - But for shared embedding columns, it is many-to-one mapping. - """ - raise NotImplementedError('not impl') - - def get_initializer(self): - """Returns the initializer.""" - raise NotImplementedError('not impl') - - def is_categorical_column_weighted(self): - """Check if the categorical column of the embedding column is weighted.""" - raise NotImplementedError('not impl') - - -class _TPUEmbeddingColumn(_TPUBaseEmbeddingColumn, fc._EmbeddingColumn): - """Core Embedding Column.""" - - def __new__(cls, - categorical_column, - dimension, - combiner='mean', - layer_creator=None, - ckpt_to_load_from=None, - tensor_name_in_ckpt=None, - max_norm=None, - trainable=True): - # Note, args ckpt_to_load_from, tensor_name_in_ckpt, max_norm and trainable - # are not supported on TPU. They are solely for matching the signature of - # __new__ of parent class fc._EmbeddingColumn. - return fc._EmbeddingColumn.__new__( - cls, - categorical_column, - dimension, - combiner=combiner, - layer_creator=layer_creator, - ckpt_to_load_from=ckpt_to_load_from, - tensor_name_in_ckpt=tensor_name_in_ckpt, - max_norm=max_norm, - trainable=trainable) - - def __init__(self, - categorical_column, - dimension, - combiner='mean', - layer_creator=None, - ckpt_to_load_from=None, - tensor_name_in_ckpt=None, - max_norm=None, - trainable=True): - _TPUBaseEmbeddingColumn.__init__(self, categorical_column) - self._key = None - - def get_combiner(self): - return self.combiner - - def get_embedding_table_size(self): - """Returns num_ids and width.""" - return (self.categorical_column._num_buckets, self.dimension) - - def get_feature_key_name(self): - """get_feature_key_name.""" - if self.is_categorical_column_weighted(): - return self.categorical_column.categorical_column.name - return self.categorical_column.name - - def get_weight_key_name(self): - """get_weight_key_name.""" - if self.is_categorical_column_weighted(): - return self.categorical_column.weight_feature_key - return None - - def get_embedding_var_name(self): - """get_embedding_var_name.""" - return self.categorical_column.name - - def get_initializer(self): - return self._tpu_initializer - - def is_categorical_column_weighted(self): - """Check if the categorical column of the embedding column is weighted.""" - if isinstance( - self.categorical_column, - ( - fc._WeightedCategoricalColumn, # pylint: disable=protected-access - fc_lib.WeightedCategoricalColumn)): - return True - return False - - def _get_dense_tensor(self, inputs, weight_collections=None, trainable=None): - if tpu.under_tpu_inference_context(): - def host_computation(): - return fc._EmbeddingColumn._get_dense_tensor( - self, inputs, weight_collections, trainable) - return tpu.outside_compilation(host_computation) - - if _is_running_on_cpu(): - return fc._EmbeddingColumn._get_dense_tensor( - self, inputs, weight_collections, trainable) - - # TPU mode - # Get the embeddings from the LazyBuilder. - tensor = inputs.get(self.get_feature_key_name()) - - # Add to collection for _create_tpu_embedding_variables_and_ops - _record_variable_scope_and_name(self.get_embedding_var_name(), - 'embedding_weights') - - return tensor - - -class _TPUSharedEmbeddingColumn(_TPUBaseEmbeddingColumn, - fc._SharedEmbeddingColumn): - """Core Shared Embedding Column.""" - - def __new__(cls, - categorical_column, - dimension, - combiner='mean', - initializer=None, - shared_embedding_collection_name=None, - ckpt_to_load_from=None, - tensor_name_in_ckpt=None, - max_norm=None, - trainable=True): - return fc._SharedEmbeddingColumn.__new__( - cls, - categorical_column, - dimension, - combiner=combiner, - initializer=initializer, - shared_embedding_collection_name=shared_embedding_collection_name, - ckpt_to_load_from=ckpt_to_load_from, - tensor_name_in_ckpt=tensor_name_in_ckpt, - max_norm=max_norm, - trainable=trainable) - - def __init__(self, - categorical_column, - dimension, - combiner='mean', - initializer=None, - shared_embedding_collection_name=None, - ckpt_to_load_from=None, - tensor_name_in_ckpt=None, - max_norm=None, - trainable=True): - - _TPUBaseEmbeddingColumn.__init__(self, categorical_column) - self._key = None - - def get_combiner(self): - return self.combiner - - def get_embedding_table_size(self): - """Returns num_ids and width.""" - return (self.categorical_column._num_buckets, self.dimension) - - def get_feature_key_name(self): - """get_feature_key_name.""" - if self.is_categorical_column_weighted(): - return self.categorical_column.categorical_column.name - return self.categorical_column.name - - def get_weight_key_name(self): - """get_weight_key_name.""" - if self.is_categorical_column_weighted(): - return self.categorical_column.weight_feature_key - return None - - def get_embedding_var_name(self): - """get_embedding_var_name.""" - return self.shared_embedding_collection_name - - def get_initializer(self): - return self.initializer - - def is_categorical_column_weighted(self): - """Check if the categorical column of the embedding column is weighted.""" - if isinstance( - self.categorical_column, - ( - fc._WeightedCategoricalColumn, # pylint: disable=protected-access - fc_lib.WeightedCategoricalColumn)): - return True - return False - - def _get_dense_tensor(self, inputs, weight_collections=None, trainable=None): - if tpu.under_tpu_inference_context(): - def host_computation(): - return fc._SharedEmbeddingColumn._get_dense_tensor( - self, inputs, weight_collections, trainable) - return tpu.outside_compilation(host_computation) - - if _is_running_on_cpu(): - return fc._SharedEmbeddingColumn._get_dense_tensor( - self, inputs, weight_collections, trainable) - - # TPU mode - # Get the embeddings from the LazyBuilder. - tensor = inputs.get(self.get_feature_key_name()) - - # Add to collection for _create_tpu_embedding_variables_and_ops - _record_variable_scope_and_name( - self.get_embedding_var_name(), - 'embedding_weights', - is_shared_embedding=True) - return tensor - - -def _record_variable_scope_and_name(embedding_var_name, - embedding_var_name_in_fc, - is_shared_embedding=False): - """Add embedding variable name and scope to collection.""" - g = ops.get_default_graph() - collection = g.get_collection_ref(_TPU_FC_TO_SCOPE) - if not collection: - collection.append({}) - - var_def_dict = collection[0] - - captured_scope = variable_scope.get_variable_scope() - captured_scope_name = captured_scope.name - - if embedding_var_name in var_def_dict: - if (var_def_dict[embedding_var_name][0] != captured_scope_name - and not is_shared_embedding): - raise ValueError( - 'For embedding var name {}, the variable scope name is different, ' - 'got {}; expected {}'.format(embedding_var_name, - captured_scope_name, - var_def_dict[embedding_var_name][0])) - if var_def_dict[embedding_var_name][1] != embedding_var_name_in_fc: - raise ValueError( - 'For embedding var name {}, the embedding name is different, ' - 'got {}; expected {}'.format(embedding_var_name, - embedding_var_name_in_fc, - var_def_dict[embedding_var_name][1])) - else: - var_def_dict[embedding_var_name] = (captured_scope_name, - embedding_var_name_in_fc) - - -def _is_running_on_cpu(): - """Returns True if the current context is CPU model.""" - return tpu_function.get_tpu_context().number_of_shards is None +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.feature_column import * +# used by tests +from tensorflow.python.tpu.feature_column import _is_running_on_cpu +from tensorflow.python.tpu.feature_column import _record_variable_scope_and_name +from tensorflow.python.tpu.feature_column import _TPU_FC_TO_SCOPE +from tensorflow.python.tpu.feature_column import _TPUBaseEmbeddingColumn +from tensorflow.python.tpu.feature_column import _TPUEmbeddingColumn +from tensorflow.python.tpu.feature_column import _TPUSharedEmbeddingColumn +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/functional.py b/tensorflow/contrib/tpu/python/tpu/functional.py index 3d04c64033b5a27b34b5aa77a8753246d35d23aa..9a5759221ed9660200cc213df69961db56f8d490 100644 --- a/tensorflow/contrib/tpu/python/tpu/functional.py +++ b/tensorflow/contrib/tpu/python/tpu/functional.py @@ -1,23 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ============================================================================= -"""Functional operations.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.tpu.python.ops import tpu_ops - -TPUPartitionedCall = tpu_ops.tpu_partitioned_call # pylint: disable=invalid-name +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.functional import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/keras_support.py b/tensorflow/contrib/tpu/python/tpu/keras_support.py index 6ad4e45e9625f191bb4c01f70b434dc2c4fba638..0f95afd6db46809136c47b58346b85fe45d3e763 100644 --- a/tensorflow/contrib/tpu/python/tpu/keras_support.py +++ b/tensorflow/contrib/tpu/python/tpu/keras_support.py @@ -63,6 +63,7 @@ from tensorflow.contrib.tpu.python.tpu import tpu_optimizer from tensorflow.contrib.tpu.python.tpu import tpu_system_metadata as tpu_system_metadata_lib from tensorflow.core.protobuf import config_pb2 from tensorflow.core.protobuf.tpu import compilation_result_pb2 as tpu_compilation_result +from tensorflow.python import tf2 from tensorflow.python.client import session as tf_session from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import iterator_ops @@ -200,13 +201,22 @@ class TPUDistributionStrategy(object): removed in future once the model replication functionality is mature enough. If `False` (default behavior), the system automatically finds the best configuration, in terms of number of TPU cores, for the model - replication, typically using all avaiable TPU cores. If overwrites as + replication, typically using all available TPU cores. If overwrites as `True`, force the model replication using single core, i.e., no replication. Raises: Exception: No TPU Found on the given worker. """ - + if tf2.enabled(): + raise RuntimeError( + 'Keras support is now deprecated in support of TPU Strategy. ' + 'Please follow the distribution strategy guide on tensorflow.org ' + 'to migrate to the 2.0 supported version.') + else: + logging.warning( + 'Keras support is now deprecated in support of TPU Strategy. ' + 'Please follow the distribution strategy guide on tensorflow.org ' + 'to migrate to the 2.0 supported version.') if tpu_cluster_resolver is None: tpu_cluster_resolver = tpu_cluster_resolver_lib.TPUClusterResolver('') @@ -725,9 +735,10 @@ class TPUDatasetInfeedManager(TPUInfeedManager): self._dataset = dataset self._tpu_assignment = tpu_assignment - dummy_x_shape = dataset.output_shapes[0].as_list() + dataset_output_shapes = dataset_ops.get_legacy_output_shapes(dataset) + dummy_x_shape = dataset_output_shapes[0].as_list() dummy_x_shape[0] *= tpu_assignment.num_towers - dummy_y_shape = dataset.output_shapes[1].as_list() + dummy_y_shape = dataset_output_shapes[1].as_list() dummy_y_shape[0] *= tpu_assignment.num_towers self._iterator = dataset_ops.make_initializable_iterator(dataset) K.get_session().run(self._iterator.initializer) @@ -743,23 +754,26 @@ class TPUDatasetInfeedManager(TPUInfeedManager): # Use dummy numpy inputs for the rest of Keras' shape checking. We # intercept them when building the model. + dataset_output_types = dataset_ops.get_legacy_output_types(dataset) self._dummy_x = np.zeros( - dummy_x_shape, dtype=dataset.output_types[0].as_numpy_dtype) + dummy_x_shape, dtype=dataset_output_types[0].as_numpy_dtype) self._dummy_y = np.zeros( - dummy_y_shape, dtype=dataset.output_types[1].as_numpy_dtype) + dummy_y_shape, dtype=dataset_output_types[1].as_numpy_dtype) input_specs = [] - if isinstance(self._iterator.output_shapes, tuple): - assert isinstance(self._iterator.output_types, tuple) - assert len(self._iterator.output_shapes) == len( - self._iterator.output_types) - for i in range(len(self._iterator.output_shapes)): - spec = tensor_spec.TensorSpec(self._iterator.output_shapes[i], - self._iterator.output_types[i]) + iterator_output_shapes = dataset_ops.get_legacy_output_shapes( + self._iterator) + iterator_output_types = dataset_ops.get_legacy_output_types(self._iterator) + if isinstance(iterator_output_shapes, tuple): + assert isinstance(iterator_output_types, tuple) + assert len(iterator_output_shapes) == len(iterator_output_types) + for i in range(len(iterator_output_shapes)): + spec = tensor_spec.TensorSpec(iterator_output_shapes[i], + iterator_output_types[i]) input_specs.append(spec) - elif isinstance(self._iterator.output_shapes, tensor_shape.TensorShape): - spec = tensor_spec.TensorSpec(self._iterator.output_shapes, - self._iterator.output_types) + elif isinstance(iterator_output_shapes, tensor_shape.TensorShape): + spec = tensor_spec.TensorSpec(iterator_output_shapes, + iterator_output_types) input_specs.append(spec) # Pre-process the inputs and get_next_ops before caching. @@ -770,24 +784,26 @@ class TPUDatasetInfeedManager(TPUInfeedManager): def _verify_dataset_shape(self, dataset): """Verifies a dataset is of an appropriate shape for TPUs.""" + dataset_output_shapes = dataset_ops.get_legacy_output_shapes(dataset) + dataset_output_classes = dataset_ops.get_legacy_output_classes(dataset) if not isinstance(dataset, dataset_ops.DatasetV2): raise ValueError('The function passed as the `x` parameter did not ' 'return a `tf.data.Dataset`.') - if not isinstance(dataset.output_classes, tuple): + if not isinstance(dataset_output_classes, tuple): raise ValueError('The dataset must return a tuple of tf.Tensors, ' - 'instead it returns: %s' % dataset.output_classes) - if len(dataset.output_classes) != 2: + 'instead it returns: %s' % dataset_output_classes) + if len(dataset_output_classes) != 2: raise ValueError('The dataset must return a 2-element tuple, got ' - '%s output classes instead.' % (dataset.output_classes,)) - for i, cls in enumerate(dataset.output_classes): + '%s output classes instead.' % (dataset_output_classes,)) + for i, cls in enumerate(dataset_output_classes): if cls != ops.Tensor: raise ValueError('The dataset returned a non-Tensor type (%s) at ' 'index %d.' % (cls, i)) - for i, shape in enumerate(dataset.output_shapes): + for i, shape in enumerate(dataset_output_shapes): if not shape: raise ValueError('The dataset returns a scalar tensor in ' 'tuple index %d. Did you forget to batch? ' - '(Output shapes: %s).' % (i, dataset.output_shapes)) + '(Output shapes: %s).' % (i, dataset_output_shapes)) for j, dim in enumerate(shape): if dim.value is None: if j == 0: @@ -800,7 +816,7 @@ class TPUDatasetInfeedManager(TPUInfeedManager): 'currently requires static shapes. The provided ' 'dataset only has a partially defined shape. ' '(Dimension %d of output tensor %d is not statically known ' - 'for output shapes: %s.%s)' % (j, i, dataset.output_shapes, hint)) + 'for output shapes: %s.%s)' % (j, i, dataset_output_shapes, hint)) @property def dummy_x(self): @@ -1367,7 +1383,16 @@ class KerasTPUModel(models.Model): outputs=cpu_model.outputs, name=cpu_model.name, ) - + if tf2.enabled(): + raise RuntimeError( + 'Keras support is now deprecated in support of TPU Strategy. ' + 'Please follow the distribution strategy guide on tensorflow.org ' + 'to migrate to the 2.0 supported version.') + else: + logging.warning( + 'Keras support is now deprecated in support of TPU Strategy. ' + 'Please follow the distribution strategy guide on tensorflow.org ' + 'to migrate to the 2.0 supported version.') # Create a mapping from numpy arrays to infeed managers. # Note: uses a list of tuples instead of a map because numpy arrays are # not hashable. diff --git a/tensorflow/contrib/tpu/python/tpu/session_support.py b/tensorflow/contrib/tpu/python/tpu/session_support.py index 5cb2ca6478a1d7589cd2aa2d52c82306b3fd11f4..ed8f9525c9b91208d39805654b01837abdbf3a77 100644 --- a/tensorflow/contrib/tpu/python/tpu/session_support.py +++ b/tensorflow/contrib/tpu/python/tpu/session_support.py @@ -1,438 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # -# Licensed under the Apache License, Version 2.0 (the 'License'); +# Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an 'AS IS' BASIS, +# distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ====================================== -"""Operations for handling session logging and shutdown notifications.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import threading - -import time -from google.protobuf import text_format - -from tensorflow.contrib.tpu.python.ops import tpu_ops -from tensorflow.core.protobuf import config_pb2 -from tensorflow.core.util import event_pb2 -from tensorflow.python.client import session as session_lib -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors -from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.platform import tf_logging as logging -from tensorflow.python.training import session_run_hook -from tensorflow.python.training import training_util - -_WATCHDOG = None - - -class CoordinatorShutdownException(Exception): - """Raised when the coordinator needs to shutdown.""" - pass - - -def _clone_session(session, graph=None): - return session_lib.Session( - target=session.sess_str, - config=session._config, # pylint: disable=protected-access - graph=graph if graph else session.graph) - - -def _make_heartbeat_op(session, device, request_ph): - """Return a heartbeat op or None if heartbeats are not supported by device.""" - try: - # Test if we can connect in a isolated graph + session - with ops.Graph().as_default(): - with _clone_session(session) as temp_session: - with ops.device(device): - heartbeat_op = tpu_ops.worker_heartbeat('') - options = config_pb2.RunOptions(timeout_in_ms=5000) - temp_session.run(heartbeat_op, options=options) - except errors.InvalidArgumentError as _: - logging.warning('Error running heartbeat on %s', device) - return None - except errors.DeadlineExceededError as _: - logging.warning('Timeout connecting to %s when testing heartbeat', device) - return None - - # If we successfully connected and pinged the worker, go ahead and construct - # the operation. - with ops.device(device): - return tpu_ops.worker_heartbeat(request_ph) - - -class WorkerHeartbeatManager(object): - """Manages the status/heartbeat monitor for a set of workers.""" - - def __init__(self, session, devices, heartbeat_ops, request_placeholder): - """Construct a new WorkerHeartbeatManager. - - (Prefer using `WorkerHeartbeatManager.from_devices` when possible.) - - Args: - session: `tf.Session`, session to use for heartbeat operations. - devices: `list[string]` Set of devices to connect to. - heartbeat_ops: `list[tf.Operation]` Heartbeat operations. - request_placeholder: `tf.Placeholder[String]` Placeholder used to specify - the WorkerHeartbeatRequest protocol buffer. - """ - self._session = session - self._devices = devices - self._ops = heartbeat_ops - self._request_placeholder = request_placeholder - - @staticmethod - def from_devices(session, devices): - """Construct a heartbeat manager for the given devices.""" - if not devices: - logging.error('Trying to create heartbeat manager with no devices?') - - logging.info('Creating heartbeat manager for %s', devices) - request_placeholder = array_ops.placeholder( - name='worker_heartbeat_request', dtype=dtypes.string) - - heartbeat_ops = [] - kept_devices = [] - for device in devices: - heartbeat_op = _make_heartbeat_op(session, device, request_placeholder) - if heartbeat_op is not None: - kept_devices.append(device) - heartbeat_ops.append(heartbeat_op) - else: - logging.warning('Heartbeat support not available for %s', device) - - return WorkerHeartbeatManager(session, kept_devices, heartbeat_ops, - request_placeholder) - - def num_workers(self): - return len(self._devices) - - def configure(self, message): - """Configure heartbeat manager for all devices. - - Args: - message: `event_pb2.WorkerHeartbeatRequest` - Returns: `None` - """ - logging.info('Configuring worker heartbeat: %s', - text_format.MessageToString(message)) - self._session.run(self._ops, - {self._request_placeholder: message.SerializeToString()}) - - def ping(self, request=None, timeout_in_ms=5000): - """Ping all workers, returning the parsed status results.""" - if request is None: - request = event_pb2.WorkerHeartbeatRequest() - - options = config_pb2.RunOptions(timeout_in_ms=timeout_in_ms) - results = self._session.run( - self._ops, - feed_dict={self._request_placeholder: request.SerializeToString()}, - options=options) - parsed_results = [ - event_pb2.WorkerHeartbeatResponse.FromString(res_pb) - for res_pb in results - ] - logging.debug('Ping results: %s', parsed_results) - return parsed_results - - def lame_workers(self): - """Ping all workers, returning manager containing lame workers (or None).""" - ping_results = self.ping() - lame_workers = [] - - for ping_response, device, op in zip(ping_results, self._devices, - self._ops): - if ping_response.health_status != event_pb2.OK: - lame_workers.append((device, op)) - - if not lame_workers: - return None - - bad_devices, bad_ops = zip(*lame_workers) - return WorkerHeartbeatManager(self._session, bad_devices, bad_ops, - self._request_placeholder) - - def __repr__(self): - return 'HeartbeatManager(%s)' % ','.join(self._devices) - - def shutdown(self, timeout_ms=10000): - """Shutdown all workers after `shutdown_timeout_secs`.""" - logging.info('Shutting down %s.', self) - req = event_pb2.WorkerHeartbeatRequest( - watchdog_config=event_pb2.WatchdogConfig(timeout_ms=timeout_ms), - shutdown_mode=event_pb2.WAIT_FOR_COORDINATOR) - self.configure(req) - - # Wait for workers to shutdown. This isn't strictly required - # but it avoids triggering multiple checkpoints with the same lame worker. - logging.info('Waiting %dms for worker shutdown.', timeout_ms) - time.sleep(timeout_ms / 1000) - - -def all_worker_devices(session): - """Return a list of devices for each worker in the system.""" - devices = session.list_devices() - return [ - device.name - for device in devices - if ':CPU:' in device.name and 'coordinator' not in device.name - ] - - -class WatchdogManager(threading.Thread): - """Configures worker watchdog timer and handles periodic pings. - - Usage: - # Ping workers every minute, shutting down workers if they haven't received - # a ping after 1 hour. - watchdog_manager = WatchdogManager( - ping_interval=60, shutdown_timeout=3600 - ) - - # Use as a context manager, resetting watchdog on context exit: - with watchdog_manager: - session.run(...) - - # Or setup globally; watchdog will remain active until program exit. - watchdog_manager.configure_and_run() - """ - - def __init__(self, - session, - devices=None, - ping_interval=60, - shutdown_timeout=3600): - """Initialize a watchdog manager. - - Args: - session: Session connected to worker devices. A cloned session and graph - will be created for managing worker pings. - devices: Set of devices to monitor. If none, all workers will be - monitored. - ping_interval: Time, in seconds, between watchdog pings. - shutdown_timeout: Time, in seconds, before watchdog timeout. - """ - threading.Thread.__init__(self) - self.ping_interval = ping_interval - self.shutdown_timeout = shutdown_timeout - self.daemon = True - self._config = session._config # pylint: disable=protected-access - self._target = session.sess_str - self._running = False - self._devices = devices - - self._graph = None - self._session = None - self._worker_manager = None - - def _reset_manager(self): - """Reset the graph, session and worker manager.""" - self._graph = ops.Graph() - self._session = session_lib.Session( - target=self._target, - graph=self._graph, - config=self._config, - ) - - if self._devices is None: - self._devices = all_worker_devices(self._session) - - with self._graph.as_default(): - self._worker_manager = WorkerHeartbeatManager.from_devices( - self._session, self._devices) - - self._worker_manager.configure( - event_pb2.WorkerHeartbeatRequest( - watchdog_config=event_pb2.WatchdogConfig( - timeout_ms=self.shutdown_timeout * 1000,), - shutdown_mode=event_pb2.WAIT_FOR_COORDINATOR)) - - def configure_and_run(self): - logging.info( - 'Enabling watchdog timer with %d second timeout ' - 'and %d second ping interval.', self.shutdown_timeout, - self.ping_interval) - self._reset_manager() - self._running = True - self.start() - - def stop(self): - logging.info('Stopping worker watchdog.') - self._worker_manager.configure( - event_pb2.WorkerHeartbeatRequest( - watchdog_config=event_pb2.WatchdogConfig(timeout_ms=-1,), - shutdown_mode=event_pb2.NOT_CONFIGURED)) - self._running = False - self.join() - - def __enter__(self): - self.configure_and_run() - - def __exit__(self, exc_type, exc_val, exc_tb): - self.stop() - - def run(self): - # Don't fetch logs or adjust timing: just ping the watchdog. - # - # If we hit an exception, reset our session as it is likely broken. - while self._running: - try: - self._worker_manager.ping(request=None) - time.sleep(self.ping_interval) - except errors.OpError as e: - # Catch any TF errors that occur so we don't stop sending heartbeats - logging.debug('Caught error while sending heartbeat: %s', e) - self._reset_manager() - - -def start_worker_watchdog(session, - devices=None, - ping_interval=60, - shutdown_timeout=3600): - """Start global worker watchdog to shutdown workers on coordinator exit.""" - global _WATCHDOG - if _WATCHDOG is None: - # Ensure we can send a few pings before we timeout! - ping_interval = min(shutdown_timeout / 10., ping_interval) - _WATCHDOG = WatchdogManager(session, devices, ping_interval, - shutdown_timeout) - _WATCHDOG.configure_and_run() - - -class GracefulShutdownHook(session_run_hook.SessionRunHook): - """Session hook that watches for shutdown events. - - If a shutdown is indicated, `saver.save(checkpoint_prefix)` is executed, and a - SystemShutdown exception is raised to terminate the main session. If `saver` - is None the `SAVERS` collection will be read to find a saver. - - `on_shutdown_hooks` is an optional list of functions that should be called - after checkpointing. The function is called with (`run_context`, - `all_workers`, `lame_workers`). - - If `heartbeat_group` is not specified, it will default to all CPU workers - in the system. - """ - - def __init__(self, checkpoint_prefix, saver=None, on_shutdown_hooks=None): - self._saver = saver - self._checkpoint_prefix = checkpoint_prefix - self._on_shutdown_hooks = on_shutdown_hooks if on_shutdown_hooks else [] - - # Worker heartbeats are managed independently of the main training graph. - self._graph = ops.Graph() - self._workers = None - self._session = None - self._heartbeat_supported = False - - def after_create_session(self, training_session, coord): # pylint: disable=unused-argument - # N.B. We have to pull the global step here to avoid it being unavailable - # at checkpoint time; the graph has been frozen at that point. - if training_util.get_global_step() is None and self.saver() is not None: - raise ValueError( - 'Saver defined but no global step. Run `get_or_create_global_step()`' - ' in your model definition to allow checkpointing.') - - with self._graph.as_default(): - logging.info('Installing graceful shutdown hook.') - self._session = _clone_session(training_session, self._graph) - self._workers = WorkerHeartbeatManager.from_devices( - self._session, all_worker_devices(self._session)) - self._heartbeat_supported = self._workers.num_workers() > 0 - if self._heartbeat_supported: - self._workers.configure( - event_pb2.WorkerHeartbeatRequest( - shutdown_mode=event_pb2.WAIT_FOR_COORDINATOR)) - else: - logging.warn( - 'No workers support hearbeats. Failure handling will be disabled.') - - def saver(self): - if self._saver: - return self._saver - - savers = ops.get_collection(ops.GraphKeys.SAVERS) - if not savers: - return None - - if not isinstance(savers, list): - return savers - - if len(savers) > 1: - logging.error( - 'Multiple savers in the SAVERS collection. On-demand checkpointing ' - 'will be disabled. Pass an explicit `saver` to the constructor to ' - 'override this behavior.') - return None - - return savers[0] - - def after_run(self, run_context, run_values): - del run_values - - if not self._heartbeat_supported: - return - - lame_workers = self._workers.lame_workers() - if lame_workers: - logging.info('ShutdownHook: lame workers found: %s', lame_workers) - - if self.saver(): - logging.info('ShutdownHook: saving checkpoint to %s', - self._checkpoint_prefix) - self.saver().save( - run_context.session, - self._checkpoint_prefix, - global_step=training_util.get_global_step(), - write_state=True, - ) - else: - logging.info('ShutdownHook: no Saver defined.') - - for fn in self._on_shutdown_hooks: - fn(run_context, self._workers, lame_workers) - - -class RestartComputation(object): - """Restart the entire computation. - - This hook shuts down all workers and returns control to the top-level by - throwing a CoordinatorShutdownException. - """ - - def __init__(self, timeout_ms=10000): - self.timeout_ms = timeout_ms - - def __call__(self, run_context, all_workers, lame_workers): - del run_context, lame_workers - all_workers.shutdown(timeout_ms=self.timeout_ms) - - logging.info('Terminating coordinator.') - raise CoordinatorShutdownException() - - -class ShutdownLameWorkers(object): - """Shutdown lamed workers. - - Processing will continue normally (typically by waiting for the down - workers to be restarted). - """ - - def __init__(self, timeout_ms=10000): - self.timeout_in_ms = timeout_ms - - def __call__(self, run_context, all_workers, lame_workers): - lame_workers.shutdown(timeout_ms=self.timeout_in_ms) +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.session_support import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py b/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py index ae0582208450919b79a7c3031c726e24986aa456..73db253fd790f26679fb05bd6e7a5da6a99da1a7 100644 --- a/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py +++ b/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py @@ -1,1638 +1,23 @@ -# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ======================================================================== -"""A utility to trace tensor values on TPU.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import os -import os.path -import re -import sys - -from tensorflow.contrib.tpu.python.ops import tpu_ops -from tensorflow.contrib.tpu.python.tpu import tpu -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import graph_io -from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_util -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import control_flow_ops -from tensorflow.python.ops import control_flow_util -from tensorflow.python.ops import gen_math_ops -from tensorflow.python.ops import init_ops -from tensorflow.python.ops import linalg_ops -from tensorflow.python.ops import logging_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import state_ops -from tensorflow.python.ops import variable_scope -from tensorflow.python.platform import gfile -from tensorflow.python.platform import tf_logging as logging - -_TRACER_LOG_PREFIX = ' [>>>TT>>>]' -_DEVICE_TYPE_TPU = 'tpu' -_DEVICE_TYPE_CPU = 'cpu' -_TRACE_MODE_NAN_INF = 'nan-inf' -_TRACE_MODE_PART_TENSOR = 'part-tensor' -_TRACE_MODE_PART_TENSOR_SIZE = 3 -_TRACE_MODE_FULL_TENSOR = 'full-tensor' -_TRACE_MODE_NORM = 'norm' -_TRACE_MODE_MAX_ABS = 'max-abs' -_SUBMODE_BRIEF = 'brief' -_SUBMODE_DETAILED = 'detailed' -_REASON_OUTSIDE_OP_RANGE = 'not-traced-outside-op-range' -_REASON_UNSAFE_OP = 'not-traced-unsafe-op' -_REASON_WHILELOOP_OP = 'not-traced-special-whileloop-op' -_REASON_UNSAFE_SCALAR = 'not-traced-unsafe-scalar' -_REASON_LESS_INTERESTING_OP = 'not-traced-less-interesting-op' -_REASON_DEVICE_MISMATCH = 'not-traced-device-mismatch' -_REASON_DYNAMIC_SHAPE = 'not-traced-dynamic-shape' -_REASON_SCALAR_GET_TRACED = 'traced-scalar' -_REASON_TENSOR_GET_TRACED = 'traced-tensor' -_REASON_USER_INCLUDED = 'traced-user-included' -_REASON_USER_EXCLUDED = 'not-traced-user-excluded' -_REASON_NOT_EXECUTED = 'not-traced-not-in-exec-path' -_REASON_NON_NUMERIC_TENSOR = 'not-traced-non-numeric-tensor' -_REASON_FEEDS_WHILELOOP_OP = 'not-traced-feeds-special-whileloop-op' -_MARKER_SECTION_BEGIN = '!!!!!!! section-begin:' -_MARKER_SECTION_END = '!!!!!!! section-end:' -_SECTION_NAME_CONFIG = 'configuration' -_SECTION_NAME_REASON = 'reason' -_SECTION_NAME_OP_LIST = 'op-list' -_SECTION_NAME_TENSOR_LIST = 'tensor-list' -_SECTION_NAME_CACHE_INDEX_MAP = 'cache-index-map' -_SECTION_NAME_GRAPH = 'graph' -_FIELD_NAME_VERSION = 'version:' -_FIELD_NAME_DEVICE = 'device:' -_FIELD_NAME_TRACE_MODE = 'trace-mode:' -_FIELD_NAME_SUBMODE = 'submode:' -_FIELD_NAME_NUM_REPLICAS = 'num-replicas:' -_FIELD_NAME_NUM_REPLICAS_PER_HOST = 'num-replicas-per-host:' -_FIELD_NAME_NUM_HOSTS = 'num-hosts:' -_FIELD_NAME_NUM_OPS = 'number-of-ops:' -_FIELD_NAME_NUM_TENSORS = 'number-of-tensors:' -_FIELD_NAME_NUM_CACHE_INDICES = 'number-of-indices:' -_FIELD_NAME_TOPOLOGICAL_SORT_SUCCEED = 'topological-sort-succeed:' -_FLAGS_ENV_VAR = 'TENSOR_TRACER_FLAGS' -_FLAG_SINGLE_QUOTE_PAT = re.compile(r"\s*--([^=]+)='([^']*)'") -_FLAG_DOUBLE_QUOTE_PAT = re.compile(r'\s*--([^=]+)="([^"]*)"') -_FLAG_NO_QUOTE_PAT = re.compile(r'\s*--([^=]+)=(\S*)') -_FLAG_NO_EQUAL_PAT = re.compile(r'\s*--([^=]+)\s*') -_FLAG_NAME_ENABLE = 'enable' -_FLAG_NAME_TRACE_MODE = 'trace_mode' -_FLAG_NAME_USE_COMPACT_TRACE = 'compact_trace' -_FLAG_NAME_SUBMODE = 'submode' -_FLAG_NAME_INCLUDE_LESS_INTERESTING_OPS = 'include_less_interesting_ops' -_FLAG_NAME_EXCLUDED_OPNAMES = 'excluded_opnames' -_FLAG_NAME_EXCLUDED_OPTYPES = 'excluded_optypes' -_FLAG_NAME_INCLUDED_OPNAMES = 'included_opnames' -_FLAG_NAME_INCLUDED_OPTYPES = 'included_optypes' -_FLAG_NAME_TRACE_DIR = 'trace_dir' -_FLAG_NAME_REPORT_FILE = 'report_file' -_FLAG_NAME_USE_TEST_UNDECLARED_OUTPUTS_DIR = 'use_test_undeclared_outputs_dir' -_FLAG_NAME_OP_RANGE = 'op_range' -# Folder to dump the pre (before tensor tracer updates) and post graphs (after -# tensor tracer updates). -_FLAG_DUMP_BEFORE_AFTER_GRAPHS = 'dump_graphs' -_OP_RANGE_PAT = re.compile(r'(\d+):(\d+)') -_OUTPUT_STREAM_ESCAPE = 'file://' -_TEST_UNDECLARED_OUTPUTS_DIR_ENV_VAR = 'TEST_UNDECLARED_OUTPUTS_DIR' -_TENSOR_TRACER_COLLECTION = 'tensor_tracer_variables' -_TENSOR_TRACER_CHECKPOINT = 'tensor_tracer_checkpoint' -_TRACE_FILE_NAME = 'trace.all' -_COMPACT_TRACE_FILE_PREFIX = 'compact_trace.' -_COMPACT_TRACE_ENTRY_INIT_VALUE = -1.0 -_TENSOR_TRACER_STORAGE = 'tensor_tracer_storage' -_TENSOR_VALUES_CACHE = 'tensor_values_cache' -_REPLICA_ID_TAG = '#replica-id: ' - - -def tensor_tracepoint(tensor, checkpoint_name): - """Adds a checkpoint with the given checkpoint name for the given tensor. - - The tensor will be added to the list of tensors that will be traced by the - tensor tracer. - - Args: - tensor: the tensor object for which the tracing is requested. - checkpoint_name: a string name for the checkpoint. This name has to be a - unique name if used within model comparison. The tensors that have the same - checkpoint identifier is compared in model comparison. - Returns: - The provided tensor. - """ - - tensor.graph.get_collection(_TENSOR_TRACER_COLLECTION) - tensor.graph.add_to_collection(_TENSOR_TRACER_COLLECTION, - (tensor, checkpoint_name)) - return tensor - - -def keras_layer_tracepoint(layer, checkpoint_name): - """An interface for adding the tensor outputs of a keras layer. - - Encapsulates tensor_tracepoint. - - Args: - layer: A keras layer. - checkpoint_name: a string name for the checkpoint. This name has to be a - unique name if used within model comparison. The tensors that have the same - checkpoint identifier is compared in model comparison. - - Returns: - The provided layer. - """ - try: - outputs = layer.output - if tensor_util.is_tensor(outputs): - tensor_tracepoint(outputs, '%s' % (checkpoint_name)) - else: - idx = 0 - for output_tensor in outputs: - if tensor_util.is_tensor(outputs): - tensor_tracepoint(output_tensor, '%s_%d' % (checkpoint_name, idx)) - idx += 1 - except AttributeError: - pass - except RuntimeError: - pass - return layer - - -def _trace_files_need_precreated(output_dir): - """Return True if trace files must be pre-created by users.""" - - if not output_dir.startswith('/'): - return False - if len(output_dir) < 5: - return False - if output_dir[2] != 'n': - return False - if output_dir[3] != 's': - return False - if output_dir[1] != 'c': - return False - if output_dir[4] != '/': - return False - return True - - -def _get_tensor_values_cache(graph=None): - """Returns the variable that implements tensor-value caching.""" - - graph = graph or ops.get_default_graph() - collection = graph.get_collection(_TENSOR_TRACER_STORAGE) - if len(collection) == 1: - return collection[0] - elif not collection: - raise RuntimeError('%s has not been created'%_TENSOR_VALUES_CACHE) - else: - raise RuntimeError('Multiple %s created'%_TENSOR_VALUES_CACHE) - return None - - -def _create_tensor_values_cache(graph, num_tensors): - """Creates a variable as the cache to store intermediate tensor values.""" - graph = graph or ops.get_default_graph() - # Create in proper graph and base name_scope. - with graph.as_default() as g, g.name_scope(None): - return variable_scope.get_variable( - _TENSOR_VALUES_CACHE, - shape=[num_tensors], - dtype=dtypes.float32, - initializer=init_ops.constant_initializer( - _COMPACT_TRACE_ENTRY_INIT_VALUE), - trainable=False, - use_resource=True, - collections=[_TENSOR_TRACER_STORAGE, ops.GraphKeys.GLOBAL_VARIABLES]) - - -class TensorTracer(object): - """A software construct for tracing tensor values in a TF graph on TPU. - - This utility is disabled by default. It can be enabled by setting - the TENSOR_TRACER_FLAGS env variable as: - export TENSOR_TRACER_FLAGS="--enable=1" - If it is enabled, it will trace the output tensor values of - selected Ops in the graph. It has two outputs: (1) the traces and (2) - a report. The traces are dumped to a specified local file on the TPU - host. The report is printed to the log.info of the TPU job. - By passing options via the env variable, users can change: - (1) the trace mode (e.g., detecting NaN/Inf, printing partial or - full tensor values) - (2) which Ops to be traced (via op.name or op.type) - (3) output trace file path. - """ - # The set of graphs that are rewritten by tensor tracer. - _traced_graphs = set() - @staticmethod - def _match_next_flag(flags, pos): - """Returns the match for the next TensorTracer flag. - - Args: - flags: a string that contains the flags. - pos: where in flags to start the search. - - Returns: - A pair where the first element is the regular-expression - match found and the second element indicates if the match - has a value. - """ - - match = _FLAG_DOUBLE_QUOTE_PAT.match(flags, pos) - if match: - return match, True - match = _FLAG_SINGLE_QUOTE_PAT.match(flags, pos) - if match: - return match, True - match = _FLAG_NO_QUOTE_PAT.match(flags, pos) - if match: - return match, True - match = _FLAG_NO_EQUAL_PAT.match(flags, pos) - if match: - # The flag is found but is not given a value. - return match, False - # The flag is not found. - return None, False - - @staticmethod - def validate_flag_names(): - """Validates if the TensorTrace flags passed are valid.""" - valid_flag_names = [_FLAG_NAME_ENABLE, _FLAG_NAME_TRACE_MODE, - _FLAG_NAME_USE_COMPACT_TRACE, - _FLAG_NAME_SUBMODE, - _FLAG_NAME_EXCLUDED_OPNAMES, - _FLAG_NAME_EXCLUDED_OPTYPES, - _FLAG_NAME_INCLUDED_OPNAMES, - _FLAG_NAME_INCLUDED_OPTYPES, - _FLAG_NAME_TRACE_DIR, - _FLAG_NAME_REPORT_FILE, - _FLAG_NAME_USE_TEST_UNDECLARED_OUTPUTS_DIR, - _FLAG_NAME_INCLUDE_LESS_INTERESTING_OPS, - _FLAG_NAME_OP_RANGE, - _FLAG_DUMP_BEFORE_AFTER_GRAPHS] - tensor_tracer_flags = os.environ.get(_FLAGS_ENV_VAR) - if not tensor_tracer_flags: - return - pos = 0 - while True: - match, _ = TensorTracer._match_next_flag(tensor_tracer_flags, pos) - if not match: - break - flag_name = match.group(1) - if flag_name not in valid_flag_names: - raise ValueError( - 'The flag name "%s" passed via the environment variable "%s" ' - 'is invalid. Valid flag names are:' - '\n%s'%(flag_name, _FLAGS_ENV_VAR, valid_flag_names)) - pos = match.end() - - @staticmethod - def print_flag_values(): - """Prints all TensorTracer flags passed via environment variables.""" - - tensor_tracer_flags = os.environ.get(_FLAGS_ENV_VAR) - if not tensor_tracer_flags: - return 'Env variable "%s" is not set'%_FLAGS_ENV_VAR - result = 'Env variable "%s" is set to "%s"\n'%(_FLAGS_ENV_VAR, - tensor_tracer_flags) - result += 'Individual flag value:\n' - pos = 0 - while True: - match, has_value = TensorTracer._match_next_flag( - tensor_tracer_flags, pos) - if not match: - break - flag_name = match.group(1) - if has_value: - flag_value = match.group(2) - else: - flag_value = None - result += ' %s: %s\n'%(flag_name, flag_value) - pos = match.end() - result += '\n' - return result - - @staticmethod - def get_flag_value(wanted_flag_name): - """Returns the value of a TensorTracer flags. - - Args: - wanted_flag_name: the name the the flag we are looking for. - - Returns: - A pair where the first element indicates if the flag is - found and the second element is the value of the flag. - - Raises: - RuntimeError: If supposedly deadcode is reached. - """ - - tensor_tracer_flags = os.getenv(_FLAGS_ENV_VAR) - if not tensor_tracer_flags: - return False, None - pos = 0 - while True: - match, has_value = TensorTracer._match_next_flag( - tensor_tracer_flags, pos) - if not match: - return False, None - flag_name = match.group(1) - if has_value: - flag_value = match.group(2) - else: - flag_value = None - if flag_name == wanted_flag_name: - return True, flag_value - pos = match.end() - raise RuntimeError('Should not reach here.') - - @staticmethod - def flag_value_to_re_list(flag_name): - """Converts list of strings to compiled RE.""" - - re_list = [] - found, flag_value = TensorTracer.get_flag_value(flag_name) - if not found or not flag_value: - return re_list - list_of_values = flag_value.split() - for v in list_of_values: - r = re.compile(v) - re_list.append(r) - return re_list - - @staticmethod - def _is_flag_on(flag_name): - """Returns True if the given flag is on.""" - - found, flag_value = TensorTracer.get_flag_value(flag_name) - if not found: - return False - if flag_value is None: - return True - # Depends on the flag value. - flag_value = flag_value.lower() - enabled = flag_value in ['1', 't', 'true', 'y', 'yes'] - return enabled - - @staticmethod - def is_enabled(): - """Returns True if TensorTracer is enabled.""" - - return TensorTracer._is_flag_on(_FLAG_NAME_ENABLE) - - @staticmethod - def use_test_undeclared_outputs_dir(): - """Decides the output directory of the report and trace files. - - Args: - None. - - Returns: - True if the output files should be written to the - test-undeclared-outputs-directory defined via an - env variable. - """ - - return TensorTracer._is_flag_on( - _FLAG_NAME_USE_TEST_UNDECLARED_OUTPUTS_DIR) - - @staticmethod - def use_compact_trace(): - return TensorTracer._is_flag_on( - _FLAG_NAME_USE_COMPACT_TRACE) - - @staticmethod - def check_device_type(device_type): - """Checks if the given device type is valid.""" - - if device_type not in [_DEVICE_TYPE_TPU, _DEVICE_TYPE_CPU]: - raise ValueError('Invalid device_type "%s"'%device_type) - - @staticmethod - def check_trace_mode(trace_mode): - """Checks if the given trace mode is valid.""" - - valid_trace_modes = [_TRACE_MODE_NAN_INF, _TRACE_MODE_PART_TENSOR, - _TRACE_MODE_FULL_TENSOR, _TRACE_MODE_NORM, - _TRACE_MODE_MAX_ABS] - if trace_mode not in valid_trace_modes: - raise ValueError('Invalid trace mode "%s" given to the Tensor_Tracer.' - 'Valid trace modes are: %s'%(trace_mode, - valid_trace_modes)) - - @staticmethod - def check_submode(submode): - """Checks if the given submode is valid.""" - - if not submode: - return - valid_submodes = [_SUBMODE_DETAILED, _SUBMODE_BRIEF] - if submode not in valid_submodes: - raise ValueError('Invalid submode "%s" given to the Tensor_Tracer.' - 'Valid submodes are: %s'%(submode, - valid_submodes)) - - @staticmethod - def loop_cond_op(op): - return op.type in ('LoopCond', 'RefLoopCond') - - @staticmethod - def while_loop_op(op): - """Returns true if op is one of the special ops of in a while loop. - - Args: - op: A tf.Operation. - - Returns: - True if the given op is one of [Switch, Merge, Enter, Exit, - NextIteration, LoopCond], which are all building blocks for TF while - loops. - """ - return (control_flow_util.IsLoopSwitch(op) or - control_flow_util.IsLoopMerge(op) or - control_flow_util.IsLoopEnter(op) or - control_flow_util.IsLoopExit(op) or - TensorTracer.loop_cond_op(op) or - op.type in ('RefNextIteration', 'NextIteration')) - - @staticmethod - def unsafe_op(op): - """Returns True if this op is not safe to be traced.""" - - if control_flow_util.IsInCond(op): - return True - # Reasons for not including following op types: - # Assign: cause incorrect result with CPU tracing. - if op.type in ['Assign']: - return True - return False - - @staticmethod - def device_mismatch(device_type, op): - if device_type == _DEVICE_TYPE_TPU: - # pylint: disable=protected-access - return tpu._TPU_REPLICATE_ATTR not in op.node_def.attr - # pylint: enable=protected-access - return False - - @staticmethod - def unsafe_scalar_trace(op): - """Return true if scalar output tensor from Op is not safe to be traced.""" - - # Tracing the following causes cycle in the graph on TPU. - if op.type in ['LoopCond', 'Enter', 'Merge', 'Const', - 'Switch', 'Less', 'ReadVariableOp']: - return True - # Tracing the following will cause casting-issue - # with the norm tracing mode or other compilation issues on CPU. - if op.type in ['VarHandleOp', 'IteratorToStringHandle', - 'IteratorGetNext', 'OneShotIterator', - 'IteratorV2', 'MakeIterator', - 'BatchDatasetV2', 'MapDataset', - 'FixedLengthRecordDataset', 'TakeDataset', 'ZipDataset', - 'Placeholder', 'PlaceholderWithDefault', 'StridedSlice']: - return True - return False - - @staticmethod - def less_interesting_op(op): - """Returns True if the given Op is not an interesting one to be traced.""" - - found, _ = TensorTracer.get_flag_value( - _FLAG_NAME_INCLUDE_LESS_INTERESTING_OPS) - if found: - # users force to include all ops. - return False - # Following ops are highly unlikey to cause bugs. - return op.type in ['Const', 'Identity', 'Cast', 'Shape'] - - @staticmethod - def reason(op_idx, details): - """Returns reason why the Op at op_idx is traced or not.""" - - return '%d %s'%(op_idx, details) - - @staticmethod - def topological_sort(g): - """Performs topological sort on the given graph. - - Args: - g: the graph. - - Returns: - A pair where the first element indicates if the topological - sort succeeded (True if there is no cycle found; False if a - cycle is found) and the second element is either the sorted - list of nodes or the cycle of nodes found. - """ - - def visit(op, cycle, permanently_marked_ops, - temporarily_marked_ops, sorted_ops): - """Recursively visits all Ops in a graph. - - Args: - op: the current Op being visited. - cycle: a cycle of Ops found. - permanently_marked_ops: the set of Ops that were already visited. - temporarily_marked_ops: the set of Ops that we have visited during - the current descent. - sorted_ops: the list of Ops sorted in topological order. - """ - - if cycle: - return - if op in permanently_marked_ops: - return - if op in temporarily_marked_ops: - cycle = temporarily_marked_ops - return - temporarily_marked_ops.add(op) - for i in range(len(op.outputs)): - out_tensor = op.outputs[i] - for consumer_op in out_tensor.consumers(): - visit(consumer_op, cycle, permanently_marked_ops, - temporarily_marked_ops, sorted_ops) - # pylint: disable=protected-access - for ctrl_output_op in op._control_outputs: - # pylint: enable=protected-access - visit(ctrl_output_op, cycle, permanently_marked_ops, - temporarily_marked_ops, sorted_ops) - temporarily_marked_ops.remove(op) - permanently_marked_ops.add(op) - sorted_ops.insert(0, op) - - graph_cycle = set([]) - sorted_ops = [] - permanently_marked_ops = set([]) - temporarily_marked_ops = set([]) - unsorted_ops = g.get_operations() - for op in unsorted_ops: - visit(op, graph_cycle, permanently_marked_ops, - temporarily_marked_ops, sorted_ops) - if graph_cycle: - return (False, graph_cycle) - else: - assert len(unsorted_ops) == len(sorted_ops) - return (True, sorted_ops) - - @staticmethod - def _make_op_and_tensor_maps(op_list): - """Creates various maps and lists from op_list. - - Args: - op_list: a list of Ops - - Returns: - opname_idx_map: a map from Op's name to its index in op_list. - tensor_list: a list of output tensors of the Ops in op_list. - tensorname_idx_map: a map from output tensor name to its index - in tensor_list. - """ - - opname_idx_map = {} - tensor_list = [] - tensorname_idx_map = {} - for op_id, op in enumerate(op_list): - if op.name in opname_idx_map: - raise ValueError('Duplicated Op name: %s'%op.name) - opname_idx_map[op.name] = op_id - for output_tensor in op.outputs: - if output_tensor.name not in tensorname_idx_map: - tensor_list.append(output_tensor) - tensorname_idx_map[output_tensor.name] = len(tensor_list)-1 - return (opname_idx_map, tensor_list, tensorname_idx_map) - - def __init__(self): - """Initializes a TensorTracer. - - Sets the various member fields from the flags (if given) or the defaults. - """ - self._version = 'use-outside-compilation' - self._device_type = None - TensorTracer.validate_flag_names() - found, self._trace_mode = TensorTracer.get_flag_value(_FLAG_NAME_TRACE_MODE) - if not found or not self._trace_mode: - self._trace_mode = _TRACE_MODE_NAN_INF - TensorTracer.check_trace_mode(self._trace_mode) - found, self._submode = TensorTracer.get_flag_value(_FLAG_NAME_SUBMODE) - if not found or not self._submode: - self._submode = _SUBMODE_DETAILED - TensorTracer.check_submode(self._submode) - self._part_tensor_size = _TRACE_MODE_PART_TENSOR_SIZE - self._instrument_records = {} - self._set_trace_dir() - self._set_report_file() - self._set_op_range() - self._set_excluded_opnames() - self._set_excluded_optypes() - self._set_included_opnames() - self._set_included_optypes() - self._num_replicas = None - self._num_replicas_per_host = None - self._num_hosts = None - self._replica_id = None - _, self._graph_dump_path = TensorTracer.get_flag_value( - _FLAG_DUMP_BEFORE_AFTER_GRAPHS) - - def _add_replica_id_to_graph(self): - """Adds nodes for computing the replica ID to the graph.""" - - if self._num_replicas: - with ops.control_dependencies(None): - # Uses None as dependency to run outside of TPU graph rewrites. - self._replica_id = tpu_ops.tpu_replicated_input( - list(range(self._num_replicas)), - name='tt_replica_id') - else: - self._replica_id = 'unknown' - - def _set_trace_dir(self): - found, self._trace_dir = TensorTracer.get_flag_value(_FLAG_NAME_TRACE_DIR) - if found and self._trace_dir \ - and TensorTracer.use_test_undeclared_outputs_dir(): - raise ValueError('Cannot not use --%s and --%s at the same time' - %(_FLAG_NAME_TRACE_DIR, - _FLAG_NAME_USE_TEST_UNDECLARED_OUTPUTS_DIR)) - if TensorTracer.use_test_undeclared_outputs_dir(): - self._trace_dir = os.environ.get(_TEST_UNDECLARED_OUTPUTS_DIR_ENV_VAR) - - def _set_report_file(self): - """Sets the path of the output report file.""" - - found, self._report_file_path = TensorTracer.get_flag_value( - _FLAG_NAME_REPORT_FILE) - if found and self._report_file_path \ - and TensorTracer.use_test_undeclared_outputs_dir(): - if os.path.isabs(self._report_file_path): - raise ValueError('If use_test_undeclared_outputs_dir is set,' - 'report_file_path cannot be an absolute path (%s)' - %self._report_file_path) - outputs_dir = os.environ.get(_TEST_UNDECLARED_OUTPUTS_DIR_ENV_VAR) - self._report_file_path = os.path.join(outputs_dir, - self._report_file_path) - if not self._report_file_path: - self._report_file = None - return - try: - self._report_file = gfile.Open(self._report_file_path, 'w') - except IOError as e: - raise e - - def _close_report_file(self): - if self._report_file: - self._report_file.close() - - def _set_op_range(self): - """Sets the index range of the Ops that we will consider tracing.""" - - found, op_range = TensorTracer.get_flag_value(_FLAG_NAME_OP_RANGE) - if not found or not op_range: - self._op_range = (-1, -1) # this means including all ops. - return - match = _OP_RANGE_PAT.match(op_range) - if not match: - self._op_range = (-1, -1) # this means including all ops. - return - self._op_range = (int(match.group(1)), int(match.group(2))) - - def _inside_op_range(self, idx): - """Return True if the given index is inside the selected range.""" - - if idx < self._op_range[0]: - return False - return self._op_range[1] < 0 or idx <= self._op_range[1] - - def _set_excluded_opnames(self): - self._excluded_opname_re_list = TensorTracer.flag_value_to_re_list( - _FLAG_NAME_EXCLUDED_OPNAMES) - - def _set_excluded_optypes(self): - self._excluded_optype_re_list = TensorTracer.flag_value_to_re_list( - _FLAG_NAME_EXCLUDED_OPTYPES) - - def _set_included_opnames(self): - self._included_opname_re_list = TensorTracer.flag_value_to_re_list( - _FLAG_NAME_INCLUDED_OPNAMES) - - def _set_included_optypes(self): - self._included_optype_re_list = TensorTracer.flag_value_to_re_list( - _FLAG_NAME_INCLUDED_OPTYPES) - - def _is_user_included_op(self, op): - for opname_re in self._included_opname_re_list: - if opname_re.match(op.name): - return True - for optype_re in self._included_optype_re_list: - if optype_re.match(op.type): - return True - return False - - def _is_user_excluded_op(self, op): - for opname_re in self._excluded_opname_re_list: - if opname_re.match(op.name): - return True - for optype_re in self._excluded_optype_re_list: - if optype_re.match(op.type): - return True - return False - - def _use_tensor_values_cache(self): - """Returns True if immediate tensors should be first saved to a cache.""" - - if self._trace_mode not in set([_TRACE_MODE_NAN_INF, - _TRACE_MODE_NORM, _TRACE_MODE_MAX_ABS]): - return False - if self._trace_dir and _trace_files_need_precreated(self._trace_dir): - return True - if TensorTracer.use_compact_trace(): - return True - return False - - def _save_tensor_value_to_cache_op(self, graph, cache_idx, updates): - """Returns an Op that will save the given updates to an entry in the cache.""" - - cache = _get_tensor_values_cache(graph) - indices = constant_op.constant([cache_idx]) - return state_ops.scatter_update(cache, indices, updates).op - - def _write_report(self, content): - """Writes the given content to the report.""" - - line = '%s %s'%(_TRACER_LOG_PREFIX, content) - if self._report_file: - self._report_file.write(line) - else: - logging.info(line) - - def _write_config_section(self): - """Writes the config section of the report.""" - - self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, _SECTION_NAME_CONFIG)) - self._write_report('%s %s\n'%(_FIELD_NAME_VERSION, self._version)) - self._write_report('%s %s\n'%(_FIELD_NAME_DEVICE, self._device_type)) - self._write_report('%s %s\n'%(_FIELD_NAME_TRACE_MODE, self._trace_mode)) - self._write_report('%s %s\n'%(_FIELD_NAME_SUBMODE, self._submode)) - self._write_report('%s %s\n'%(_FIELD_NAME_NUM_REPLICAS, self._num_replicas)) - self._write_report('%s %s\n'%(_FIELD_NAME_NUM_REPLICAS_PER_HOST, - self._num_replicas_per_host)) - self._write_report('%s %s\n'%(_FIELD_NAME_NUM_HOSTS, self._num_hosts)) - self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_CONFIG)) - - def _write_reason_section(self): - """Writes the reason section of the report.""" - - self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, _SECTION_NAME_REASON)) - for key in sorted(self._instrument_records): - self._write_report('"%s" %s\n'%(key, self._instrument_records[key])) - self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_REASON)) - - def _write_op_list_section(self, op_list): - """Writes the Op-list section of the report.""" - - self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, _SECTION_NAME_OP_LIST)) - self._write_report('%s %d\n'%(_FIELD_NAME_NUM_OPS, len(op_list))) - for i in range(0, len(op_list)): - op = op_list[i] - line = '%d "%s" %s'%(i, op.name, op.type) - for out_tensor in op.outputs: - if out_tensor.name not in self._tensorname_idx_map: - raise ValueError( - 'out_tensor %s is not in tensorname_idx_map'%out_tensor.name) - line += ' %d'%self._tensorname_idx_map[out_tensor.name] - line += '\n' - self._write_report(line) - self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_OP_LIST)) - - def _write_tensor_list_section(self, tensor_list, opname_idx_map): - """Writes the tensor-list section of the report.""" - - self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, - _SECTION_NAME_TENSOR_LIST)) - self._write_report('%s %d\n'%(_FIELD_NAME_NUM_TENSORS, len(tensor_list))) - for i in range(0, len(tensor_list)): - tensor = tensor_list[i] - line = '%d "%s"'%(i, tensor.name) - for consumer_op in tensor.consumers(): - if consumer_op.name not in opname_idx_map: - raise ValueError( - 'consumer_op %s is not in opname_idx_map'%consumer_op.name) - line += ' %d'%opname_idx_map[consumer_op.name] - line += '\n' - self._write_report(line) - self._write_report('%s %s\n'%(_MARKER_SECTION_END, - _SECTION_NAME_TENSOR_LIST)) - - def _write_cache_index_map_section(self): - """Writes the mapping from cache index to tensor index to the report.""" - - self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, - _SECTION_NAME_CACHE_INDEX_MAP)) - self._write_report('%s %d\n'%(_FIELD_NAME_NUM_CACHE_INDICES, - len(self._cache_idx_to_tensor_idx))) - for cache_idx in range(0, len(self._cache_idx_to_tensor_idx)): - tensor_idx = self._cache_idx_to_tensor_idx[cache_idx] - line = '%d %d\n'%(cache_idx, tensor_idx) - self._write_report(line) - self._write_report('%s %s\n'%(_MARKER_SECTION_END, - _SECTION_NAME_CACHE_INDEX_MAP)) - - def _write_graph_section(self, succeed, sorted_or_cycle): - """Writes the graph section of the report.""" - - self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, _SECTION_NAME_GRAPH)) - self._write_report('%s %s\n'%(_FIELD_NAME_TOPOLOGICAL_SORT_SUCCEED, - succeed)) - l = list(sorted_or_cycle) - for i in range(0, len(l)): - self._write_report('%d "%s"\n'%(i, l[i].name)) - self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_GRAPH)) - - def _preprocess_traced_tensor(self, tensor): - """Computes NAN/Norm/Max on TPUs before sending to CPU. - - Args: - tensor: The tensor to be traced. - Returns: - A tensor that should be input to the trace_function. - Raises: - RuntimeError: If the trace mode is invalid. - """ - - def _detect_nan_inf(tensor): - """Trace function for detecting any NaN/Inf in the tensor.""" - - if tensor.dtype.is_floating: - mask = math_ops.reduce_any( - gen_math_ops.logical_or( - gen_math_ops.is_nan(tensor), gen_math_ops.is_inf(tensor))) - output_tensor = control_flow_ops.cond(mask, - lambda: constant_op.constant(1.0), - lambda: constant_op.constant(0.0)) - else: - output_tensor = constant_op.constant(0.0) - # The shape has to be 1. Set it if it does not have the information. - output_tensor = array_ops.reshape(output_tensor, [1]) - return output_tensor - - def _show_norm(tensor): - tensor = math_ops.cast(tensor, dtypes.float32) - output_tensor = linalg_ops.norm(tensor) - # The shape has to be 1. Set it if it does not have the information. - output_tensor = array_ops.reshape(output_tensor, [1]) - return output_tensor - - def _show_max_abs(tensor): - tensor = math_ops.cast(tensor, dtypes.float32) - output_tensor = math_ops.reduce_max(math_ops.abs(tensor)) - zero = constant_op.constant(0, dtypes.float32) - output_tensor = gen_math_ops.maximum(zero, output_tensor) - # The shape has to be 1. Set it if it does not have the information. - output_tensor = array_ops.reshape(output_tensor, [1]) - return output_tensor - - if self._trace_mode == _TRACE_MODE_NAN_INF: - return _detect_nan_inf(tensor) - if self._trace_mode == _TRACE_MODE_PART_TENSOR: - return tensor - if self._trace_mode == _TRACE_MODE_FULL_TENSOR: - return tensor - if self._trace_mode == _TRACE_MODE_NORM: - return _show_norm(tensor) - if self._trace_mode == _TRACE_MODE_MAX_ABS: - return _show_max_abs(tensor) - raise RuntimeError( - 'Tensor trace fun for %s is not yet implemented' % self._trace_mode) - - def _make_tensor_trace_fun(self, tensor_name): - """Makes the tensor tracing function called by outside compilation. - - Args: - tensor_name: name of the tensor being traced. - - Returns: - A function to be passed as the first argument to outside compilation. - - Raises: - RuntimeError: If the trace mode is invalid. - """ - - def _print_tensor(tensor_name, num_elements, tensor, output_tensor): - """Prints a tensor value to a file. - - Args: - tensor_name: name of the tensor being traced. - num_elements: number of elements to print (-1 means print all). - tensor: the tensor needs to be returned. - output_tensor: the tensor needs to be printed. - - Returns: - The same tensor passed via the "tensor" argument. - - Raises: - ValueError: If tensor_name is not already in - self._tensorname_idx_map. - """ - - if self._submode == _SUBMODE_BRIEF: - if tensor_name not in self._tensorname_idx_map: - raise ValueError( - 'Tensor name %s is not in the tensorname_idx_map'%tensor_name) - msg = '%d'%self._tensorname_idx_map[tensor_name] - else: - msg = '"%s"'%tensor_name - - if self._trace_dir: - output_path = os.path.join(self._trace_dir, _TRACE_FILE_NAME) - output_stream = _OUTPUT_STREAM_ESCAPE + output_path - else: - output_stream = sys.stderr - return logging_ops.print_v2(msg, array_ops.shape(output_tensor), - '@', self._replica_id, - '\n', output_tensor, '\n', - summarize=num_elements, - output_stream=output_stream) - - def _show_part_tensor(tensor): - """Trace function for printing part of the tensor.""" - - return _print_tensor(tensor_name, self._part_tensor_size, - tensor, tensor) - - def _show_full_tensor(tensor): - """Trace function for printing the entire tensor.""" - - return _print_tensor(tensor_name, -1, tensor, tensor) - - if self._trace_mode == _TRACE_MODE_PART_TENSOR: - return _show_part_tensor - # The input tensor has a shape of "[1]" for _TRACE_MODE_NAN_INF, - # _TRACE_MODE_NORM, and _TRACE_MODE_MAX_ABS, as related computations are - # performed within TPUs and only their results are transferred to CPU. - # Simply, print the full tensor for these trace modes. - if self._trace_mode in [ - _TRACE_MODE_NAN_INF, _TRACE_MODE_NORM, _TRACE_MODE_FULL_TENSOR, - _TRACE_MODE_MAX_ABS - ]: - return _show_full_tensor - - raise RuntimeError('Tensor trace fun for %s is not yet implemented' - %self._trace_mode) - - def _skip_op(self, op_id, op, user_included, user_excluded, - in_exec_path=True): - """Returns True if we should not trace Op.""" - - if TensorTracer.while_loop_op(op): - self._instrument_records[op.name] = TensorTracer.reason( - op_id, _REASON_WHILELOOP_OP) - return True - if TensorTracer.unsafe_op(op): - self._instrument_records[op.name] = TensorTracer.reason( - op_id, _REASON_UNSAFE_OP) - return True - if TensorTracer.device_mismatch(self._device_type, op): - self._instrument_records[op.name] = TensorTracer.reason( - op_id, _REASON_DEVICE_MISMATCH) - return True - if not in_exec_path: - self._instrument_records[op.name] = TensorTracer.reason( - op_id, _REASON_NOT_EXECUTED) - return True - - if not self._inside_op_range(op_id): - self._instrument_records[op.name] = TensorTracer.reason( - op_id, _REASON_OUTSIDE_OP_RANGE) - return True - if TensorTracer.less_interesting_op(op): - self._instrument_records[op.name] = TensorTracer.reason( - op_id, _REASON_LESS_INTERESTING_OP) - return True - if user_included: - self._instrument_records[op.name] = TensorTracer.reason( - op_id, _REASON_USER_INCLUDED) - return False - if user_excluded: - self._instrument_records[op.name] = TensorTracer.reason( - op_id, _REASON_USER_EXCLUDED) - return True - return False - - def _skip_tensor(self, op_id, out_tensor, user_included, - user_excluded): - """Returns True if we should not trace out_tensor.""" - - # Skips a tensor if the tensor has a non-numeric type. - # Note: we cannot use check_ops.is_numeric_tensor(out_tensor) - # because it also excludes tensors with dtypes, bool, and - # float32_ref, which we actually want to trace. - non_numeric_tensor_types = set([dtypes.variant, dtypes.resource, - dtypes.string]) - if out_tensor.dtype in non_numeric_tensor_types: - self._instrument_records[out_tensor.name] = TensorTracer.reason( - op_id, _REASON_NON_NUMERIC_TENSOR) - return True - # Skip a tensor if it feeds a special while loop op. - if [consumer for consumer in out_tensor.consumers() if - TensorTracer.while_loop_op(consumer)]: - self._instrument_records[out_tensor.name] = TensorTracer.reason( - op_id, _REASON_FEEDS_WHILELOOP_OP) - return True - if user_included: - self._instrument_records[out_tensor.name] = TensorTracer.reason( - op_id, _REASON_USER_INCLUDED) - return False - if user_excluded: - self._instrument_records[out_tensor.name] = TensorTracer.reason( - op_id, _REASON_USER_EXCLUDED) - return True - if not out_tensor.get_shape().is_fully_defined(): - # If trace mode is nan-inf, norm or max, then the tensor will be reduced - # to a scalar before the outside compilation call. - if self._trace_mode in [ - _TRACE_MODE_NAN_INF, _TRACE_MODE_NORM, _TRACE_MODE_MAX_ABS - ]: - self._instrument_records[out_tensor.name] = TensorTracer.reason( - op_id, _REASON_TENSOR_GET_TRACED) - return False - else: - self._instrument_records[out_tensor.name] = TensorTracer.reason( - op_id, _REASON_DYNAMIC_SHAPE) - return True - rank = len(out_tensor.shape) - if rank < 1: - # scalar - if TensorTracer.unsafe_scalar_trace(out_tensor.op): - self._instrument_records[out_tensor.name] = TensorTracer.reason( - op_id, _REASON_UNSAFE_SCALAR) - return True - else: - self._instrument_records[out_tensor.name] = TensorTracer.reason( - op_id, _REASON_SCALAR_GET_TRACED) - return False - else: - # tensor - self._instrument_records[out_tensor.name] = TensorTracer.reason( - op_id, _REASON_TENSOR_GET_TRACED) - return False - - def _filter_execution_path_operations(self, operations, fetches): - """Returns the set of ops in the execution path to compute given fetches.""" - - # If no fetch provided, then return all operations. - if fetches is None: - return set(operations) - # Convert to list, if a single element is provided. - if not isinstance(fetches, (list, tuple)): - fetches = [fetches] - # If a tensor is given as fetch, convert it to op. - op_fetches = [] - for fetch in fetches: - if isinstance(fetch, ops.Operation): - op_fetches.append(fetch) - elif isinstance(fetch, ops.Tensor): - op_fetches.append(fetch.op) - else: - raise RuntimeError('Given fetch:%s is neither a tensor nor an op.' - %fetch) - - execution_path_operations = set(op_fetches) - traverse_stack = list(op_fetches) - while True: - if not traverse_stack: - break - head_op = traverse_stack.pop() - input_ops = [tensor_input.op for tensor_input in head_op.inputs] - input_ops.extend(head_op.control_inputs) - - for input_op in input_ops: - if input_op not in execution_path_operations: - # Filter out loop condition operations, tracing them causes a cycle. - # Trace only the loop-body. - if TensorTracer.loop_cond_op(input_op): - continue - execution_path_operations.add(input_op) - traverse_stack.append(input_op) - return execution_path_operations - - def _determine_traced_tensors(self, graph, ops_in_exec_path): - """Determines the tensors that will be traced.""" - - self._traced_tensorname_to_cache_idx_map = {} - self._cache_idx_to_tensor_idx = [] - operations = graph.get_operations() - checkpoint_operations = self._get_checkpoints(graph) - for op_id, op in enumerate(operations): - if checkpoint_operations and op.name not in checkpoint_operations: - continue - user_included = self._is_user_included_op(op) - user_excluded = self._is_user_excluded_op(op) - in_exec_path = op in ops_in_exec_path - if self._skip_op(op_id, op, user_included, user_excluded, in_exec_path): - continue - for i in range(len(op.outputs)): - out_tensor = op.outputs[i] - if self._skip_tensor(op_id, out_tensor, user_included, - user_excluded): - continue - tensor_name = out_tensor.name - if tensor_name in self._traced_tensorname_to_cache_idx_map: - raise ValueError( - 'Tensor name %s should not be already in ' - 'traced_tensorname_to_cache_idx_map'%tensor_name) - if tensor_name not in self._tensorname_idx_map: - raise ValueError( - 'Tensor name %s is not in the tensorname_idx_map'%tensor_name) - tensor_idx = self._tensorname_idx_map[tensor_name] - cache_idx = len(self._traced_tensorname_to_cache_idx_map) - self._traced_tensorname_to_cache_idx_map[tensor_name] = cache_idx - self._cache_idx_to_tensor_idx.append(tensor_idx) - if len(self._traced_tensorname_to_cache_idx_map) != len( - self._cache_idx_to_tensor_idx): - raise RuntimeError('len(self._traced_tensorname_to_cache_idx_map) != ' - 'len(self._cache_idx_to_tensor_idx') - - def _check_trace_files(self): - """Checks if any requirements for trace files are satisfied.""" - - if not self._trace_dir: - # traces will be written to stderr. No need to check trace files. - return - if _trace_files_need_precreated(self._trace_dir): - for replica_id in range(0, self._num_replicas): - trace_file_path = os.path.join( - self._trace_dir, - _COMPACT_TRACE_FILE_PREFIX) + '%d'%replica_id - if not gfile.Exists(trace_file_path): - raise RuntimeError( - '%s must be pre-created with the ' - 'appropriate properties.'%trace_file_path) - else: - if not gfile.Exists(self._trace_dir): - gfile.MkDir(self._trace_dir) - if not gfile.Exists(self._trace_dir): - raise RuntimeError('Failed to create %s'%self._trace_dir) - - def _pre_tracing(self, graph, fetches): - """Work needs to be done prior to TPU or CPU tracing.""" - - self._check_trace_files() - operations = graph.get_operations() - (opname_idx_map, tensor_list, self._tensorname_idx_map) = ( - TensorTracer._make_op_and_tensor_maps(operations)) - self._write_config_section() - self._write_op_list_section(operations) - self._write_tensor_list_section(tensor_list, opname_idx_map) - # Filter out the operations that won't be executed. - # if fetches=None, then ops_in_exec_path = set(operations) - ops_in_exec_path = self._filter_execution_path_operations(operations, - fetches) - self._determine_traced_tensors(graph, ops_in_exec_path) - self._write_cache_index_map_section() - # Does the topological sort before adding any nodes to the graph. - (succeed, sorted_or_cycle) = TensorTracer.topological_sort(graph) - if self._use_tensor_values_cache(): - _create_tensor_values_cache(graph, - len(self._cache_idx_to_tensor_idx)) - return (ops_in_exec_path, succeed, sorted_or_cycle) - - def _post_tracing(self, succeed, sorted_or_cycle): - """Work needs to be done after TPU or CPU tracing.""" - - self._write_reason_section() - self._write_graph_section(succeed, sorted_or_cycle) - self._close_report_file() - - def _get_checkpoints(self, graph): - """Returns the list of Ops that produce the tensors traced with API. - - Args: - graph: the graph of Ops. - - Returns: - A set of operation names which should be traced. - """ - - self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, - _TENSOR_TRACER_CHECKPOINT)) - checkpoint_operations = set() - tensor_tracer_variables = graph.get_collection(_TENSOR_TRACER_COLLECTION) - for (tensor, checkpoint_name) in tensor_tracer_variables: - self._write_report('%s %s\n'%(tensor.name, checkpoint_name)) - checkpoint_operations.add(tensor.op.name) - self._write_report('%s %s\n'%(_MARKER_SECTION_END, - _TENSOR_TRACER_CHECKPOINT)) - return checkpoint_operations - - def _generate_flush_cache_op(self, graph, start_replica, on_tpu): - """Generates an Op that will flush the cache to file. - - Args: - graph: the graph of Ops - start_replica: the ID of the first replica being flushed by this Op. - on_tpu: if the graph is executed on TPU. - - Returns: - The Op to flush the cache to file. - """ - def _make_flush_fun(replica_id): - """Makes a function for flushing the cache for the given replica.""" - - def _fun(): - """A function that flushes the cache to a file.""" - - def _flush_fun(cache): - """Flushes the cache to a file.""" - - if isinstance(replica_id, str): - replica_id_str = replica_id - else: - replica_id_str = '%d'%replica_id - if self._trace_dir: - output_path = os.path.join(self._trace_dir, - _COMPACT_TRACE_FILE_PREFIX) \ - + replica_id_str - output_stream = _OUTPUT_STREAM_ESCAPE + output_path - else: - output_stream = sys.stderr - new_step_line = _REPLICA_ID_TAG + replica_id_str - print_op = logging_ops.print_v2( - new_step_line, '\n', - cache, '\n', - summarize=-1, - output_stream=output_stream) - with ops.control_dependencies([print_op]): - return constant_op.constant(0).op - - cache = _get_tensor_values_cache(graph) - if on_tpu: - flush_op = tpu.outside_compilation(_flush_fun, cache.value()) - else: - flush_op = _flush_fun(cache.value()) - with ops.control_dependencies([flush_op]): - reset_value = constant_op.constant(_COMPACT_TRACE_ENTRY_INIT_VALUE, - dtype=cache.dtype, - shape=cache.shape) - assign_op = state_ops.assign(cache, reset_value).op - with ops.control_dependencies([assign_op]): - return flush_op.outputs[0] - - return _fun - - def _f(replica_id): - return _make_flush_fun(replica_id) - def _eq(x): - return math_ops.equal(x, self._replica_id) - def _do_nothing(): - return constant_op.constant(0) - - return control_flow_ops.case({\ - _eq(start_replica): _f(start_replica), \ - _eq(start_replica+1): _f(start_replica+1), \ - _eq(start_replica+2): _f(start_replica+2), \ - _eq(start_replica+3): _f(start_replica+3), \ - _eq(start_replica+4): _f(start_replica+4), \ - _eq(start_replica+5): _f(start_replica+5), \ - _eq(start_replica+6): _f(start_replica+6), \ - _eq(start_replica+7): _f(start_replica+7), \ - }, - default=_do_nothing, - exclusive=True).op - - def _flush_tensor_values_cache(self, graph, tensor_fetches, op_fetches, - on_tpu): - """Flushes the intermediate tensor values in the graph to the cache. - - Args: - graph: the graph of Ops - tensor_fetches: list of tensor results returned by the model_fn. - op_fetches: list of ops that are returned by the model_fn, e.g., train_op. - on_tpu: if the graph is executed on TPU. - - Returns: - An identical copy of tensor_fetches. - """ - # Add a dependency to op and tensor fetches to make sure that all tracing - # ops are executed before flushing trace results. - with ops.control_dependencies(op_fetches + - [tensor.op for tensor in tensor_fetches]): - flush_cache_op_list = [] - for host in range(self._num_hosts): - start_replica = host * 8 - flush_op = self._generate_flush_cache_op(graph, start_replica, on_tpu) - flush_cache_op_list.append(flush_op) - return control_flow_ops.tuple(tensor_fetches, - control_inputs=flush_cache_op_list) - - def _process_tensor_fetches(self, tensor_fetches): - """Check that tensor_fetches is not empty and have valid tensors.""" - # If none or empty list. - if tensor_fetches is None: - raise RuntimeError('tensor_fetches provided to tensor_tracer cannot be ' - 'None.') - if not isinstance(tensor_fetches, (list, tuple)): - tensor_fetches = [tensor_fetches] - elif not tensor_fetches: - raise RuntimeError('tensor_fetches provided to tensor_tracer cannot be ' - 'empty list.') - fetches = [] - for fetch in tensor_fetches: - if isinstance(fetch, ops.Tensor): - fetches.append(fetch) - else: - raise RuntimeError('Given tensor_fetch:%s is not a tensor.' % fetch) - return fetches - - def _process_op_fetches(self, op_fetches): - """Check that op_fetches have valid ops.""" - if op_fetches is None: - return [] - - if not isinstance(op_fetches, (list, tuple)): - op_fetches = [op_fetches] - - fetches = [] - for fetch in op_fetches: - if isinstance(fetch, ops.Operation): - fetches.append(fetch) - else: - logging.warning('Ignoring the given op_fetch:%s, which is not an op.' % - fetch) - return fetches - - def _convert_fetches_to_input_format(self, input_fetches, current_fetches): - """Changes current_fetches' format, so that it matches input_fetches.""" - if isinstance(input_fetches, ops.Tensor): - if len(current_fetches) != 1: - raise RuntimeError('Tensor tracer input/output fetches do not match.') - return current_fetches[0] - else: - if len(current_fetches) != len(current_fetches): - raise RuntimeError('Tensor tracer input/output fetches do not match.') - elif isinstance(input_fetches, tuple): - return tuple(current_fetches) - else: - return current_fetches - - def _get_op_control_flow_context(self, op): - """Returns the control flow of the given op. - - Args: - op: tf.Operation for which the control flow context is requested. - Returns: - op_control_flow_context: which the is control flow context of the given - op. If the operation type is LoopExit, returns the outer control flow - context. - """ - # pylint: disable=protected-access - op_control_flow_context = op._control_flow_context - # pylint: enable=protected-access - if control_flow_util.IsLoopExit(op): - op_control_flow_context = op_control_flow_context.outer_context - return op_control_flow_context - - def _trace_execution(self, graph, - tensor_fetches, - op_fetches=None, - on_tpu=True): - """Commong tracing function for both CPU and TPUs. - - The caller function should set _device_type, _num_replicas, - _num_replicas_per_host, _num_hosts and _replica_id before calling - _trace_execution. - - - Args: - graph: the graph of Ops executed on the TPU. - tensor_fetches: a (list,tuple,or a single object) of tensor fetches - returned by model_fn given to session.run. Function must be provided - with as least one tensor to fetch. - op_fetches: A list of op fetches returned by model_fn given to - session.run. op_fetches and tensor_fetches are used to determine the - nodes that will be executed. Can be None. - on_tpu: True if executing on TPU. - - Returns: - tensor_fetches: an exact copy of tensor_fetches that has additional - dependencies. - Raises: - RuntimeError: If tensor_fetches is None or empty. - """ - def _cast_unsupported_dtypes(tensor): - """Casts tensor to a supported type.""" - - if tensor.dtype.__eq__(dtypes.int64): - # outside-compilation doesn't support int64 input yet. - return math_ops.cast(tensor, dtypes.int32) - if tensor.dtype.__eq__(dtypes.bfloat16) or tensor.dtype.__eq__( - dtypes.float16): - # Since host can't handle bf16, convert tensor to f32. - return math_ops.cast(tensor, dtypes.float32) - return tensor - - TensorTracer.check_device_type(self._device_type) - # Check in_tensor_fetches, and op_fetches and convert them to lists. - processed_t_fetches = self._process_tensor_fetches(tensor_fetches) - op_fetches = self._process_op_fetches(op_fetches) - all_fetches = op_fetches + [tensor.op for tensor in processed_t_fetches] - - # Filter the set of ops that will be executed, and topological sort. - (exec_op_set, succeed, sorted_or_cycle) = self._pre_tracing(graph, - all_fetches) - - tensor_fetch_set = set(processed_t_fetches) - tracing_ops = [] - - # pylint: disable=protected-access - current_control_flow_context = graph._get_control_flow_context() - # pylint: enable=protected-access - - # Trace ops only if they are in the execution path. - for op in exec_op_set: - for i in range(len(op.outputs)): - out_tensor = op.outputs[i] - tensor_name = out_tensor.name - if tensor_name not in self._traced_tensorname_to_cache_idx_map: - continue - # Create the list of consumers before calling _preprocess_traced_tensor. - # Otherwise, adding control input below, will introduce a cycle in the - # graph. - consumers = out_tensor.consumers() - # Not all consumers may be in the exec path. Filter out the consumers - # to keep the graph simpler. - consumers = [cop for cop in consumers if cop in exec_op_set] - - # If there is no consumer of the tensor, there is no need to trace it; - # unless the tensor itself is one of the fetches. - is_a_fetched_tensor = out_tensor in tensor_fetch_set - if (not consumers) and (not is_a_fetched_tensor): - continue - - op_control_flow_context = self._get_op_control_flow_context(op) - # pylint: disable=protected-access - graph._set_control_flow_context(op_control_flow_context) - # pylint: enable=protected-access - processed_out_tensor = self._preprocess_traced_tensor(out_tensor) - - if on_tpu: - processed_out_tensor = _cast_unsupported_dtypes(processed_out_tensor) - - if self._use_tensor_values_cache(): - cache_idx = self._traced_tensorname_to_cache_idx_map[tensor_name] - trace_op = self._save_tensor_value_to_cache_op(graph, - cache_idx, - processed_out_tensor) - elif on_tpu: - trace_op = tpu.outside_compilation( - self._make_tensor_trace_fun(tensor_name), processed_out_tensor) - else: - trace_fun = self._make_tensor_trace_fun(tensor_name) - trace_op = trace_fun(processed_out_tensor) - - if is_a_fetched_tensor: - tracing_ops.append(trace_op) - continue - # Add it to all consumers, as some consumers may not be executed if they - # are in a control flow. - for consumer_op in consumers: - # pylint: disable=protected-access - consumer_op._add_control_input(trace_op) - # pylint: enable=protected-access - - # pylint: disable=protected-access - graph._set_control_flow_context(current_control_flow_context) - # pylint: enable=protected-access - if tracing_ops: - # If we are tracing a fetched tensor, their dependency is stored in - # tracing_ops. - processed_t_fetches = control_flow_ops.tuple(processed_t_fetches, - control_inputs=tracing_ops) - if self._use_tensor_values_cache(): - processed_t_fetches = self._flush_tensor_values_cache(graph, - processed_t_fetches, - op_fetches, - on_tpu=on_tpu) - self._post_tracing(succeed, sorted_or_cycle) - # processed_t_fetches is a list at this point. Convert it to the same - # format as given in tensor_fetches. - return self._convert_fetches_to_input_format(tensor_fetches, - processed_t_fetches) - - def trace_tpu(self, graph, - tensor_fetches, - op_fetches=None, - num_replicas=None, - num_replicas_per_host=None, - num_hosts=None): - """Traces the tensors generated by TPU Ops in a TF graph. - - Args: - graph: the graph of Ops executed on the TPU. - tensor_fetches: a (list,tuple,or a single object) of tensor fetches - returned by model_fn given to session.run. Function must be provided - with as least one tensor to fetch. - op_fetches: A list of op fetches returned by model_fn given to - session.run. op_fetches and tensor_fetches are used to determine the - nodes that will be executed. Can be None. - num_replicas: number of replicas used on the TPU. - num_replicas_per_host: number of replicas per TPU host. - num_hosts: total number of TPU hosts. - - Returns: - tensor_fetches: an exact copy of tensor_fetches that has additional - dependencies. - Raises: - RuntimeError: If num_replicas_per_host > 8. - RuntimeError: If tensor_fetches is None or empty. - """ - - if graph in TensorTracer._traced_graphs: - logging.warning('Graph is already rewritten with tensor tracer, ignoring ' - 'multiple calls.') - return tensor_fetches - else: - TensorTracer._traced_graphs.add(graph) - self._device_type = _DEVICE_TYPE_TPU - self._num_replicas = num_replicas - self._num_replicas_per_host = num_replicas_per_host - self._num_hosts = num_hosts - if self._num_replicas is not None: - if self._num_replicas_per_host is None: - self._num_replicas_per_host = 8 - if self._num_hosts is None: - self._num_hosts = num_replicas // self._num_replicas_per_host + \ - (num_replicas % self._num_replicas_per_host > 0) - - if self._num_replicas_per_host > 8: - # Checks for the assumption in _generate_flush_cache_op(). - raise RuntimeError('num_replicas_per_host (%d) is ' - 'greater than 8'%self._num_replicas_per_host) - if self._graph_dump_path: - graph_io.write_graph(graph, self._graph_dump_path, - 'graph_before_tt.pbtxt') - with graph.as_default(): - self._add_replica_id_to_graph() - tensor_fetches = self._trace_execution(graph, tensor_fetches, op_fetches, - on_tpu=True) - if self._graph_dump_path: - graph_io.write_graph(graph, self._graph_dump_path, - 'graph_after_tt.pbtxt') - return tensor_fetches - - def trace_cpu(self, graph, tensor_fetches, op_fetches=None): - """Traces the tensors generated by CPU Ops in a TF graph. - - Args: - graph: the graph of Ops executed on the CPU. - tensor_fetches: a (list,tuple,or a single object) of tensor fetches - returned by model_fn given to session.run. Function must be provided - with as least one tensor to fetch. - op_fetches: A list of op fetches returned by model_fn given to - session.run. op_fetches and tensor_fetches are used to determine the - nodes that will be executed. Can be None. - - Returns: - tensor_fetches: an exact copy of tensor_fetches that has additional - dependencies. - Raises: - RuntimeError: If tensor_fetches is None or empty. - """ - - if graph in TensorTracer._traced_graphs: - logging.warning('Graph is already rewritten with tensor tracer, ignoring ' - 'multiple calls.') - return tensor_fetches - else: - TensorTracer._traced_graphs.add(graph) - - self._device_type = _DEVICE_TYPE_CPU - self._num_replicas = 1 - self._num_replicas_per_host = 1 - self._num_hosts = 1 - self._replica_id = 0 - if self._graph_dump_path: - graph_io.write_graph(graph, self._graph_dump_path, - 'graph_before_tt.pbtxt') - with graph.as_default(): - tensor_fetches = self._trace_execution(graph, tensor_fetches, op_fetches, - on_tpu=False) - if self._graph_dump_path: - graph_io.write_graph(graph, self._graph_dump_path, - 'graph_after_tt.pbtxt') - return tensor_fetches - - +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.tensor_tracer import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/topology.py b/tensorflow/contrib/tpu/python/tpu/topology.py index 00ee21e694d15d2e795b0b35289e1e116b9e76cf..5bf805752cf51b0a0f4b7400b18b63aae93cf831 100644 --- a/tensorflow/contrib/tpu/python/tpu/topology.py +++ b/tensorflow/contrib/tpu/python/tpu/topology.py @@ -1,220 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ====================================== -"""Defines the `Topology` class, that describes a TPU fabric topology.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import numpy as np -from six.moves import xrange # pylint: disable=redefined-builtin - -from tensorflow.core.protobuf.tpu import topology_pb2 - - -def _tpu_device_name(job, task, device): - """Returns the device name for the TPU `device` on `task` of `job`.""" - if job is None: - return "/task:%d/device:TPU:%d" % (task, device) - else: - return "/job:%s/task:%d/device:TPU:%d" % (job, task, device) - - -def _tpu_host_device_name(job, task): - """Returns the device name for the CPU device on `task` of `job`.""" - if job is None: - return "/task:%d/device:CPU:0" % task - else: - return "/job:%s/task:%d/device:CPU:0" % (job, task) - - -class Topology(object): - """Describes a set of TPU devices. - - Represents both the shape of the physical mesh, and the mapping between - TensorFlow TPU devices to physical mesh coordinates. - """ - - def __init__(self, serialized=None, mesh_shape=None, device_coordinates=None): - """Builds a Topology object. - - If `serialized` is not `None`, the topology is parsed from `serialized` and - the other arguments are ignored. Otherwise, the topology is computed from - `mesh_shape` and `device_coordinates`. - - Args: - serialized: A serialized `TopologyProto`, or `None`. If not `None`, the - serialized proto is parsed to discover the topology. - mesh_shape: A sequence of 3 positive integers, or `None`. If not `None`, - the shape of the TPU topology, in number of cores. Ignored if - `serialized` is not `None`. - device_coordinates: A rank 3 numpy array that describes the mapping from - TensorFlow TPU devices to TPU fabric coordinates, or `None`. Ignored - if `serialized is not `None`. - - Raises: - ValueError: If `serialized` does not describe a well-formed topology. - ValueError: If `serialized` is `None` and `mesh_shape` is not a sequence - of 3 positive integers. - ValueError: If `serialized` is `None` and `device_coordinates` is not a - rank 3 numpy int32 array that describes a valid coordinate mapping. - """ - - self._serialized = serialized - - if serialized: - self._parse_topology(serialized) - else: - self._mesh_shape = np.asarray(mesh_shape, dtype=np.int32) - self._device_coordinates = np.asarray(device_coordinates, np.int32) - if len(self._mesh_shape) != 3 or any(self._mesh_shape < 1): - raise ValueError("`mesh_shape` must be a sequence of 3 positive " - "entries; got {}".format(self._mesh_shape)) - - if (len(self._device_coordinates.shape) != 3 or - self._device_coordinates.shape[2] != len(self._mesh_shape)): - raise ValueError("`device_coordinates` must be a rank 3 int32 array " - "with minor dimension equal to the mesh shape rank") - - self._topology_tasks, self._topology_devices = self._invert_topology() - - def _parse_topology(self, serialized): - """Parses a serialized `TopologyProto` into `self`.""" - proto = topology_pb2.TopologyProto() - proto.ParseFromString(serialized) - - self._mesh_shape = np.array(proto.mesh_shape, dtype=np.int32) - if len(self._mesh_shape) != 3 or any(self._mesh_shape < 1): - raise ValueError("`mesh_shape` must be a vector of size 3 with positive " - "entries; got {}".format(self._mesh_shape)) - - if proto.num_tasks < 0: - raise ValueError("`num_tasks` must be >= 0; got {}".format( - proto.num_tasks)) - if proto.num_tpu_devices_per_task < 0: - raise ValueError("`num_tpu_devices_per_task` must be >= 0; got {}".format( - proto.num_tpu_devices_per_task)) - - expected_coordinates_size = ( - proto.num_tasks * proto.num_tpu_devices_per_task * len( - proto.mesh_shape)) - if len(proto.device_coordinates) != expected_coordinates_size: - raise ValueError("`device_coordinates` must have shape num_tasks ({}) * " - "num_tpu_devices_per_task ({}) * len(mesh_shape) ({}); " - "got shape {}".format(proto.num_tasks, - proto.num_tpu_devices_per_task, - proto.mesh_shape, - len(proto.device_coordinates))) - - coords = np.array(proto.device_coordinates, dtype=np.int32) - if any(coords < 0): - raise ValueError("`device_coordinates` must be >= 0") - coords = coords.reshape((proto.num_tasks, proto.num_tpu_devices_per_task, - len(proto.mesh_shape))) - self._device_coordinates = coords - - def _invert_topology(self): - """Inverts a [task,device,axis] topology to [x,y,z] -> task/device maps.""" - tasks = np.full(list(self.mesh_shape), -1, dtype=np.int32) - devices = np.full(list(self.mesh_shape), -1, dtype=np.int32) - for task in xrange(self.device_coordinates.shape[0]): - for device in xrange(self.device_coordinates.shape[1]): - x, y, z = self.device_coordinates[task, device, :] - tasks[x, y, z] = task - devices[x, y, z] = device - return tasks, devices - - @property - def mesh_shape(self): - """A rank 1 int32 array describing the shape of the TPU topology.""" - return self._mesh_shape - - @property - def mesh_rank(self): - """Returns the number of dimensions in the mesh.""" - return len(self._mesh_shape) - - @property - def device_coordinates(self): - """Describes the mapping from TPU devices to topology coordinates. - - Returns: - A rank 3 int32 array with shape `[tasks, devices, axis]`. - `tasks` is the number of tasks in the TPU cluster, `devices` is the number - of TPU devices per task, and `axis` is the number of axes in the TPU - cluster topology. Each entry gives the `axis`-th coordinate in the - topology of a task/device pair. TPU topologies are 3-dimensional, with - dimensions `(x, y, core number)`. - """ - return self._device_coordinates - - def task_ordinal_at_coordinates(self, device_coordinates): - """Returns the TensorFlow task number attached to `device_coordinates`. - - Args: - device_coordinates: An integer sequence describing a device's physical - coordinates in the TPU fabric. - - Returns: - Returns the TensorFlow task number that contains the TPU device with those - physical coordinates. - """ - return self._topology_tasks[tuple(device_coordinates)] - - def tpu_device_ordinal_at_coordinates(self, device_coordinates): - """Returns the TensorFlow device number at `device_coordinates`. - - Args: - device_coordinates: An integer sequence describing a device's physical - coordinates in the TPU fabric. - - Returns: - Returns the TensorFlow device number within the task corresponding to - attached to the device with those physical coordinates. - """ - return self._topology_devices[tuple(device_coordinates)] - - def cpu_device_name_at_coordinates(self, device_coordinates, job=None): - """Returns the CPU device attached to a logical core.""" - return _tpu_host_device_name( - job, self._topology_tasks[tuple(device_coordinates)]) - - def tpu_device_name_at_coordinates(self, device_coordinates, job=None): - """Returns the name of the TPU device assigned to a logical core.""" - return _tpu_device_name(job, - self._topology_tasks[tuple(device_coordinates)], - self._topology_devices[tuple(device_coordinates)]) - - @property - def num_tasks(self): - """Returns the number of TensorFlow tasks in the TPU slice.""" - return self._device_coordinates.shape[0] - - @property - def num_tpus_per_task(self): - """Returns the number of TPU devices per task in the TPU slice.""" - return self._device_coordinates.shape[1] - - def serialized(self): - """Returns the serialized form of the topology.""" - if self._serialized is None: - proto = topology_pb2.TopologyProto() - proto.mesh_shape[:] = list(self._mesh_shape) - proto.num_tasks = self._device_coordinates.shape[0] - proto.num_tpu_devices_per_task = self._device_coordinates.shape[1] - proto.device_coordinates.extend(list(self._device_coordinates.flatten())) - self._serialized = proto.SerializeToString() - - return self._serialized +# pylint: disable=wildcard-import,unused-import,redefined-builtin +from tensorflow.python.tpu.topology import * +# pylint: enable=wildcard-import,unused-import,redefined-builtin diff --git a/tensorflow/contrib/tpu/python/tpu/tpu.py b/tensorflow/contrib/tpu/python/tpu/tpu.py index 3b2d0534773fa0cce3c515cfaa7102cec195fcc3..5364b20f231ac7af8adf943c3d5e21921b7a06a9 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu.py @@ -1,1579 +1,25 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ====================================== - -"""Library of TPU helper functions.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections -from six.moves import xrange # pylint: disable=redefined-builtin - -from tensorflow.contrib.compiler import xla -from tensorflow.contrib.framework.python.framework import experimental -from tensorflow.contrib.tpu.python.ops import tpu_ops -from tensorflow.contrib.tpu.python.tpu import tpu_function -from tensorflow.core.framework import attr_value_pb2 -from tensorflow.core.protobuf.tpu import dynamic_padding_pb2 as dynamic_padding -from tensorflow.python.compat import compat as api_compat -from tensorflow.python.framework import device as pydev -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors -from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_shape -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import control_flow_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import variable_scope -from tensorflow.python.platform import tf_logging as logging -from tensorflow.python.util import compat -from tensorflow.python.util import nest - - -# Operations that indicate some error in the users graph, e.g. a placeholder -# that's introduced outside of the infeed. -_BLACKLISTED_OPS = set([ - "Placeholder", -]) - -# XLA doesn't currently support reading of intermediate tensors, thus some ops -# are not supported. -_UNSUPPORTED_OPS = set([ - "AudioSummary", - "AudioSummaryV2", - "HistogramSummary", - "ImageSummary", - "MergeSummary", - "Print", - "ScalarSummary", - "TensorSummary", - "TensorSummaryV2", - ]) - -_MAX_WARNING_LINES = 5 - -_TPU_REPLICATE_ATTR = "_tpu_replicate" -_TPU_COMPILATION_STATUS_ATTR = "_tpu_compilation_status" -_OUTSIDE_COMPILATION_ATTR = "_xla_outside_compilation" - - -def _tpu_system_device_name(job): - """Returns the device name for the TPU_SYSTEM device of `job`.""" - if job is None: - return "/device:TPU_SYSTEM:0" - else: - return "/job:%s/device:TPU_SYSTEM:0" % job - - -def initialize_system(embedding_config=None, job=None): - """Initializes a distributed TPU system for use with TensorFlow. - - Args: - embedding_config: If not None, a `TPUEmbeddingConfiguration` proto - describing the desired configuration of the hardware embedding lookup - tables. If embedding_config is None, no hardware embeddings can be used. - job: The job (the XXX in TensorFlow device specification /job:XXX) that - contains the TPU devices that will be initialized. If job=None it is - assumed there is only one job in the TensorFlow flock, and an error will - be returned if this assumption does not hold. - Returns: - A serialized `TopologyProto` that describes the TPU system. Note: - the topology must be evaluated using `Session.run` before it can be used. - """ - config_string = ("" if embedding_config is None else - embedding_config.SerializeToString()) - with ops.device(_tpu_system_device_name(job)): - return tpu_ops.configure_distributed_tpu(embedding_config=config_string) - - -def shutdown_system(job=None): - """Shuts down a running a distributed TPU system.""" - with ops.device(_tpu_system_device_name(job)): - shutdown_distributed_tpu = tpu_ops.shutdown_distributed_tpu() - return shutdown_distributed_tpu - - -def core(num): - """Returns the device name for a core in a replicated TPU computation. - - Args: - num: the virtual core number within each replica to which operators should - be assigned. - Returns: - A device name, suitable for passing to `tf.device()`. - """ - return "device:TPU_REPLICATED_CORE:{}".format(num) - - -class TPUReplicateContext(control_flow_ops.XLAControlFlowContext): - """A `ControlFlowContext` for nodes inside a TPU computation. - - The primary role of `TPUReplicateContext` is to mark operators inside a - tpu.replicate() computation with the attribute "_tpu_replicate=XYZ", where XYZ - is a unique name. - - We use a `ControlFlowContext` to perform the annotation since it integrates - with Tensorflow constructs like ResourceVariables. For example, if a - `ResourceVariable` is constructed inside a tpu.replicate() block, the - `ResourceVariable` implementation can use - `with ops.control_dependencies(None)` to build the variable's definition - outside the replicated computation. - """ - - def __init__(self, name, num_replicas, pivot): - """Builds a new TPUReplicateContext. - - Args: - name: a unique name for the context, used to populate the `_tpu_replicate` - attribute. - num_replicas: an integer that gives the number of replicas for the - computation. - pivot: a pivot node. Nodes in the TPUReplicateContext that do not have any - inputs will have a control dependency on the pivot node. This ensures - that nodes are correctly included in any enclosing control flow - contexts. - """ - super(TPUReplicateContext, self).__init__() - self._num_replicas = num_replicas - self._outer_device_function_stack = None - self._oc_dev_fn_stack = None - self._outside_compilation_cluster = None - self._outside_compilation_counter = 0 - self._in_gradient_colocation = None - self._gradient_colocation_stack = [] - self._host_compute_core = [] - self._name = name - self._name_as_bytes = compat.as_bytes(name) - self._unsupported_ops = [] - self._pivot = pivot - self._replicated_vars = {} - - def get_replicated_var_handle(self, name, vars_): - """Returns a variable handle for replicated TPU variable 'var'. - - This is a method used by an experimental replicated variable implementation - and is not intended as a public API. - - Args: - name: The common name of the variable. - vars_: The replicated TPU variables. - - Returns: - The handle of the TPU replicated input node. - """ - handle = self._replicated_vars.get(name) - if handle is not None: - return handle - - # Builds a TPUReplicatedInput node for the variable, if one does not already - # exist. The TPUReplicatedInput node must belong to the enclosing - # control-flow scope of the TPUReplicateContext. - # TODO(phawkins): consider changing the contract of the TPU encapsulation - # so the TPUReplicatedInput nodes go inside the TPUReplicateContext scope - # instead. - - # pylint: disable=protected-access - graph = ops.get_default_graph() - saved_context = graph._get_control_flow_context() - graph._set_control_flow_context(self.outer_context) - handle = tpu_ops.tpu_replicated_input( - [v.handle for v in vars_], name=name + "/handle") - graph._set_control_flow_context(saved_context) - # pylint: enable=protected-access - self._replicated_vars[name] = handle - return handle - - def report_unsupported_operations(self): - if self._unsupported_ops: - op_str = "\n".join([" %s (%s)" % (op.type, op.name) - for op in self._unsupported_ops[:_MAX_WARNING_LINES]]) - logging.warning("%d unsupported operations found: \n%s", - len(self._unsupported_ops), op_str) - if len(self._unsupported_ops) > _MAX_WARNING_LINES: - logging.warning("... and %d more" % - (len(self._unsupported_ops) - _MAX_WARNING_LINES)) - - def EnterGradientColocation(self, op, gradient_uid): - if op is not None: - self._gradient_colocation_stack.append(op) - if not self._outside_compilation_cluster: - try: - outside_attr = op.get_attr(_OUTSIDE_COMPILATION_ATTR) - if self._in_gradient_colocation: - raise NotImplementedError( - "Cannot nest gradient colocation operations outside compilation" - ) - if gradient_uid == "__unsupported__": - raise NotImplementedError( - "No gradient_uid calling gradient within outside_compilation") - # When we take the gradient of an op X in an outside_compilation - # cluster C in a forward computation we would like to put the ops - # corresponding to the gradient of X into a new outside_compilation - # cluster C'. However, if we take the gradient of X twice, the second - # one should get yet another new outside_compilation cluster C''. - # - # The mechanism we adopt is to use a 'root_cluster' which is the - # cluster that X was in before we took gradients, and a 'gradient_uid' - # which is different for every invocation of gradients, and put the - # gradient of X in cluster 'root_cluster.gradient_uid'. - # - # When taking a gradient of a gradient, some ops will be colocated - # with Op in the forward pass (e.g., cluster root_cluster) and some in - # the backward pass (e.g., cluster root_cluster.initial_gradient_uid). - # We need all of the grad-of-grad ops to be in the same cluster to - # avoid cyclic dependencies between clusters. We adopt a heuristic - # that puts any op clustered with root_cluster. in - # root_cluster.gradient_uid, even if xxx was initial_gradient_uid. - self._in_gradient_colocation = op - parts = outside_attr.split(".") - cluster = parts[0] + "." + gradient_uid - self._EnterOutsideCompilationScope(cluster=cluster) - except ValueError: - # The attr was not present: do nothing. - pass - - def ExitGradientColocation(self, op, gradient_uid): - if op is not None: - if not self._gradient_colocation_stack: - raise errors.InternalError( - op.node_def, op, - "Badly nested gradient colocation: empty stack when popping Op " + - op.name) - last_op = self._gradient_colocation_stack.pop() - if op is last_op: - if op is self._in_gradient_colocation: - self._in_gradient_colocation = None - self._ExitOutsideCompilationScope() - else: - raise errors.InternalError( - op.node_def, op, "Badly nested gradient colocation, expected " + - last_op + ", got " + op.name) - - def _EnterOutsideCompilationScope(self, cluster=None): - - class FakeOp(object): - """A helper class to determine the current device. - - Supports only the type and device set/get methods needed to run the - graph's _apply_device_function method. - """ - - def __init__(self): - self._device = "" - - @property - def type(self): - return "FakeOp" - - @property - def device(self): - return self._device - - def _set_device(self, device): - if isinstance(device, pydev.DeviceSpec): - self._device = device.to_string() - else: - self._device = device - - if self._outside_compilation_cluster: - raise NotImplementedError("Cannot nest outside_compilation clusters") - if cluster: - self._outside_compilation_cluster = cluster - else: - self._outside_compilation_cluster = str(self._outside_compilation_counter) - self._outside_compilation_counter += 1 - graph = ops.get_default_graph() - fake_op = FakeOp() - graph._apply_device_functions(fake_op) # pylint: disable=protected-access - device = pydev.DeviceSpec.from_string(fake_op.device) - if (device.device_type == "TPU_REPLICATED_CORE" and - device.device_index is not None): - self._host_compute_core.append(self._outside_compilation_cluster + ":" + - str(device.device_index)) - self._oc_dev_fn_stack = graph._device_function_stack # pylint: disable=protected-access - graph._device_function_stack = self._outer_device_function_stack # pylint: disable=protected-access - - def _ExitOutsideCompilationScope(self): - if not self._outside_compilation_cluster: - raise NotImplementedError( - "Attempted to exit outside_compilation scope when not in scope") - self._outside_compilation_cluster = None - graph = ops.get_default_graph() - graph._device_function_stack = self._oc_dev_fn_stack # pylint: disable=protected-access - - def Enter(self): - if not self._outer_device_function_stack: - # Capture the device function stack at the time of first entry - # since that is the stack that will be used outside_compilation. - graph = ops.get_default_graph() - # pylint: disable=protected-access - self._outer_device_function_stack = graph._device_function_stack.copy() - # pylint: enable=protected-access - super(TPUReplicateContext, self).Enter() - - def HostComputeCore(self): - return self._host_compute_core - - def _RemoveExternalControlEdges(self, op): - """Remove any external control dependency on this op.""" - internal_control_inputs = [] - external_control_inputs = [] - for x in op.control_inputs: - # pylint: disable=protected-access - is_internal_op = False - ctxt = x._get_control_flow_context() - while ctxt is not None: - if ctxt == self: - is_internal_op = True - break - ctxt = ctxt._outer_context - if is_internal_op: - internal_control_inputs.append(x) - else: - external_control_inputs.append(x) - # pylint: enable=protected-access - # pylint: disable=protected-access - op._remove_all_control_inputs() - op._add_control_inputs(internal_control_inputs) - # pylint: enable=protected-access - return internal_control_inputs, external_control_inputs - - def AddOp(self, op): - # pylint: disable=protected-access - if op.type in _BLACKLISTED_OPS: - logging.error("Operation of type %s (%s) is not supported on the TPU. " - "Execution will fail if this op is used in the graph. " % - (op.type, op.name)) - - if op.type in _UNSUPPORTED_OPS: - self._unsupported_ops.append(op) - - if any(x.dtype._is_ref_dtype for x in op.inputs): - raise NotImplementedError( - "Non-resource Variables are not supported inside TPU computations " - "(operator name: %s)" % op.name) - if _TPU_REPLICATE_ATTR in op.node_def.attr: - raise ValueError("TPU computations cannot be nested") - op._set_attr(_TPU_REPLICATE_ATTR, - attr_value_pb2.AttrValue(s=self._name_as_bytes)) - if self._outside_compilation_cluster: - op._set_attr( - _OUTSIDE_COMPILATION_ATTR, - attr_value_pb2.AttrValue( - s=compat.as_bytes(self._outside_compilation_cluster))) - if self._num_replicas > 1 or not self._outside_compilation_cluster: - # Prevent feeding or fetching anything that is being compiled, - # and any replicated outside_compilation Op. - op.graph.prevent_feeding(op) - op.graph.prevent_fetching(op) - - # Remove any control edges from outer control flow contexts. These may cause - # mismatched frame errors. - (internal_control_inputs, - external_control_inputs) = self._RemoveExternalControlEdges(op) - - if not op.inputs: - # Add a control edge from the control pivot to this op. - if not internal_control_inputs: - # pylint: disable=protected-access - op._add_control_input(self.GetControlPivot()) - # pylint: enable=protected-access - else: - for index in xrange(len(op.inputs)): - x = op.inputs[index] - real_x = self.AddValue(x) - if real_x != x: - op._update_input(index, real_x) # pylint: disable=protected-access - - if external_control_inputs: - # Use an identity to pull control inputs as data inputs. Note that we - # ignore ops which don't have outputs. TODO(phawkins): fix that. - with ops.control_dependencies(None): - self.Enter() - external_control_inputs = [ - array_ops.identity(x.outputs[0]).op - for x in external_control_inputs - if x.outputs - ] - self.Exit() - # pylint: disable=protected-access - op._add_control_inputs(external_control_inputs) - # pylint: enable=protected-access - - # Mark op's outputs as seen by this context and any outer contexts. - output_names = [x.name for x in op.outputs] - context = self - while context is not None: - # pylint: disable=protected-access - context._values.update(output_names) - context = context._outer_context - # pylint: enable=protected-access - - if self._outer_context: - self._outer_context.AddInnerOp(op) - - def AddValue(self, val): - """Add `val` to the current context and its outer context recursively.""" - if val.name in self._values: - # Use the real value if it comes from outer context. - result = self._external_values.get(val.name) - return val if result is None else result - - result = val - self._values.add(val.name) - if self._outer_context: - result = self._outer_context.AddValue(val) - self._values.add(result.name) - - self._external_values[val.name] = result - - return result - - def AddInnerOp(self, op): - self.AddOp(op) - if self._outer_context: - self._outer_context.AddInnerOp(op) - - @property - def grad_state(self): - # Define the gradient loop state associated with the TPUReplicateContext to - # be None as the TPUReplicateContext does not get nested nor does the - # grad_state outside the TPUReplicateContext affect the graph inside so the - # grad_state should be as if this is the top-level gradient state. - return None - - @property - def back_prop(self): - """Forwards to the enclosing while context, if any.""" - if self.GetWhileContext(): - return self.GetWhileContext().back_prop - return False - - def GetControlPivot(self): - return self._pivot - - -def outside_compilation(computation, *args, **kwargs): - """Builds part of a computation outside any current TPU replicate scope. - - Args: - computation: A Python function that builds the computation to - place on the host. - *args: the positional arguments for the computation. - **kwargs: the keyword arguments for the computation. - - Returns: - The Tensors returned by computation. - """ - args = [] if args is None else args - graph = ops.get_default_graph() - - # If we are in a TPUReplicateContext, signal that we are now - # outside_compilation - initial_context = graph._get_control_flow_context() # pylint: disable=protected-access - context = initial_context - while context: - if isinstance(context, TPUReplicateContext): - context._EnterOutsideCompilationScope() # pylint: disable=protected-access - context = context.outer_context - - retval = computation(*args, **kwargs) - - # If we are in a TPUReplicateContext, signal that we are no longer - # outside_compilation - final_context = graph._get_control_flow_context() # pylint: disable=protected-access - if initial_context is not final_context: - raise NotImplementedError( - "Control-flow context cannot be different at start and end of an " - "outside_compilation scope") - context = initial_context - while context: - if isinstance(context, TPUReplicateContext): - context._ExitOutsideCompilationScope() # pylint: disable=protected-access - context = context.outer_context - - return retval - - -def replicate(computation, - inputs=None, - infeed_queue=None, - device_assignment=None, - name=None, - maximum_shapes=None): - """Builds a graph operator that runs a replicated TPU computation. - - Args: - computation: A Python function that builds the computation to replicate. - inputs: A list of lists of input tensors or `None` (equivalent to - `[[]]`), indexed by `[replica_num][input_num]`. All replicas must - have the same number of inputs. Each input can be a nested structure - containing values that are convertible to tensors. Note that passing an - N-dimension list of compatible values will result in a N-dimention list of - scalar tensors rather than a single Rank-N tensors. If you need different - behavior, convert part of inputs to tensors with `tf.convert_to_tensor`. - infeed_queue: If not `None`, the `InfeedQueue` from which to append a tuple - of arguments as inputs to computation. - device_assignment: If not `None`, a `DeviceAssignment` describing the - mapping between logical cores in the computation with physical cores in - the TPU topology. Uses a default device assignment if `None`. The - `DeviceAssignment` may be omitted if each replica of the computation uses - only one core, and there is either only one replica, or the number of - replicas is equal to the number of cores in the TPU system. - name: (Deprecated) Does nothing. - maximum_shapes: A nested structure of tf.TensorShape representing the shape - to which the respective component of each input element in each replica - should be padded. Any unknown dimensions (e.g. tf.Dimension(None) in a - tf.TensorShape or -1 in a tensor-like object) will be padded to the - maximum size of that dimension over all replicas. Note that if the input - dimension is already static, we won't do padding on it and we require the - maximum_shapes to have the same value or None on that dimension. The - structure of `maximum_shapes` needs to be the same as `inputs[0]`. - Returns: - A list of outputs, indexed by `[replica_num]` each output can be a nested - structure same as what computation() returns with a few exceptions. - - Exceptions include: - 1) None output: a NoOp would be returned which control-depends on - computation. - 2) Single value output: A tuple containing the value would be returned. - 3) Operation-only outputs: a NoOp would be returned which - control-depends on computation. - TODO(b/121383831): Investigate into removing these special cases. - - Raises: - ValueError: If all replicas do not have equal numbers of input tensors. - ValueError: If the number of inputs per replica does not match - the number of formal parameters to `computation`. - ValueError: If the static `inputs` dimensions don't match with the values - given in `maximum_shapes`. - ValueError: If the structure of inputs per replica does not match - the structure of `maximum_shapes`. - """ - return split_compile_and_replicate( - computation, - inputs, - infeed_queue, - device_assignment, - name, - maximum_shapes=maximum_shapes)[1] - - -def _pad_all_input(inputs, padded_shapes): - """Pad all input tensors given padded_shapes. - - The real shape tensors will be concatenated with the padded original inputs. - - Args: - inputs: The original inputs. - padded_shapes: A list of padded shapes for each input. - - Returns: - The padded inputs and a PaddingMap list which maps the padded input - dimension to the real shape argument index. - """ - input_shape_tensors = [] - for core_idx, inputs_per_core in enumerate(inputs): - for idx, input_tensor in enumerate(inputs_per_core): - if core_idx == 0: - input_shape_tensors.append([]) - input_shape_tensors[idx].append(array_ops.shape(input_tensor)) - - maximum_shapes = [] - for shapes_per_input in input_shape_tensors: - maximum_shapes.append( - math_ops.reduce_max(array_ops.stack(shapes_per_input), axis=0)) - - padded_inputs = [] - real_shapes = [] - padding_maps = [] - for core_idx, inputs_per_core in enumerate(inputs): - padded_inputs.append([]) - real_shapes.append([]) - real_shape_idx = len(inputs_per_core) - 1 - for idx, input_tensor in enumerate(inputs_per_core): - input_shape_tensor = input_shape_tensors[idx][core_idx] - input_shape = input_tensor.get_shape() - padded_shape = padded_shapes[idx] - - # The static shape of inputs should be compatible with the given padded - # shapes. - input_shape.assert_is_compatible_with(padded_shape) - - if input_shape.is_fully_defined(): - # Do nothing if the shape of the whole tensor is already static. - padded_inputs[core_idx].append(input_tensor) - else: - # Only pad the non static shape dimension. - for i, s in enumerate(input_shape): - if s.value is None: - if core_idx == 0: - real_shape_idx += 1 - padding_map = dynamic_padding.PaddingMap() - padding_map.arg_index = idx - padding_map.shape_index = i - padding_map.padding_arg_index = real_shape_idx - padding_maps.append(padding_map) - real_shapes[core_idx].append( - math_ops.cast(input_shape_tensor[i], dtypes.uint32)) - - paddings = [] - for i, s in enumerate(padded_shape): - if input_shape[i].value: - # Don't pad if input shape is already static. - padding = [0, 0] - else: - if s.value: - # Pad to the given maximum value. - padding = [0, s.value - input_shape_tensor[i]] - else: - # If maximum value is not given, then pad to the maximum dimension - # among all the cores. - padding = [0, maximum_shapes[idx][i] - input_shape_tensor[i]] - paddings.append(padding) - - padded_input = array_ops.pad(input_tensor, paddings) - padded_inputs[core_idx].append(padded_input) - - num_replicas = len(padded_inputs) - for i in range(num_replicas): - padded_inputs[i].extend(real_shapes[i]) - - return padded_inputs, padding_maps - - -def split_compile_and_replicate(computation, - inputs=None, - infeed_queue=None, - device_assignment=None, - name=None, - use_tpu=True, - maximum_shapes=None): - """Builds graph operators that runs compilation and replicated computation. - - This is a lower level interface than replicate that returns a separate compile - and execute output tensor. In the generated graph the compile op feeds into - the execute op and no additional compilation is incurred when running the - compile op before the execute op. The compile op returns additional - information about the compilation but does not return the compiled program. - - Args: - computation: A Python function that builds the computation to replicate. - inputs: A list of lists of input tensors or `None` (equivalent to - `[[]]`), indexed by `[replica_num][input_num]`. All replicas must - have the same number of inputs. Each input can be a nested structure - containing values that are convertible to tensors. Note that passing an - N-dimension list of compatible values will result in a N-dimention list of - scalar tensors rather than a single Rank-N tensors. If you need different - behavior, convert part of inputs to tensors with `tf.convert_to_tensor`. - infeed_queue: If not `None`, the `InfeedQueue` from which to append a tuple - of arguments as inputs to computation. - device_assignment: If not `None`, a `DeviceAssignment` describing the - mapping between logical cores in the computation with physical cores in - the TPU topology. Uses a default device assignment if `None`. The - `DeviceAssignment` may be omitted if each replica of the computation uses - only one core, and there is either only one replica, or the number of - replicas is equal to the number of cores in the TPU system. - name: (Deprecated) Does nothing. - use_tpu: When false, the input `computation` is executed on the XLA CPU/GPU - backends. Currently, only supports a default placement (computation is - placed on GPU if one is available, and on CPU if not). - maximum_shapes: A nested structure of tf.TensorShape representing the shape - to which the respective component of each input element in each replica - should be padded. Any unknown dimensions (e.g. tf.Dimension(None) in a - tf.TensorShape or -1 in a tensor-like object) will be padded to the - maximum size of that dimension over all replicas. Note that if the input - dimension is already static, we won't do padding on it and we require the - maximum_shapes to have the same value or None on that dimension. The - structure of `maximum_shapes` needs to be the same as `inputs[0]`. - - Returns: - A list of lists with the first list corresponding to the compile op and the - second a list of output tensors, indexed by `[replica_num][output_num]`. - Raises: - ValueError: If all replicas do not have equal numbers of input tensors. - ValueError: If the number of inputs per replica does not match - the number of formal parameters to `computation`. - ValueError: If the static `inputs` dimensions don't match with the values - given in `maximum_shapes`. - ValueError: If the structure of inputs per replica does not match - the structure of `maximum_shapes`. - """ - del name - inputs = [[]] if inputs is None else inputs - - metadata_kwargs = {} - if device_assignment is not None: - # Turn the Numpy array into a flattened list so we can pass it as an - # operator attribute. - metadata_kwargs = { - "topology": - device_assignment.topology.serialized(), - "device_assignment": - device_assignment.core_assignment.flatten().tolist() - } - # TODO(phawkins): remove this case after the forward compatibility window - # expires on 2018-10-5. - if api_compat.forward_compatible(2018, 10, 5): - metadata_kwargs["num_cores_per_replica"] = ( - device_assignment.num_cores_per_replica) - else: - metadata_kwargs["computation_shape"] = [ - device_assignment.num_cores_per_replica - ] - - if ((not isinstance(inputs, list)) or - any(not isinstance(inp, (list, tuple)) for inp in inputs)): - raise TypeError("tpu.replicate() inputs must be a list of lists/tuples") - - num_replicas = len(inputs) - - # No replicas? Nothing to do. - if num_replicas == 0: - return [] - - # Checks all replicas have the same structure. - for i in xrange(1, num_replicas): - nest.assert_same_structure(inputs[0], inputs[i]) - - # Flatten inputs. - flat_inputs = [ - nest.flatten(per_replica_input) for per_replica_input in inputs - ] - # Converts inputs to Tensors. - flat_inputs = [[ops.convert_to_tensor(x) for x in inp] for inp in flat_inputs] - - # Verifies that all replicas have matching numbers and types of inputs - flat_input_types = [x.dtype for x in flat_inputs[0]] - input_arity = len(inputs[0]) - flat_input_arity = len(flat_input_types) - for i in range(num_replicas): - if len(inputs[i]) != input_arity: - raise ValueError("Replicas must have the same number of inputs. " - "Replica 0 had {} inputs, replica {} had {} " - "inputs.".format(input_arity, i, len(inputs[i]))) - - types = [x.dtype for x in flat_inputs[i]] - if types != flat_input_types: - raise ValueError("Replicas must have matching input types. Replica 0 had " - "input types {}, replica {} had input types {}".format( - flat_input_types, i, types)) - - arg_error = xla.check_function_argument_count( - computation, input_arity, infeed_queue) - if arg_error is not None: - if infeed_queue is None: - raise TypeError( - "Supplied computation cannot be called with the specified inputs. " - "You specified %d inputs: %s, but the computation needs %s" % ( - input_arity, str([i.name for i in inputs[0]]), arg_error)) - else: - raise TypeError( - "Supplied computation cannot be called with the specified inputs. " - "You specified %d inputs: %s and %d additional inputs from infeed," - " but the computation needs %s" % (input_arity, str( - [i.name - for i in inputs[0]]), infeed_queue.number_of_tuple_elements, - arg_error)) - - if maximum_shapes: - if infeed_queue: - raise ValueError( - "Dynamic input shapes are not supported with infeed queues") - - # Make sure maximum_shapes has the same structure as inputs. - nest.assert_same_structure(inputs[0], maximum_shapes, check_types=False) - - # Flatten padded shapes. - flat_maximum_shapes = nest.flatten(maximum_shapes) - flat_maximum_shapes = [ - tensor_shape.TensorShape(s) for s in flat_maximum_shapes - ] - - flat_inputs, padding_maps = _pad_all_input(flat_inputs, flat_maximum_shapes) - - serialized_padding_maps = [] - for padding_map in padding_maps: - serialized_padding_maps.append(padding_map.SerializeToString()) - metadata_kwargs["padding_map"] = serialized_padding_maps - - metadata_kwargs["step_marker_location"] = getattr( - computation, "step_marker_location", "STEP_MARK_AT_ENTRY") - - graph = ops.get_default_graph() - - # Fan-in: Builds a TPUReplicatedInput node for each input. - flat_replicated_inputs = [] - for i in range(0, len(flat_inputs[0])): - replicas = [flat_inputs[replica][i] for replica in xrange(num_replicas)] - flat_replicated_inputs.append( - tpu_ops.tpu_replicated_input(replicas, name="input{}".format(i))) - - cluster_name = graph.unique_name("cluster") - pivot = control_flow_ops.no_op(name=cluster_name + "/pivot") - context = TPUReplicateContext( - name=cluster_name, num_replicas=num_replicas, pivot=pivot) - try: - context.Enter() - - metadata = tpu_ops.tpu_replicate_metadata( - num_replicas=num_replicas, use_tpu=use_tpu, **metadata_kwargs) - - with tpu_function.tpu_shard_context( - num_replicas), ops.control_dependencies([metadata]): - - # Add identity ops so even unused inputs are "consumed" by the - # computation. This is to avoid orphaned TPUReplicatedInput nodes. - # TODO(phawkins): consider instead pruning unused TPUReplicatedInput - # and eliding trivial TPUReplicatedInput/TPUReplicatedOutput pairs. - flat_replicated_inputs = [ - array_ops.identity(x, name="replicated_input_{}".format(i)) - for i, x in enumerate(flat_replicated_inputs) - ] - for i in flat_replicated_inputs: - # pylint: disable=protected-access - # Add an attribute to the identity node so that they could be removed in - # encapsulate TPU computation pass if unused. However we don't remove - # inputs when dynamic padding is enabled. - # TODO(rxsang): Use other ways except argument index in padding_map so - # outside compilation can work with dynamic padding correctly. - if maximum_shapes is None: - i.op._set_attr("_tpu_input_identity", - attr_value_pb2.AttrValue(b=True)) - # pylint: enable=protected-access - - # Unflatten the computation inputs to match original input structure. - computation_inputs = nest.pack_sequence_as( - structure=inputs[0], - flat_sequence=flat_replicated_inputs[:flat_input_arity]) - - # If there is an infeed queue, adds the dequeued values to the - # computation's inputs. - if infeed_queue is not None: - infeed_queue.set_number_of_shards(num_replicas) - for t in infeed_queue.generate_dequeue_op(): - computation_inputs.append(t) - - # Only resource variables work inside a TPU computation, so turn on - # resource variables for the computation. - # TODO(phawkins): consider removing this code. It will - # be less confusing to clients if they knowingly choose to use resource - # variables. - # Partitioned variables is not supported (b/112311320). - vscope = variable_scope.get_variable_scope() - saved_use_resource = vscope.use_resource - saved_custom_getter = vscope.custom_getter - - def custom_getter(getter, name, *args, **kwargs): - """Variables on TPU have a few restrictions.""" - partitioner = kwargs["partitioner"] - if partitioner is not None: - kwargs["partitioner"] = None - logging.warning( - "Partitioned variables are not supported on TPU. Got " - "`partitioner` that is {} for variable {}. " - "Setting `partitioner` to `None`." - .format(partitioner, name)) - if saved_custom_getter is None: - return getter(name, *args, **kwargs) - else: - return saved_custom_getter(getter, name, *args, **kwargs) - - vscope.set_use_resource(True) - vscope.set_custom_getter(custom_getter) - - outputs = computation(*computation_inputs) - - vscope.set_use_resource(saved_use_resource) - vscope.set_custom_getter(saved_custom_getter) - - outputs_is_flat = xla.is_flat(outputs) - if outputs_is_flat: - output_tensors, control_deps = _postprocess_flat_outputs(outputs) - else: - output_tensors, control_deps = _postprocess_non_flat_outputs(outputs) - - # tensor_tracer imports tpu.py. Local import to tensor_tracer to avoid - # import-cycle - # pylint: disable=g-import-not-at-top - from tensorflow.contrib.tpu.python.tpu import tensor_tracer - # pylint: enable=g-import-not-at-top - if tensor_tracer.TensorTracer.is_enabled(): - tt = tensor_tracer.TensorTracer() - output_tensors = tt.trace_tpu(ops.get_default_graph(), - output_tensors, control_deps, - num_replicas) - - context.ExitResult(output_tensors) - finally: - context.report_unsupported_operations() - context.Exit() - host_compute_core = context.HostComputeCore() - - if host_compute_core: - attr_value = attr_value_pb2.AttrValue() - attr_value.list.s.extend([compat.as_bytes(x) for x in host_compute_core]) - metadata._set_attr("host_compute_core", attr_value) # pylint: disable=protected-access - - with ops.control_dependencies([metadata]): - if use_tpu: - compile_status = tpu_ops.tpu_compilation_result() - op = compile_status.op - attr_value = attr_value_pb2.AttrValue(s=compat.as_bytes(cluster_name)) - op._set_attr(_TPU_COMPILATION_STATUS_ATTR, attr_value) # pylint: disable=protected-access - else: - compile_status = control_flow_ops.no_op(name="compilation_status") - - if not output_tensors: - # Returns a list of NoOps dependent on the replication Op, indexed by - # [replica_num]. - return [ - compile_status, - [ - control_flow_ops.group(control_deps, name="shard_%d" % i) - for i in range(num_replicas) - ] - ] - - # Fan-out: Builds a TPUReplicatedOutput node for each output. - replicated_outputs = [[] for i in xrange(num_replicas)] - for i, t in enumerate(output_tensors): - # Fan-out: Builds a TPUReplicatedOutput node for each output. - ys = tpu_ops.tpu_replicated_output( - t, num_replicas, name="output{}".format(i)) - - # Wraps the outputs in identity operators so the names of any possible - # `fetch` nodes are preserved by the replication rewrite. - with ops.control_dependencies(control_deps): - for replica in xrange(num_replicas): - replicated_outputs[replica].append( - array_ops.identity( - ys[replica], name="output_%d_shard_%d" % (i, replica))) - - if not outputs_is_flat: - replicated_outputs = [ - nest.pack_sequence_as(outputs, replica_outs) - for replica_outs in replicated_outputs - ] - - return [compile_status, replicated_outputs] - - -def _postprocess_flat_outputs(outputs): - """Validates non-flat outputs, add backs device assignments and other attrs. - - Args: - outputs: Output from `computation` inside `tpu.rewrite`. - - Returns: - Tensors and Operations extracted from outputs. - """ - # Following code segment is to preserve legacy behavior. Previously we only - # supported flat outputs and thus for consistency it was nice to convert even - # single element into a tuple. But now that we support arbitrary output - # structure, this is no longer necessary. - # TODO(b/121383831): Migrate all legacy use cases and delete this special - # case. - # If the computation returns `None`, make it an empty tuple. - if outputs is None: - outputs = tuple() - # If the computation only returned one value, makes it a tuple. - if not isinstance(outputs, collections.Sequence): - outputs = (outputs,) - - # Append `no_op` here so that fetching any return value of this function - # will trigger TPUExecute node. - outputs += (control_flow_ops.no_op(),) - try: - with ops.device(core(0)): - outputs = [ - o if isinstance(o, ops.Operation) else ops.convert_to_tensor(o) - for o in outputs - ] - except Exception as e: - raise ValueError( - "TPU function return values must all either be Operations or " - "convertible to Tensors. Got '%s'" % str(e)) - - # Separates the returned Operations and Tensors. - output_operations = [o for o in outputs if isinstance(o, ops.Operation)] - output_tensors = [o for o in outputs if not isinstance(o, ops.Operation)] - - if outputs != output_tensors + output_operations: - raise ValueError( - "TPU functions must return zero-or more Tensor values followed by " - "zero or more Operations.") - - # Wraps outputs in Identity ops. Otherwise a replicated input copied - # straight to an output would bypass the replicate(). This would be bad - # because the TPUReplicatedInput/TPUReplicatedOutput operator would not - # be rewritten away, leading to a runtime error. - # TODO(phawkins): extend the rewrite to elide these nodes instead. - new_output_tensors = [] - for t in output_tensors: - with ops.device(t.device if t.device else core(0)): - o = array_ops.identity(t) - # pylint: disable=protected-access - o.op._set_attr("_tpu_output_identity", attr_value_pb2.AttrValue(b=True)) - # pylint: enable=protected-access - new_output_tensors.append(o) - return new_output_tensors, output_operations - - -def _postprocess_non_flat_outputs(outputs): - """Validates non-flat outputs, add backs device assignments and other attrs. - - Args: - outputs: Output from `computation` inside `tpu.rewrite`. - - Returns: - Tensors extracted from outputs and an empty list because Operations are not - allowed in non-flat outputs.. - """ - - # Flatten output items. - flat_outputs = nest.flatten(outputs) - - # Convert all non-Operation outputs to Tensors. - for i, o in enumerate(flat_outputs): - if isinstance(o, ops.Operation): - raise ValueError( - "tpu.rewrite does not support Operation as return value in non-flat " - "output structure. You can set returned Operations as control " - "dependencies of returned Tensors so Operations are triggered when " - 'Tensors are evaluated. Operation found: "%s"' % o.name) - - try: - o = ops.convert_to_tensor(o) - except Exception as e: - raise ValueError( - "TPU function return values must all either be Operations or " - 'convertible to Tensors. Got error: "%s"' % str(e)) - - # Wraps outputs in Identity ops. Otherwise a replicated input copied - # straight to an output would bypass the replicate(). This would be bad - # because the TPUReplicatedInput/TPUReplicatedOutput operator would not - # be rewritten away, leading to a runtime error. - # TODO(phawkins): extend the rewrite to elide these nodes instead. - with ops.device(core(0)): - o = array_ops.identity(o) - # pylint: disable=protected-access - o.op._set_attr("_tpu_output_identity", attr_value_pb2.AttrValue(b=True)) - # pylint: enable=protected-access - flat_outputs[i] = array_ops.identity(o) - - # All flat_outputs are Tensors, and no Operations. - return flat_outputs, [] - - -def split_compile_and_shard(computation, - inputs=None, - num_shards=1, - input_shard_axes=None, - outputs_from_all_shards=True, - output_shard_axes=None, - infeed_queue=None, - device_assignment=None, - name=None): - """Shards `computation` for parallel execution. - - `inputs` must be a list of Tensors or None (equivalent to an empty list), each - of which has a corresponding split axis (from `input_shard_axes`). Each input - is split into `num_shards` pieces along the corresponding axis, and - computation is applied to each shard in parallel. - - Tensors are broadcast to all shards if they are lexically captured by - `computation`. e.g., - - x = tf.constant(7) - def computation(): - return x + 3 - ... = shard(computation, ...) - - If `outputs_from_all_shards` is true, the outputs from all shards of - `computation` are concatenated back together along their `output_shards_axes`. - Otherwise, each output is taken from an arbitrary shard. - - Inputs and outputs of the computation must be at least rank-1 Tensors. - - Args: - computation: A Python function that builds a computation to apply to each - shard of the input. - inputs: A list of input tensors or None (equivalent to an empty list). Each - input tensor has a corresponding shard axes, given by `input_shard_axes`, - which must have size divisible by `num_shards`. - num_shards: The number of shards. - input_shard_axes: A list of dimensions along which to shard `inputs`, or - `None`. `None` means "shard all inputs along dimension 0". If not `None`, - there must be one dimension per input. - outputs_from_all_shards: Boolean or list of boolean. For each output, if - `True`, outputs from all shards are concatenated along the corresponding - `output_shard_axes` entry. Otherwise, each output is taken - from an arbitrary shard. If the argument is a boolean, the argument's - value is used for each output. - output_shard_axes: A list of dimensions along which to concatenate the - outputs of `computation`, or `None`. `None` means "concatenate all outputs - along dimension 0". If not `None`, there must be one dimension per output. - Ignored if `outputs_from_all_shards` is False. - infeed_queue: If not `None`, the `InfeedQueue` to use to augment the inputs - of `computation`. - device_assignment: If not `None`, a `DeviceAssignment` describing the - mapping between logical cores in the computation with physical cores in - the TPU topology. Uses a default device assignment if `None`. The - `DeviceAssignment` may be omitted if each shard of the computation uses - only one core, and there is either only one shard, or the number of shards - is equal to the number of cores in the TPU system. - name: (Deprecated) Does nothing. - Returns: - A tuple of (compile op, [output tensors]). - Raises: - ValueError: If num_shards <= 0 - ValueError: If len(input_shard_axes) != len(inputs) - ValueError: If len(output_shard_axes) != len(outputs from `computation`) - """ - # TODO(phawkins): consider adding support for broadcasting Tensors passed as - # inputs. - - if num_shards <= 0: - raise ValueError("num_shards must be a positive integer.") - - inputs = [] if inputs is None else inputs - if not isinstance(inputs, list): - raise TypeError("tpu.shard()'s inputs must be a list of Tensors or None.") - - # Converts inputs to Tensors. - inputs = [ops.convert_to_tensor(x) for x in inputs] - - if input_shard_axes is None: - input_shard_axes = [0] * len(inputs) - if len(inputs) != len(input_shard_axes): - raise ValueError("Length of input_shard_axes must be equal to the number " - "of inputs.") - - if inputs: - # Splits the `inputs` along the corresponding `input_shard_axes`, giving - # lists with layout [input][shard] - split_inputs = [ - array_ops.split(x, num_shards, axis=axis) - for (axis, x) in zip(input_shard_axes, inputs)] - - # Transposes the input lists to have layout [shard][input] - transposed_inputs = [list(i) for i in zip(*split_inputs)] - else: - transposed_inputs = [[]] * num_shards - - compile_op, outputs = split_compile_and_replicate( - computation, - transposed_inputs, - infeed_queue=infeed_queue, - device_assignment=device_assignment, - name=name) - - # There must be at least one shard since num_shards > 0. - # TODO(b/36647078) remove disable when pylint bug is fixed. - # pylint: disable=indexing-exception - if isinstance(outputs[0], ops.Operation): - # pylint: enable=indexing-exception - # There were no outputs from the computation and replicate returned a list - # of NoOps with control dependencies on the computation. Return the first - # one so it can be used as a control dependency or fetch node. - # TODO(b/36647078) remove disable when pylint bug is fixed. - # pylint: disable=indexing-exception - return compile_op, [outputs[0]] - # pylint: enable=indexing-exception - - # TODO(b/36647078) remove disable when pylint bug is fixed. - # pylint: disable=indexing-exception - num_outputs = len(outputs[0]) - # pylint: enable=indexing-exception - - if output_shard_axes is None: - output_shard_axes = [0] * num_outputs - if num_outputs != len(output_shard_axes): - raise ValueError("Length of output_shard_axes must be equal to the number " - "of outputs.") - - if isinstance(outputs_from_all_shards, bool): - outputs_from_all_shards = [outputs_from_all_shards] * num_outputs - - if num_outputs != len(outputs_from_all_shards): - raise ValueError("Length of outputs_from_all_shards must be equal to the " - "number of outputs.") - - results = [] - for (axis, all_shards, x) in zip(output_shard_axes, outputs_from_all_shards, - zip(*outputs)): - if all_shards: - # Concatenate all of the outputs together (use stack for scalars). - shape = x[0].shape - is_scalar = shape is not None and (shape.ndims == 0) - results.append((array_ops.stack(list(x)) if is_scalar - else array_ops.concat(list(x), axis=axis))) - else: - # TODO(phawkins): use a smarter policy, e.g., round-robin across shards. - results.append(x[0]) - - return compile_op, results - - -def shard(computation, - inputs=None, - num_shards=1, - input_shard_axes=None, - outputs_from_all_shards=True, - output_shard_axes=None, - infeed_queue=None, - device_assignment=None, - name=None): - """Shards `computation` for parallel execution. - - `inputs` must be a list of Tensors or None (equivalent to an empty list), each - of which has a corresponding split axis (from `input_shard_axes`). Each input - is split into `num_shards` pieces along the corresponding axis, and - computation is applied to each shard in parallel. - - Tensors are broadcast to all shards if they are lexically captured by - `computation`. e.g., - - x = tf.constant(7) - def computation(): - return x + 3 - ... = shard(computation, ...) - - TODO(phawkins): consider adding support for broadcasting Tensors passed - as inputs. - - If `outputs_from_all_shards` is true, the outputs from all shards of - `computation` are concatenated back together along their `output_shards_axes`. - Otherwise, each output is taken from an arbitrary shard. - - Inputs and outputs of the computation must be at least rank-1 Tensors. - - Args: - computation: A Python function that builds a computation to apply to each - shard of the input. - inputs: A list of input tensors or None (equivalent to an empty list). Each - input tensor has a corresponding shard axes, given by `input_shard_axes`, - which must have size divisible by `num_shards`. - num_shards: The number of shards. - input_shard_axes: A list of dimensions along which to shard `inputs`, or - `None`. `None` means "shard all inputs along dimension 0". If not `None`, - there must be one dimension per input. - outputs_from_all_shards: Boolean or list of boolean. For each output, if - `True`, outputs from all shards are concatenated along the corresponding - `output_shard_axes` entry. Otherwise, each output is taken - from an arbitrary shard. If the argument is a boolean, the argument's - value is used for each output. - output_shard_axes: A list of dimensions along which to concatenate the - outputs of `computation`, or `None`. `None` means "concatenate all outputs - along dimension 0". If not `None`, there must be one dimension per output. - Ignored if `outputs_from_all_shards` is False. - infeed_queue: If not `None`, the `InfeedQueue` to use to augment the inputs - of `computation`. - device_assignment: If not `None`, a `DeviceAssignment` describing the - mapping between logical cores in the computation with physical cores in - the TPU topology. Uses a default device assignment if `None`. The - `DeviceAssignment` may be omitted if each shard of the computation uses - only one core, and there is either only one shard, or the number of shards - is equal to the number of cores in the TPU system. - name: (Deprecated) Does nothing. - Returns: - A list of output tensors. - Raises: - ValueError: If num_shards <= 0 - ValueError: If len(input_shard_axes) != len(inputs) - ValueError: If len(output_shard_axes) != len(outputs from `computation`) - """ - return split_compile_and_shard( - computation, - inputs=inputs, - num_shards=num_shards, - input_shard_axes=input_shard_axes, - outputs_from_all_shards=outputs_from_all_shards, - output_shard_axes=output_shard_axes, - infeed_queue=infeed_queue, - device_assignment=device_assignment, - name=name)[1] - - -def batch_parallel(computation, - inputs=None, - num_shards=1, - infeed_queue=None, - device_assignment=None, - name=None): - """Shards `computation` along the batch dimension for parallel execution. - - Convenience wrapper around shard(). - - `inputs` must be a list of Tensors or None (equivalent to an empty list). - Each input is split into `num_shards` pieces along the 0-th dimension, and - computation is applied to each shard in parallel. - - Tensors are broadcast to all shards if they are lexically captured by - `computation`. e.g., - - x = tf.constant(7) - def computation(): - return x + 3 - ... = shard(computation, ...) - - The outputs from all shards are concatenated back together along their 0-th - dimension. - - Inputs and outputs of the computation must be at least rank-1 Tensors. - - Args: - computation: A Python function that builds a computation to apply to each - shard of the input. - inputs: A list of input tensors or None (equivalent to an empty list). The - 0-th dimension of each Tensor must have size divisible by `num_shards`. - num_shards: The number of shards. - infeed_queue: If not `None`, the `InfeedQueue` from which to append a tuple - of arguments as inputs to `computation`. - device_assignment: If not `None`, a `DeviceAssignment` describing the - mapping between logical cores in the computation with physical cores in - the TPU topology. Uses a default device assignment if `None`. The - `DeviceAssignment` may be omitted if each shard of the computation uses - only one core, and there is either only one shard, or the number of shards - is equal to the number of cores in the TPU system. - name: (Deprecated) Does nothing. - Returns: - A list of output tensors. - Raises: - ValueError: If `num_shards <= 0` - """ - return shard( - computation, - inputs, - num_shards=num_shards, - infeed_queue=infeed_queue, - device_assignment=device_assignment, - name=name) - - -def rewrite(computation, - inputs=None, - infeed_queue=None, - device_assignment=None, - name=None): - """Rewrites `computation` for execution on a TPU system. - - Args: - computation: A Python function that builds a computation to apply to the - input. If the function takes n inputs, 'inputs' should be a list of n - tensors. - - `computation` may return a list of operations and tensors. Tensors must - come before operations in the returned list. The return value of - `rewrite` is a list of tensors corresponding to the tensors from the - output of `computation`. - - All `Operation`s constructed during `computation` will be executed when - evaluating any of the returned output tensors, not just the ones returned. - inputs: A list of input tensors or `None` (equivalent to an empty list). - Each input can be a nested structure containing values that are - convertible to tensors. Note that passing an N-dimension list of - compatible values will result in a N-dimention list of scalar tensors - rather than a single Rank-N tensors. If you need different behavior, - convert part of inputs to tensors with `tf.convert_to_tensor`. - infeed_queue: If not `None`, the `InfeedQueue` from which to append a tuple - of arguments as inputs to `computation`. - device_assignment: if not `None`, a `DeviceAssignment` describing the - mapping between logical cores in the computation with physical cores in - the TPU topology. May be omitted for a single-core computation, in which - case the core attached to task 0, TPU device 0 is used. - name: (Deprecated) Does nothing. - Returns: - Same data structure as if computation(*inputs) is called directly with some - exceptions for correctness. Exceptions include: - 1) None output: a NoOp would be returned which control-depends on - computation. - 2) Single value output: A tuple containing the value would be returned. - 3) Operation-only outputs: a NoOp would be returned which - control-depends on computation. - TODO(b/121383831): Investigate into removing these special cases. - """ - # TODO(b/36647078) remove disable when pylint bug is fixed. - # pylint: disable=indexing-exception - return replicate( - computation, - None if inputs is None else [inputs], - infeed_queue=infeed_queue, - device_assignment=device_assignment, - name=name)[0] - # pylint: enable=indexing-exception - - # Operations that indicate some error in the user's inference graph. -_BLACKLISTED_INFERENCE_OPS = set([ - "ReadVariableOp", - "AssignVariableOp", - "AssignAddVariableOp", - "AssignSubVariableOp", - "VarHandleOp", - "Variable", - "VariableV2", -]) - - -def under_tpu_inference_context(): - """Check if it is currently under `tpu.rewrite_for_inference()`.""" - graph = ops.get_default_graph() - - context = graph._get_control_flow_context() # pylint: disable=protected-access - while context: - if isinstance(context, _TPUInferenceContext): - return True - context = context.outer_context - - return False - - -class _TPUInferenceContext(control_flow_ops.XLAControlFlowContext): - """A `ControlFlowContext` for nodes inside a TPU inference computation. - - The primary role of `TPUReplicateContext` is to sanity check operators inside - a tpu.rewrite_for_inference() computation. - """ - - def __init__(self, name): - super(_TPUInferenceContext, self).__init__() - self._name = name - - def AddOp(self, op): - self._AddOpInternal(op) - - def _AddOpInternal(self, op): - # pylint: disable=protected-access - if op.type in _BLACKLISTED_INFERENCE_OPS: - raise NotImplementedError( - "Operation of type %s (%s) is not supported on the TPU for inference." - " Execution will fail if this op is used in the graph. Make sure your" - " variables are using variable_scope." % (op.type, op.name)) - if self._outer_context: - self._outer_context.AddInnerOp(op) - - def AddValue(self, val): - result = val - if self._outer_context: - result = self._outer_context.AddValue(val) - return result - - def AddInnerOp(self, op): - self._AddOpInternal(op) - - @property - def grad_state(self): - return None - - -@experimental -def validate_inference_rewrite_for_variables(graph): - """Validates whether rewrite_for_inference() 'worked' for variables. - - The rewrite_for_inference() method is supposed to append GuaranteeConstOps - after ReadVariableOps, but this mechanism works only if you are using - tf.get_variable() to create and access variables in your tpu computation. - This validation method can be called immediately after calling - tpu.rewrite_for_inference() to check whether GuaranteeConstOps where added - to the graph. - - Typical usages: - tpu.validate_inference_rewrite_for_variables(tf.get_default_graph()) - - tpu.validate_inference_rewrite_for_variables(sess.graph) - - Args: - graph: The graph which needs to be validated. - Raises: - RuntimeError: if validation failed. - """ - if not any(x.type == "GuaranteeConst" for x in graph.get_operations()): - raise RuntimeError( - "No GuaranteeConst ops found in the graph after running " - "tpu.rewrite_for_inference(...). Please check that you are using " - "tf.get_variable() to create and access variables in your tpu " - "computation.") - - -@experimental -def rewrite_for_inference(computation, - inputs=None, - infeed_queue=None, - device_assignment=None, - name=None): - """Rewrites `computation` for inference on a TPU system. - - Other than 'rewriting' the computation to run on a TPU, if using variables - in your computation, it moves the ReadVariableOps outside the TPU - computation, and adds GuaranteeConst ops just after the ReadVariableOps. - This mechanism works only if you are using tf.get_variable() to create and - access variables in your tpu computation. You can validate whether this - worked, by calling validate_inference_rewrite_for_variables() method - immediately after this method to check whether GuaranteeConstOps where - added to the graph. - - Args: - computation: A Python function that builds a computation to apply to the - input. If the function takes n inputs, 'inputs' should be a list of n - tensors. If the function returns m outputs, rewrite will return a list of - m tensors. - inputs: A list of input tensors or `None` (equivalent to an empty list). - infeed_queue: If not `None`, the `InfeedQueue` from which to append a tuple - of arguments as inputs to `computation`. - device_assignment: if not `None`, a `DeviceAssignment` describing the - mapping between logical cores in the computation with physical cores in - the TPU topology. May be omitted for a single-core computation, in which - case the core attached to task 0, TPU device 0 is used. - name: The name of the operator. - Returns: - A list of output tensors. - """ - - def guarantee_const_getter(getter, name, *args, **kwargs): - with ops.control_dependencies(None): - return array_ops.guarantee_const( - getter(name, *args, **kwargs), name=name + "/GuaranteeConst") - - def wrapped_computation(*args, **kwargs): - """Execute computation under `_TPUInferenceContext`.""" - context = _TPUInferenceContext( - name=ops.get_default_graph().unique_name("rewrite_for_inference")) - try: - context.Enter() - - vscope = variable_scope.get_variable_scope() - prev_custom_getter = vscope.custom_getter - prev_caching_device = vscope.caching_device - vscope.set_custom_getter(guarantee_const_getter) - vscope.set_caching_device(lambda op: op.device) - - result = computation(*args, **kwargs) - - vscope.set_custom_getter(prev_custom_getter) - vscope.set_caching_device(prev_caching_device) - finally: - context.Exit() - return result - - # pylint: disable=undefined-variable - return rewrite( - wrapped_computation, - inputs=inputs, - infeed_queue=infeed_queue, - device_assignment=device_assignment, - name=name) - # pylint: enable=undefined-variable +# pylint: disable=wildcard-import,unused-import,redefined-builtin +from tensorflow.python.tpu.tpu import * +# used by tests +from tensorflow.python.tpu.tpu import _TPU_REPLICATE_ATTR +# pylint: enable=wildcard-import,unused-import,redefined-builtin diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_config.py b/tensorflow/contrib/tpu/python/tpu/tpu_config.py index 82ebb9a37f22db0c150481d3e0ce3877f882d245..c36aaa38c0e4823bfc438773e4aa5b5109794da4 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_config.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_config.py @@ -1,276 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# =================================================================== - -"""A RunConfig subclass with TPU support.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections -import json -import os - -from tensorflow.contrib.tpu.python.tpu import util as util_lib -from tensorflow.core.protobuf import config_pb2 -from tensorflow.python.estimator import run_config as run_config_lib -from tensorflow.python.platform import tf_logging as logging - -# pylint: disable=protected-access -_TF_CONFIG_ENV = run_config_lib._TF_CONFIG_ENV -_SERVICE_KEY = run_config_lib._SERVICE_KEY -_TPU_WORKER_JOB_NAME = 'tpu_worker_job_name' -# pylint: enable=protected-access - - -class InputPipelineConfig(object): - r"""Please see the definition of these values in TPUConfig.""" - PER_SHARD_V1 = 1 - PER_HOST_V1 = 2 - PER_HOST_V2 = 3 - BROADCAST = 4 - - -class TPUConfig( - collections.namedtuple('TPUConfig', [ - 'iterations_per_loop', - 'num_shards', - 'num_cores_per_replica', - 'per_host_input_for_training', - 'tpu_job_name', - 'initial_infeed_sleep_secs', - 'input_partition_dims', - ])): - r"""TPU related configuration required by `TPUEstimator`. - - Args: - iterations_per_loop: This is the number of train steps running in TPU - system before returning to CPU host for each `Session.run`. This means - global step is increased `iterations_per_loop` times in one `Session.run`. - It is recommended to be set as number of global steps for next checkpoint. - num_shards: (Deprecated, ignored by TPUEstimator). - The number of model replicas in the system. For non-model-parallelism - case, this number equals the total number of TPU cores. For - model-parallelism, the total number of TPU cores equals - num_cores_per_replica * num_shards. - num_cores_per_replica: Defaults to `None`, which disables model parallelism. - An integer which describes the number of TPU cores per model replica. This - is required by model-parallelism which enables partitioning - the model to multiple cores. Currently num_cores_per_replica must be - 1, 2, 4, or 8. - per_host_input_for_training: If `True`, `PER_HOST_V1`, or `PER_HOST_V2`, - `input_fn` is invoked once on each host. With the per-core input pipeline - configuration, it is invoked once for each core. - With a global batch size `train_batch_size` in `TPUEstimator` constructor, - the batch size for each shard is `train_batch_size` // #hosts in the - `True` or `PER_HOST_V1` mode. In `PER_HOST_V2` mode, it is - `train_batch_size` // #cores. In `BROADCAST` mode, `input_fn` is only - invoked once on host 0 and the tensors are broadcasted to all other - replicas. The batch size equals to train_batch_size`. With the per-core - input pipeline configuration, the shard batch size is also - `train_batch_size` // #cores. - Note: per_host_input_for_training==PER_SHARD_V1 only supports mode.TRAIN. - tpu_job_name: The name of the TPU job. Typically, this name is auto-inferred - within TPUEstimator, however when using ClusterSpec propagation in more - esoteric cluster configurations, you may need to specify the job name as a - string. - initial_infeed_sleep_secs: The number of seconds the infeed thread should - wait before enqueueing the first batch. This helps avoid timeouts for - models that require a long compilation time. - input_partition_dims: A nested list to describe the partition dims - for all the tensors from input_fn(). The structure of - input_partition_dims must match the structure of `features` and - `labels` from input_fn(). The total number of partitions must match - `num_cores_per_replica`. For example, if input_fn() returns two tensors: - images with shape [N, H, W, C] and labels [N]. - input_partition_dims = [[1, 2, 2, 1], None] will split the images to 4 - pieces and feed into 4 TPU cores. labels tensor are directly broadcasted - to all the TPU cores since the partition dims is `None`. - Current limitations: This feature is only supported with the PER_HOST_V2 - input mode. - - Raises: - ValueError: If `num_cores_per_replica` is not 1, 2, 4, 8 or 16. - """ - - def __new__(cls, - iterations_per_loop=2, - num_shards=None, - num_cores_per_replica=None, - per_host_input_for_training=True, - tpu_job_name=None, - initial_infeed_sleep_secs=None, - input_partition_dims=None): - - # Check iterations_per_loop. - util_lib.check_positive_integer(iterations_per_loop, - 'TPUConfig iterations_per_loop') - - # Check num_shards. - if num_shards is not None: - util_lib.check_positive_integer(num_shards, 'TPUConfig num_shards') - - if input_partition_dims is not None: - if len(input_partition_dims) != 1 and len(input_partition_dims) != 2: - raise ValueError( - 'input_partition_dims must be a list/tuple with one or two' - ' elements.') - - if per_host_input_for_training is not InputPipelineConfig.PER_HOST_V2: - raise ValueError( - 'input_partition_dims is only supported in PER_HOST_V2 mode.') - - if num_cores_per_replica is None: - raise ValueError( - 'input_partition_dims requires setting num_cores_per_replica.') - - # Check num_cores_per_replica - if num_cores_per_replica is not None: - if num_cores_per_replica not in [1, 2, 4, 8, 16]: - raise ValueError( - 'num_cores_per_replica must be 1, 2, 4, 8, or 16; got {}'.format( - str(num_cores_per_replica))) - - # per_host_input_for_training may be True, False, or integer in [1..3]. - # Map legacy values (True, False) to numeric values. - if per_host_input_for_training is False: - per_host_input_for_training = InputPipelineConfig.PER_SHARD_V1 - elif per_host_input_for_training is True: - per_host_input_for_training = InputPipelineConfig.PER_HOST_V1 - - # Check initial_infeed_sleep_secs. - if initial_infeed_sleep_secs: - util_lib.check_positive_integer(initial_infeed_sleep_secs, - 'TPUConfig initial_infeed_sleep_secs') - - tpu_job_name = tpu_job_name or _get_tpu_job_name_from_tf_config() - - return super(TPUConfig, cls).__new__( - cls, - iterations_per_loop=iterations_per_loop, - num_shards=num_shards, - num_cores_per_replica=num_cores_per_replica, - per_host_input_for_training=per_host_input_for_training, - tpu_job_name=tpu_job_name, - initial_infeed_sleep_secs=initial_infeed_sleep_secs, - input_partition_dims=input_partition_dims) - - -class RunConfig(run_config_lib.RunConfig): - """RunConfig with TPU support.""" - - def __init__(self, - tpu_config=None, - evaluation_master=None, - master=None, - cluster=None, - **kwargs): - """Constructs a RunConfig. - - Args: - tpu_config: the TPUConfig that specifies TPU-specific configuration. - evaluation_master: a string. The address of the master to use for eval. - Defaults to master if not set. - master: a string. The address of the master to use for training. - cluster: a ClusterResolver - **kwargs: keyword config parameters. - - Raises: - ValueError: if cluster is not None and the provided session_config has a - cluster_def already. - """ - super(RunConfig, self).__init__(**kwargs) - self._tpu_config = tpu_config or TPUConfig() - self._cluster = cluster - - # If user sets master and/or evaluation_master explicitly, including empty - # string '', take it. Otherwise, take the values set by parent class. - if master is not None: - if cluster is not None: - raise ValueError('Both master and cluster are set.') - self._master = master - else: - if cluster: - self._master = cluster.master() - - if evaluation_master is not None: - self._evaluation_master = evaluation_master - elif (not self._evaluation_master and - self.task_type != run_config_lib.TaskType.EVALUATOR): - # If the task type is EVALUATOR, it means some cluster manager sets the - # TF_CONFIG. In that case, we respect the configuration in TF_CONFIG. - # - # Otherwise, it means user executes the code without external cluster - # manager. For that, we optimize the user experience by setting - # evaluation_master to master, unless user overwrites it. - self._evaluation_master = self._master - - # Set the ClusterSpec to use - if cluster: - self._cluster_spec = cluster.cluster_spec() - - # Merge the cluster_def into the ConfigProto. - if self._session_config is None: # pylint: disable=access-member-before-definition - self._session_config = config_pb2.ConfigProto( - allow_soft_placement=True, isolate_session_state=True) - if self._session_config.HasField('cluster_def'): - raise ValueError( - 'You cannot provide a ClusterResolver and ' - 'session_config.cluster_def.') - if self._cluster_spec: - self._session_config.cluster_def.CopyFrom( - self._cluster_spec.as_cluster_def()) - - def _maybe_overwrite_session_config_for_distributed_training(self): - # Overrides the parent class session_config overwrite for between-graph. TPU - # runs with in-graph, which should not have device filter. Doing nothing - # ("pass") basically disables it. - pass - - @property - def evaluation_master(self): - return self._evaluation_master - - @property - def master(self): - return self._master - - @property - def tpu_config(self): - return self._tpu_config - - @property - def cluster(self): - return self._cluster - - def replace(self, **kwargs): - if 'tpu_config' not in kwargs: - return super(RunConfig, self).replace(**kwargs) - - tpu_config = kwargs.pop('tpu_config') - new_instance = super(RunConfig, self).replace(**kwargs) - new_instance._tpu_config = tpu_config # pylint: disable=protected-access - return new_instance - - -def _get_tpu_job_name_from_tf_config(): - """Extracts the TPU job name from TF_CONFIG env variable.""" - # TODO(xiejw): Extends this to support both TF_CONFIG env variable and cluster - # spec propagation. - tf_config = json.loads(os.environ.get(_TF_CONFIG_ENV, '{}')) - tpu_job_name = tf_config.get(_SERVICE_KEY, {}).get(_TPU_WORKER_JOB_NAME) - if tpu_job_name: - logging.info('Load TPU job name from TF_CONFIG: %s', tpu_job_name) - return tpu_job_name +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.tpu_config import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_context.py b/tensorflow/contrib/tpu/python/tpu/tpu_context.py index ed1e0f0401a96c34e6ff9323685857b64e10bd14..b77b010cba6bf32c3b6d170bc522eebfb6a04f77 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_context.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_context.py @@ -1,763 +1,23 @@ -# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# =================================================================== -"""TPU system metadata and associated tooling.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -from contextlib import contextmanager -import copy - -from tensorflow.contrib.tpu.python.tpu import _tpu_estimator_embedding -from tensorflow.contrib.tpu.python.tpu import device_assignment as tpu_device_assignment -from tensorflow.contrib.tpu.python.tpu import tpu_config -from tensorflow.contrib.tpu.python.tpu import tpu_system_metadata as tpu_system_metadata_lib -from tensorflow.python.estimator import model_fn as model_fn_lib -from tensorflow.python.platform import tf_logging as logging - - -_DEFAULT_JOB_NAME = 'tpu_worker' -_DEFAULT_COORDINATOR_JOB_NAME = 'coordinator' -_LOCAL_MASTERS = ('', 'local') -_NUM_CORES_TO_COMPUTATION_SHAPE = { - 1: [1, 1, 1], - 2: [1, 1, 2], - 4: [1, 2, 2], - 8: [2, 2, 2], - 16: [4, 2, 2], -} - - -class TPUContext(object): - """A context that holds the current configuration of the TPU computation.""" - - def __init__(self, - internal_ctx, - input_device=None, - invocation_index=None, - call_from_input_fn=True): - self._internal_ctx = internal_ctx - self._input_device = input_device - self._invocation_index = invocation_index - self._call_from_input_fn = call_from_input_fn - - def current_input_fn_deployment(self): - """The configuration of the current input_fn invocation. - - The configuration depends on `TPUConfig.per_host_input_for_training`. See - `TPUConfig` for details. - - Only set in params dict of input_fn - - Returns: - A tuple of - 1. Device spec string: String, is the current CPU host where the - input_fn is invoked. - 2. Current invocation index: Int, 0-based index of the input_fn - invocation. See next item for details. - 3. Total invocation count: Int, the total number of times to invoke the - input_fn on all CPU hosts. Each invocation will be passed with a new - `TPUContext` instance with current invocation index set properly. - 4. Total number of replicas consumed by current_invocation: Int, the - number of replicas fed by the data returned by current input_fn. For - example, for per_core input pipeline deployment - and non-model-parallelism, total invocation count is equal to - the number of cores in the system and num replicas consumed by - current invocation is 1. For per-host v2 input pipeline deployment, - total invocation count is equal to the number of hosts in the system - and num replicas consumed by current invocation is equal to number of - cores per host. - - Raises: - RuntimeError: If this method must not be called from input_fn. - """ - if not self._call_from_input_fn: - raise RuntimeError('This TPUContext instance must not be called from' - ' model_fn.') - - if self._internal_ctx.is_input_sharded_per_core(): - total_invocation_count = (self._internal_ctx.num_hosts - * self._internal_ctx.num_of_replicas_per_host) - replicas_consumed = 1 - elif self._internal_ctx.is_input_broadcast_with_iterators(): - total_invocation_count = 1 - replicas_consumed = self._internal_ctx.num_replicas - else: - total_invocation_count = self._internal_ctx.num_hosts - replicas_consumed = self._internal_ctx.num_of_replicas_per_host - return (self._input_device, self._invocation_index, - total_invocation_count, replicas_consumed) - - @property - def num_replicas(self): - """The total number of replicas. - - For non-model-parallelism, num_replicas should be the total num of TPU - cores in the system. - - Returns: - The number of replicas. - """ - return self._internal_ctx.num_replicas - - @property - def num_hosts(self): - """The number of hosts for the TPU system.""" - return self._internal_ctx.num_hosts - - @property - def current_host(self): - """The current host index for the TPU system.""" - return self._invocation_index - - @property - def num_of_replicas_per_host(self): - """The number of replicas for each host.""" - if self._internal_ctx.model_parallelism_enabled: - raise ValueError( - 'num_of_replicas_per_host is not supported for model_parallelism') - return self._internal_ctx.num_of_replicas_per_host - - @property - def device_assignment(self): - """Returns device_assignment object.""" - if self._call_from_input_fn: - raise RuntimeError('This TPUContext instance must not be called from' - ' input_fn.') - return self._internal_ctx.device_assignment - - def device_for_replica(self, replica_id): - """Returns the tuple of (CPU device and device ordinal) for replica. - - This should be used for full replicate for non-model-parallelism. - - Args: - replica_id: Int, the replica index. - - Returns: - A tuple of device spec for CPU device and int device ordinal. - """ - # Note that: For the non-model parallelism, the mapping could be - # a random permutation. The order should not matter in most cases - # as far as model is replicated to all cores in the system. - return self._internal_ctx.device_for_replica(replica_id) - - @property - def tpu_host_placement_function(self): - """Returns the TPU host place function. - - The place function takes host_id as the input and returns the TF device - for the correspoding host. - """ - - def _placement_function(host_id): - """Return the host device given host_id.""" - return self._internal_ctx.tpu_host_placement_function(host_id=host_id) - - return _placement_function - - -class _InternalTPUContext(object): - """A context holds immutable states of TPU computation. - - This immutable object holds TPUEstimator config, train/eval batch size, and - `TPUEstimator.use_tpu`, which is expected to be passed around. It also - provides utility functions, based on the current state, to determine other - information commonly required by TPU computation, such as TPU device names, - TPU hosts, shard batch size, etc. - - if eval_on_tpu is False, then execution of eval on TPU is disabled. - if eval_on_tpu is True, but use_tpu is False, a warning is issued, - and TPU execution is disabled for all modes. - - N.B. As `mode` is not immutable state in Estimator, but essential to - distinguish between TPU training and evaluation, a common usage for - _InternalTPUContext with `mode` is as follows: - ``` - with _ctx.with_mode(mode) as ctx: - if ctx.is_running_on_cpu(): - ... - ``` - """ - - def __init__(self, - config, - train_batch_size, - eval_batch_size, - predict_batch_size, - use_tpu, - eval_on_tpu=True, - embedding_config_spec=None): - self._config = config - self._train_batch_size = train_batch_size - self._eval_batch_size = eval_batch_size - self._predict_batch_size = predict_batch_size - self._use_tpu = use_tpu - logging.info('_TPUContext: eval_on_tpu %s', eval_on_tpu) - if not use_tpu and eval_on_tpu: - logging.warning('eval_on_tpu ignored because use_tpu is False.') - - self._eval_on_tpu = eval_on_tpu - self._model_parallelism_enabled = ( - use_tpu and config.tpu_config.num_cores_per_replica) - self._mode = None - num_cores_per_replica = config.tpu_config.num_cores_per_replica - if self._model_parallelism_enabled: - self._computation_shape = _NUM_CORES_TO_COMPUTATION_SHAPE[ - num_cores_per_replica] - else: - self._computation_shape = None - self._lazy_tpu_system_metadata_dict = {} # key by master address - self._lazy_device_assignment_dict = {} # key by master address - self._lazy_validation_dict = {} # key by ModeKeys - self._embedding_config_spec = embedding_config_spec - self._lazy_embedding_config_dict = {} # key by master address - - def _assert_mode(self): - if self._mode is None: - raise RuntimeError( - '`mode` needs to be set via contextmanager `with_mode`.') - return self._mode - - @contextmanager - def with_mode(self, mode): - # NOTE(xiejw): Shallow copy is enough. It will share he lazy dictionaries, - # such as _lazy_tpu_system_metadata_dict between new copy and the original - # one. Note that all lazy states stored in properties _lazy_foo are sort of - # immutable as they should be same for the process lifetime. - new_ctx = copy.copy(self) - new_ctx._mode = mode # pylint: disable=protected-access - yield new_ctx - - @property - def mode(self): - return self._assert_mode() - - def _get_master_address(self): - mode = self._assert_mode() - config = self._config - master = ( - config.master - if mode != model_fn_lib.ModeKeys.EVAL else config.evaluation_master) - return master - - def _get_tpu_system_metadata(self): - """Gets the (maybe cached) TPU system metadata.""" - master = self._get_master_address() - tpu_system_metadata = self._lazy_tpu_system_metadata_dict.get(master) - if tpu_system_metadata is not None: - return tpu_system_metadata - - cluster_def = None - if (self._config.session_config and - self._config.session_config.cluster_def.job): - cluster_def = self._config.session_config.cluster_def - - # pylint: disable=protected-access - tpu_system_metadata = ( - tpu_system_metadata_lib._query_tpu_system_metadata( - master, - cluster_def=cluster_def, - query_topology=self.model_parallelism_enabled)) - - self._lazy_tpu_system_metadata_dict[master] = tpu_system_metadata - return tpu_system_metadata - - def _get_device_assignment(self): - """Gets the (maybe cached) TPU device assignment.""" - master = self._get_master_address() - device_assignment = self._lazy_device_assignment_dict.get(master) - if device_assignment is not None: - return device_assignment - - tpu_system_metadata = self._get_tpu_system_metadata() - - device_assignment = tpu_device_assignment.device_assignment( - tpu_system_metadata.topology, - computation_shape=self._computation_shape, - num_replicas=self.num_replicas) - - logging.info('num_cores_per_replica: %s', - str(self._config.tpu_config.num_cores_per_replica)) - logging.info('computation_shape: %s', str(self._computation_shape)) - logging.info('num_replicas: %d', self.num_replicas) - logging.info('device_assignment.topology.device_coordinates: %s', - str(device_assignment.topology.device_coordinates)) - logging.info('device_assignment.core_assignment: %s', - str(device_assignment.core_assignment)) - - self._lazy_device_assignment_dict[master] = device_assignment - return device_assignment - - @property - def embedding_config(self): - """Returns the embedding config based on current mode.""" - master = self._get_master_address() - if master in self._lazy_embedding_config_dict: - embedding_config = self._lazy_embedding_config_dict[master] - else: - embedding_config = None - if self._use_tpu and self._embedding_config_spec: - embedding_config = _tpu_estimator_embedding.EmbeddingConfig( - self._embedding_config_spec, self._train_batch_size, - self._eval_batch_size, self.num_hosts, self.num_cores, master) - if not embedding_config.has_embedding_tables(): - embedding_config = None - self._lazy_embedding_config_dict[master] = embedding_config - - if embedding_config is not None: - mode = self._assert_mode() - # Dynamically attach tpu_embedding based on mode. With - # this, we could keep embedding_config immutable but call site always - # accesses the unified API '.tpu_embedding'. - embedding_config.tpu_embedding = embedding_config.get_tpu_embedding(mode) - return embedding_config - - @property - def model_parallelism_enabled(self): - return self._model_parallelism_enabled - - @property - def input_partition_dims(self): - return self._config.tpu_config.input_partition_dims - - @property - def device_assignment(self): - return (self._get_device_assignment() - if self._model_parallelism_enabled else None) - - @property - def num_of_cores_per_host(self): - metadata = self._get_tpu_system_metadata() - return metadata.num_of_cores_per_host - - @property - def num_cores(self): - metadata = self._get_tpu_system_metadata() - return metadata.num_cores - - @property - def num_of_replicas_per_host(self): - """Return the number of replicas per host.""" - if self.model_parallelism_enabled: - return self.num_replicas // self.num_hosts - else: - return self.num_of_cores_per_host - - @property - def num_replicas(self): - num_cores_in_system = self.num_cores - - if self.model_parallelism_enabled: - num_cores_per_replica = self._config.tpu_config.num_cores_per_replica - if num_cores_per_replica > num_cores_in_system: - raise ValueError( - 'The num of cores required by the model parallelism, specified by ' - 'TPUConfig.num_cores_per_replica, is larger than the total num of ' - 'TPU cores in the system. num_cores_per_replica: {}, num cores ' - 'in the system: {}'.format(num_cores_per_replica, - num_cores_in_system)) - - if num_cores_in_system % num_cores_per_replica != 0: - raise RuntimeError( - 'The num of cores in the system ({}) is not divisible by the num ' - 'of cores ({}) required by the model parallelism, specified by ' - 'TPUConfig.num_cores_per_replica. This should never happen!'.format( - num_cores_in_system, num_cores_per_replica)) - - return num_cores_in_system // num_cores_per_replica - else: - return num_cores_in_system - - @property - def num_hosts(self): - metadata = self._get_tpu_system_metadata() - return metadata.num_hosts - - @property - def config(self): - return self._config - - def is_input_sharded_per_core(self): - """Return true if input_fn is invoked per-core (other than per-host).""" - mode = self._assert_mode() - return (mode == model_fn_lib.ModeKeys.TRAIN and - (self._config.tpu_config.per_host_input_for_training is - tpu_config.InputPipelineConfig.PER_SHARD_V1)) - - def is_input_per_host_with_iterators(self): - """Return true if input_fn should be run in the per-host v2 config.""" - return (self._config.tpu_config.per_host_input_for_training is - tpu_config.InputPipelineConfig.PER_HOST_V2) - - def is_input_broadcast_with_iterators(self): - """Return true if input_fn should be run in the full_replicae config.""" - return (self._config.tpu_config.per_host_input_for_training is - tpu_config.InputPipelineConfig.BROADCAST) - - def is_running_on_cpu(self, is_export_mode=False): - """Determines whether the input_fn and model_fn should be invoked on CPU. - - This API also validates user provided configuration, such as batch size, - according the lazy initialized TPU system metadata. - - Args: - is_export_mode: Indicates whether the current mode is for exporting the - model, when mode == PREDICT. Only with this bool, we could - tell whether user is calling the Estimator.predict or - Estimator.export_savedmodel, which are running on TPU and CPU - respectively. Parent class Estimator does not distinguish these two. - - Returns: - bool, whether current input_fn or model_fn should be running on CPU. - - Raises: - ValueError: any configuration is invalid. - """ - - is_running_on_cpu = self._is_running_on_cpu(is_export_mode) - if not is_running_on_cpu: - self._validate_tpu_configuration() - return is_running_on_cpu - - def _is_running_on_cpu(self, is_export_mode): - """Determines whether the input_fn and model_fn should be invoked on CPU.""" - mode = self._assert_mode() - - if not self._use_tpu: - return True - - if mode == model_fn_lib.ModeKeys.EVAL and not self._eval_on_tpu: - logging.info('_is_running_on_cpu: eval_on_tpu disabled') - return True - - if is_export_mode: - return True - - return False - - @property - def global_batch_size(self): - mode = self._assert_mode() - if mode == model_fn_lib.ModeKeys.TRAIN: - return self._train_batch_size - elif mode == model_fn_lib.ModeKeys.EVAL: - return self._eval_batch_size - elif mode == model_fn_lib.ModeKeys.PREDICT: - return self._predict_batch_size - else: - return None - - @property - def batch_size_for_input_fn(self): - """Returns the shard batch size for `input_fn`.""" - global_batch_size = self.global_batch_size - - if (self.is_running_on_cpu() or self.is_input_broadcast_with_iterators()): - return global_batch_size - - # On TPU - if self.is_input_sharded_per_core() or ( - self.is_input_per_host_with_iterators()): - return global_batch_size // self.num_replicas - else: - return global_batch_size // self.num_hosts - - @property - def batch_size_for_model_fn(self): - """Returns the shard batch size for `model_fn`.""" - global_batch_size = self.global_batch_size - - if (self.is_running_on_cpu() or self.is_input_broadcast_with_iterators()): - return global_batch_size - - # On TPU. always sharded per shard. - return global_batch_size // self.num_replicas - - @property - def master_job(self): - """Returns the job name to use to place TPU computations on. - - Returns: - A string containing the job name, or None if no job should be specified. - - Raises: - ValueError: If the user needs to specify a tpu_job_name, because we are - unable to infer the job name automatically, or if the user-specified job - names are inappropriate. - """ - run_config = self._config - # If the user specifies the tpu_job_name, use that. - if run_config.tpu_config.tpu_job_name: - return run_config.tpu_config.tpu_job_name - - # The tpu job is determined by the run_config. Right now, this method is - # required as tpu_config is not part of the RunConfig. - mode = self._assert_mode() - master = ( - run_config.evaluation_master - if mode == model_fn_lib.ModeKeys.EVAL else run_config.master) - if master in _LOCAL_MASTERS: - return None - - if (not run_config.session_config or - not run_config.session_config.cluster_def.job): - return _DEFAULT_JOB_NAME - cluster_def = run_config.session_config.cluster_def - job_names = set([job.name for job in cluster_def.job]) - if _DEFAULT_JOB_NAME in job_names: - # b/37868888 tracks allowing ClusterSpec propagation to reuse job names. - raise ValueError('Currently, tpu_worker is not an allowed job name.') - if len(job_names) == 1: - return cluster_def.job[0].name - if len(job_names) == 2: - if _DEFAULT_COORDINATOR_JOB_NAME in job_names: - job_names.remove(_DEFAULT_COORDINATOR_JOB_NAME) - return job_names.pop() - # TODO(b/67716447): Include more sophisticated heuristics. - raise ValueError( - 'Could not infer TPU job name. Please specify a tpu_job_name as part ' - 'of your TPUConfig.') - - @property - def tpu_host_placement_function(self): - """Returns the TPU host place function.""" - - master = self.master_job - - def _placement_function(_sentinal=None, replica_id=None, host_id=None): # pylint: disable=invalid-name - """Return the host device given replica_id or host_id.""" - assert _sentinal is None - if replica_id is not None and host_id is not None: - raise RuntimeError( - 'replica_id and host_id can have only one non-None value.') - - if master is None: - return '/replica:0/task:0/device:CPU:0' - else: - if replica_id is not None: - if self.model_parallelism_enabled: - return self.device_assignment.host_device( - replica=replica_id, job=master) - else: - host_id = replica_id / self.num_of_cores_per_host - - return '/job:%s/task:%d/device:CPU:0' % (master, host_id) - - return _placement_function - - @property - def tpu_device_placement_function(self): - """Returns a TPU device placement Fn.""" - master = self.master_job - job_device = '' if master is None else ('/job:%s' % master) - - def _placement_function(i): - if self.model_parallelism_enabled: - return self.device_assignment.tpu_device(replica=i, job=master) - else: - num_of_cores_per_host = self.num_of_cores_per_host - host_id = i / num_of_cores_per_host - ordinal_id = i % num_of_cores_per_host - return '%s/task:%d/device:TPU:%d' % (job_device, host_id, ordinal_id) - - return _placement_function - - def tpu_ordinal_function(self, host_id): - """Returns the TPU ordinal fn.""" - - def _tpu_ordinal_function(shard_index_in_host): - """Return the TPU ordinal associated with a shard. - - Required because the enqueue ops are placed on CPU. - - Args: - shard_index_in_host: the shard index - - Returns: - The ordinal of the TPU device the shard's infeed should be placed on. - """ - if self.model_parallelism_enabled: - # We put both enqueue/dequeue ops at tpu.core(0) in each replica. - replica = self.device_assignment.lookup_replicas(host_id, - 0)[shard_index_in_host] - return self.device_assignment.tpu_ordinal(replica=replica) - else: - return shard_index_in_host % self.num_of_cores_per_host - - return _tpu_ordinal_function - - def _validate_tpu_configuration(self): - """Validates the configuration based on the TPU system metadata.""" - mode = self._assert_mode() - if self._lazy_validation_dict.get(mode): - return - - # All following information is obtained from TPU system metadata. - num_cores = self.num_cores - num_replicas = self.num_replicas - num_hosts = self.num_hosts - - if not num_cores: - tpu_system_metadata = self._get_tpu_system_metadata() - raise RuntimeError( - 'Cannot find any TPU cores in the system. Please double check ' - 'Tensorflow master address and TPU worker(s). Available devices ' - 'are {}.'.format(tpu_system_metadata.devices)) - - if self._config.tpu_config.num_shards: - user_provided_num_replicas = self._config.tpu_config.num_shards - if user_provided_num_replicas != num_replicas: - message = ( - 'TPUConfig.num_shards is not set correctly. According to TPU ' - 'system metadata for Tensorflow master ({}): num_replicas should ' - 'be ({}), got ({}). For non-model-parallelism, num_replicas should ' - 'be the total num of TPU cores in the system. For ' - 'model-parallelism, the total number of TPU cores should be ' - 'num_cores_per_replica * num_replicas. Please set it ' - 'accordingly or leave it as `None`'.format( - self._get_master_address(), num_replicas, - user_provided_num_replicas)) - - raise ValueError(message) - - if self._config.tpu_config.num_cores_per_replica: - num_cores_per_replica = self._config.tpu_config.num_cores_per_replica - num_cores_per_host = self._get_tpu_system_metadata().num_of_cores_per_host - if num_cores_per_replica > num_cores_per_host: - raise ValueError( - 'The num of cores required by the model parallelism, specified by ' - 'TPUConfig.num_cores_per_replica, is larger than the ' - 'num_cores_per_host. num_cores_per_replica: {}, ' - 'num_cores_per_host: {}'.format(num_cores_per_replica, - num_cores_per_host)) - - if mode == model_fn_lib.ModeKeys.TRAIN: - if (self._train_batch_size % num_replicas != 0 and - not self.is_input_broadcast_with_iterators()): - raise ValueError( - 'train batch size {} must be divisible by number of replicas {}' - .format(self._train_batch_size, num_replicas)) - - elif mode == model_fn_lib.ModeKeys.EVAL: - if self._eval_batch_size is None: - raise ValueError( - 'eval_batch_size in TPUEstimator constructor cannot be `None`' - 'if .evaluate is running on TPU.') - if (self._eval_batch_size % num_replicas != 0 and - not self.is_input_broadcast_with_iterators()): - raise ValueError( - 'eval batch size {} must be divisible by number of replicas {}' - .format(self._eval_batch_size, num_replicas)) - if num_hosts > 1 and not self.is_input_broadcast_with_iterators(): - raise ValueError( - 'TPUEstimator.evaluate should be running on single TPU' - ' instead of a Pod.') - else: - assert mode == model_fn_lib.ModeKeys.PREDICT - if self._predict_batch_size is None: - raise ValueError( - 'predict_batch_size in TPUEstimator constructor should not be ' - '`None` if .predict is running on TPU.') - if (self._predict_batch_size % num_replicas != 0 and - not self.is_input_broadcast_with_iterators()): - raise ValueError( - 'predict batch size {} must be divisible by number of replicas {}' - .format(self._predict_batch_size, num_replicas)) - if num_hosts > 1 and not self.is_input_broadcast_with_iterators(): - raise ValueError( - 'TPUEstimator.predict should be running on single TPU worker. ' - 'got {}.'.format(num_hosts)) - - # Record the state "validated" into lazy dictionary. - self._lazy_validation_dict[mode] = True - - def device_for_replica(self, replica_id): - """Returns the tuple of (CPU device and device ordinal) for replica. - - This should be used for full replicate for non-model-parallelism. - - Args: - replica_id: Int, the replica index. - - Returns: - A tuple of device spec for CPU device and int device ordinal. - """ - master = self.master_job - - if self.model_parallelism_enabled: - return (self.device_assignment.host_device( - replica=replica_id, job=master), - self.device_assignment.tpu_ordinal(replica=replica_id)) - - job_device = '' if master is None else ('/job:%s' % master) - - num_of_replicas_per_host = self.num_of_replicas_per_host - host_id = replica_id / num_of_replicas_per_host - ordinal_id = replica_id % num_of_replicas_per_host - - host_device = '%s/task:%d/device:CPU:0' % (job_device, host_id) - return (host_device, ordinal_id) - - -class _OneCoreTPUContext(_InternalTPUContext): - """Special _InternalTPUContext for one core usage.""" - - def __init__(self, config, train_batch_size, eval_batch_size, - predict_batch_size, use_tpu): - - super(_OneCoreTPUContext, self).__init__( - config, train_batch_size, eval_batch_size, - predict_batch_size, use_tpu) - - def _get_tpu_system_metadata(self): - """Gets the (maybe cached) TPU system metadata.""" - master = self._get_master_address() - tpu_system_metadata = self._lazy_tpu_system_metadata_dict.get(master) - if tpu_system_metadata is not None: - return tpu_system_metadata - - tpu_system_metadata = ( - tpu_system_metadata_lib._TPUSystemMetadata( # pylint: disable=protected-access - num_cores=1, - num_hosts=1, - num_of_cores_per_host=1, - topology=None, - devices=[])) - - self._lazy_tpu_system_metadata_dict[master] = tpu_system_metadata - return tpu_system_metadata - - -def _get_tpu_context(config, train_batch_size, eval_batch_size, - predict_batch_size, use_tpu, eval_on_tpu, - embedding_config_spec): - """Returns an instance of `_InternalTPUContext`.""" - - if (config.tpu_config.num_shards == 1 and - config.tpu_config.num_cores_per_replica is None): - if embedding_config_spec is not None: - raise ValueError('Setting TPUConfig.num_shards==1 is unsupported ' - 'when embedding_config_spec is not None.') - logging.warning( - 'Setting TPUConfig.num_shards==1 is an unsupported behavior. ' - 'Please fix as soon as possible (leaving num_shards as None.)') - return _OneCoreTPUContext(config, train_batch_size, eval_batch_size, - predict_batch_size, use_tpu) - - return _InternalTPUContext(config, train_batch_size, eval_batch_size, - predict_batch_size, use_tpu, eval_on_tpu, - embedding_config_spec) +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.tpu_context import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_embedding.py b/tensorflow/contrib/tpu/python/tpu/tpu_embedding.py index 1ba8017cda834436cbcc72f03a1f8b88295bf80c..cb38a8f1a6bee3c2adfbefc203c1d143303c3368 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_embedding.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_embedding.py @@ -1,10 +1,10 @@ -# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, @@ -12,1097 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""TPU embedding APIs.""" +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections -import copy -import math -import re -import six - -from tensorflow.contrib.framework.python.framework import experimental -from tensorflow.contrib.tpu.python.ops import tpu_ops -from tensorflow.contrib.tpu.python.tpu import tpu_system_metadata as tpu_system_metadata_lib -from tensorflow.core.protobuf.tpu import optimization_parameters_pb2 -from tensorflow.core.protobuf.tpu import tpu_embedding_configuration_pb2 as elc -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops -from tensorflow.python.framework import sparse_tensor -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import control_flow_ops -from tensorflow.python.ops import init_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import partitioned_variables -from tensorflow.python.ops import state_ops -from tensorflow.python.ops import variable_scope - -TRAINING = elc.TPUEmbeddingConfiguration.TRAINING -INFERENCE = elc.TPUEmbeddingConfiguration.INFERENCE - - -class TableConfig( - collections.namedtuple( - 'TableConfig', - ['vocabulary_size', 'dimension', 'initializer', 'combiner'])): - """Embedding table configuration.""" - - @experimental - def __new__(cls, - vocabulary_size, - dimension, - initializer=None, - combiner='mean'): - """Embedding table configuration. - - Args: - vocabulary_size: Number of vocabulary (/rows) in the table. - dimension: The embedding dimension. - initializer: A variable initializer function to be used in embedding - variable initialization. If not specified, defaults to - `tf.truncated_normal_initializer` with mean `0.0` and standard deviation - `1/sqrt(dimension)`. - combiner: A string specifying how to reduce if there are multiple entries - in a single row. Currently 'mean', 'sqrtn' and 'sum' are supported, with - 'mean' the default. 'sqrtn' often achieves good accuracy, in particular - with bag-of-words columns. For more information, see - `tf.nn.embedding_lookup_sparse`. - - Returns: - `TableConfig`. - - Raises: - ValueError: if `vocabulary_size` is not positive integer. - ValueError: if `dimension` is not positive integer. - ValueError: if `initializer` is specified and is not callable. - ValueError: if `combiner` is not supported. - """ - if not isinstance(vocabulary_size, int) or vocabulary_size < 1: - raise ValueError('Invalid vocabulary_size {}.'.format(vocabulary_size)) - - if not isinstance(dimension, int) or dimension < 1: - raise ValueError('Invalid dimension {}.'.format(dimension)) - - if (initializer is not None) and (not callable(initializer)): - raise ValueError('initializer must be callable if specified.') - if initializer is None: - initializer = init_ops.truncated_normal_initializer( - mean=0.0, stddev=1 / math.sqrt(dimension)) - - if combiner not in ('mean', 'sum', 'sqrtn'): - raise ValueError('Invalid combiner {}'.format(combiner)) - - return super(TableConfig, cls).__new__(cls, vocabulary_size, dimension, - initializer, combiner) - - -AdamSlotVariableNames = collections.namedtuple( - 'AdamSlotVariableNames', ['m', 'v']) - -AdagradSlotVariableName = collections.namedtuple( - 'AdagradSlotVariableName', ['accumulator']) - -AdamSlotVariables = collections.namedtuple( - 'AdamSlotVariables', ['m', 'v']) - -AdagradSlotVariable = collections.namedtuple( - 'AdagradSlotVariable', ['accumulator']) - -VariablesAndOps = collections.namedtuple( - 'VariablesAndOps', - ['embedding_variables_by_table', 'slot_variables_by_table', - 'load_ops', 'retrieve_ops'] -) - - -# TODO(shizhiw): Factor `use_gradient_accumulation` and -# `pipeline_execution_with_tensor_core` out of `_OptimizationParameters`. -class _OptimizationParameters(object): - """Parameters common to all optimizations.""" - - def __init__(self, learning_rate, use_gradient_accumulation, - pipeline_execution_with_tensor_core): - self.learning_rate = learning_rate - self.use_gradient_accumulation = use_gradient_accumulation - self.pipeline_execution_with_tensor_core = ( - pipeline_execution_with_tensor_core) - - -class AdagradParameters(_OptimizationParameters): - """Optimization parameters for Adagrad.""" - - def __init__(self, learning_rate, initial_accumulator, - use_gradient_accumulation=False, - pipeline_execution_with_tensor_core=True): - """Optimization parameters for Adagrad. - - Args: - learning_rate: used for updating embedding table. - initial_accumulator: initial accumulator for Adagrad. - use_gradient_accumulation: setting this to `True` makes embedding - gradients calculation more accurate but slower. Please see - `optimization_parameters.proto` for details. - for details. - pipeline_execution_with_tensor_core: setting this to `True` makes training - faster, but trained model will be different if step N and step N+1 - involve the same set of embedding ID. Please see - `tpu_embedding_configuration.proto` for details. - """ - super(AdagradParameters, self).__init__(learning_rate, - use_gradient_accumulation, - pipeline_execution_with_tensor_core) - self.initial_accumulator = initial_accumulator - - -class AdamParameters(_OptimizationParameters): - """Optimization parameters for Adam.""" - - def __init__(self, learning_rate, - beta1=0.9, - beta2=0.999, - epsilon=1e-08, - lazy_adam=True, - sum_inside_sqrt=True, - use_gradient_accumulation=False, - pipeline_execution_with_tensor_core=True): - """Optimization parameters for Adam. - - Args: - learning_rate: a floating point value. The learning rate. - beta1: A float value. - The exponential decay rate for the 1st moment estimates. - beta2: A float value. - The exponential decay rate for the 2nd moment estimates. - epsilon: A small constant for numerical stability. - lazy_adam: Use lazy Adam instead of Adam. Lazy Adam trains faster. - Please see `optimization_parameters.proto` for details. - sum_inside_sqrt: This improves training speed. Please see - `optimization_parameters.proto` for details. - use_gradient_accumulation: setting this to `True` makes embedding - gradients calculation more accurate but slower. Please see - `optimization_parameters.proto` for details. - for details. - pipeline_execution_with_tensor_core: setting this to `True` makes training - faster, but trained model will be different if step N and step N+1 - involve the same set of embedding ID. Please see - `tpu_embedding_configuration.proto` for details. - """ - super(AdamParameters, self).__init__(learning_rate, - use_gradient_accumulation, - pipeline_execution_with_tensor_core) - self.beta1 = beta1 - self.beta2 = beta2 - self.epsilon = epsilon - self.lazy_adam = lazy_adam - self.sum_inside_sqrt = sum_inside_sqrt - - -class StochasticGradientDescentParameters(_OptimizationParameters): - """Optimization parameters for stochastic gradient descent. - - Args: - learning_rate: a floating point value. The learning rate. - use_gradient_accumulation: setting this to `True` makes embedding - gradients calculation more accurate but slower. Please see - `optimization_parameters.proto` for details. - pipeline_execution_with_tensor_core: setting this to `True` makes training - faster, but trained model will be different if step N and step N+1 - involve the same set of embedding ID. Please see - `tpu_embedding_configuration.proto` for details. - """ - - def __init__(self, learning_rate, use_gradient_accumulation=False, - pipeline_execution_with_tensor_core=True): - super(StochasticGradientDescentParameters, self).__init__( - learning_rate, use_gradient_accumulation, - pipeline_execution_with_tensor_core) - - -class TPUEmbedding(object): - """API for using TPU for embedding. - - Example: - ``` - table_config_user = tpu_embedding.TableConfig( - vocabulary_size=4, dimension=2, - initializer=initializer, combiner='mean') - table_to_config_dict = {'video': table_config_video, - 'user': table_config_user} - feature_to_table_dict = {'watched': 'video', - 'favorited': 'video', - 'friends': 'user'} - batch_size = 4 - num_hosts = 1 - optimization_parameters = tpu_embedding.AdagradParameters(1., 1.) - mode = tpu_embedding.TRAINING - embedding = tpu_embedding.TPUEmbedding( - table_to_config_dict, feature_to_table_dict, - batch_size, num_hosts, mode, optimization_parameters) - - batch_size_per_core = embedding.batch_size_per_core - sparse_features_list = [] - for host in hosts: - with ops.device(host): - for _ in range(embedding.num_cores_per_host): - sparse_features = {} - sparse_features['watched'] = sparse_tensor.SparseTensor(...) - sparse_features['favorited'] = sparse_tensor.SparseTensor(...) - sparse_features['friends'] = sparse_tensor.SparseTensor(...) - sparse_features_list.append(sparse_features) - - enqueue_ops = embedding.generate_enqueue_ops(sparse_features_list) - embedding_variables_and_ops = embedding.create_variables_and_ops() - - def computation(): - activations = embedding.get_activations() - loss = compute_loss(activations) - - base_optimizer = gradient_descent.GradientDescentOptimizer( - learning_rate=1) - cross_shard_optimizer = tpu_optimizer.CrossShardOptimizer( - base_optimizer) - - train_op = cross_shard_optimizer.minimize(loss) - gradients = ( - tpu_embedding_gradient.get_gradients_through_compute_gradients( - cross_shard_optimizer, loss, activations) - send_gradients_op = embedding.generate_send_gradients_op(gradients) - with ops.control_dependencies([train_op, send_gradients_op]): - loss = array_ops.identity(loss) - - loss = tpu.shard(computation, - num_shards=embedding.num_cores) - - with self.test_session() as sess: - sess.run(tpu.initialize_system(embedding_config= - embedding.config_proto)) - sess.run(variables.global_variables_initializer()) - sess.run(embedding_variables_and_ops.load_ops()) - sess.run(enqueue_ops) - loss_val = sess.run(loss) - ``` - """ - - # TODO(shizhiw): Instead of `feature_to_table_dict` which maps to table - # name, consider `feature_to_config_dict` which maps to `FeatureConfig`. - # `FeatureConfig` could have fields other than table name. For example, it - # could have a field to indicate that the feature should not be used to - # update embedding table (cr/204852758, cr/204940540). Also, this can support - # different combiners for different features within the same table. - # TODO(shizhiw, b/118512626): Remove `batch_size` from `__init__` and move it - # to `FeatureConfig`? - - # TODO(shizhiw): will it be cleaner to make `table_to_config_dict` and - # `feature_to_table_dict` lists of `TableSpec` and `FeatureSpec` respectively? - - # TODO(shizhiw): Consider adding `input_fn` as an option to remove boilerplate - # for-loops around construction of inputs. - - # `optimization_parameter` applies to all tables. If the need arises, - # we can add `optimization_parameters` to `TableConfig` to override this - # global setting. - @experimental - def __init__(self, - table_to_config_dict, - feature_to_table_dict, - batch_size, - mode, - master, - optimization_parameters=None): - """API for using TPU for embedding lookups. - - Args: - table_to_config_dict: A dictionary mapping from string of table name to - `TableConfig`. Table refers to an embedding table, e.g. `params` - argument to `tf.nn.embedding_lookup_sparse()`. - feature_to_table_dict: A dictionary mapping from string of feature name - to string of table name. Feature refers to ids to lookup in embedding - table, e.g. `sp_ids` argument to `tf.nn.embedding_lookup_sparse()`. - batch_size: An `int` representing the global batch size. - mode: `TRAINING` or `INFERENCE`. - master: A `string` representing the TensorFlow master to use. - optimization_parameters: `AdagradParameters`, `AdamParameters`, - `Stochasticgradientdescentparameters`. Must be set in training and must - be `None` in inference. - - Raises: - ValueError: if any input is invalid. - """ - _validate_table_to_config_dict(table_to_config_dict) - # Avoid nondeterminism from `Dict` iteration order by using `OrderedDict`. - self._table_to_config_dict = _create_ordered_dict(table_to_config_dict) - self._combiners = _create_combiners(self._table_to_config_dict) - - _validate_feature_to_table_dict(table_to_config_dict, feature_to_table_dict) - self._feature_to_table_dict = _create_ordered_dict(feature_to_table_dict) - self._table_to_features_dict = _create_table_to_features_dict( - self._feature_to_table_dict) - - self._batch_size = batch_size - - self._master = master - self._tpu_system_metadata = ( - tpu_system_metadata_lib._query_tpu_system_metadata(self._master)) # pylint: disable=protected-access - if self._tpu_system_metadata.num_cores == 0: - raise ValueError('TPUEmbedding needs TPUs, but master {} does not have ' - 'TPUs.'.format(self._master)) - self._num_hosts = self._tpu_system_metadata.num_hosts - self._hosts = [device.name for device in self._tpu_system_metadata.devices - if 'device:CPU:' in device.name] - self._num_cores_per_host = self._tpu_system_metadata.num_of_cores_per_host - self._num_cores = self._tpu_system_metadata.num_cores - - _validate_batch_size(self._batch_size, self._num_cores) - self._batch_size_per_core = self._batch_size // self._num_cores - - # TODO(shizhiw): remove `mode`? - if mode == TRAINING: - _validate_optimization_parameters(optimization_parameters) - self._optimization_parameters = optimization_parameters - elif mode == INFERENCE: - if optimization_parameters is not None: - raise ValueError('`optimization_parameters` should be `None` ' - 'for inference mode.') - self._optimization_parameters = ( - StochasticGradientDescentParameters(1.)) - else: - raise ValueError('`mode` only supports {} and {}; got {}.' - .format(TRAINING, INFERENCE, mode)) - self._mode = mode - - # TODO(shizhiw): move `optimization_parameters` into `_optimizer_handler` - # and create special handler for inference that inherits from - # StochasticGradientDescentHandler with more user-friendly error message - # on get_slot(). - self._optimizer_handler = _get_optimization_handler( - self._optimization_parameters) - - self._config_proto = self._create_config_proto() - - @property - def hosts(self): - """A list of device names for CPU hosts. - - Returns: - A list of device names for CPU hosts. - """ - return copy.copy(self._hosts) - - # TODO(shizhiw): change to num_tensor_cores_per_host to be more explicit and - # to be consistent with `tpu_embedding_configuration.proto`. - @property - def num_cores_per_host(self): - """Number of TPU cores on a CPU host. - - Returns: - Number of TPU cores on a CPU host. - """ - return self._num_cores_per_host - - @property - def num_cores(self): - """Total number of TPU cores on all hosts. - - Returns: - Total number of TPU cores on all hosts. - """ - return self._num_cores - - @property - def batch_size_per_core(self): - """Batch size for each TPU core. - - The sparse tensors in `sparse_features_list` to `generate_enqueue_ops` - must have batch dimension equal to this. - - Returns: - Batch size for each TPU core. - """ - return self._batch_size_per_core - - @property - def config_proto(self): - """Create embedding config proto for `tpu.initialize_system()`. - - Returns: - an `TPUEmbeddingConfiguration` proto describing the desired - configuration of the hardware embedding lookup tables, which - is passed to `tpu.initialize_system()`. - """ - return self._config_proto - - @property - def table_to_config_dict(self): - return copy.copy(self._table_to_config_dict) - - @property - def feature_to_table_dict(self): - return copy.copy(self._feature_to_table_dict) - - @property - def table_to_features_dict(self): - return copy.copy(self._table_to_features_dict) - - @property - def optimization_parameters(self): - return self._optimization_parameters - - def _create_config_proto(self): - """Create `TPUEmbeddingConfiguration`.""" - config_proto = elc.TPUEmbeddingConfiguration() - for table in self._table_to_config_dict: - table_descriptor = config_proto.table_descriptor.add() - table_descriptor.name = table - - table_config = self._table_to_config_dict[table] - table_descriptor.vocabulary_size = table_config.vocabulary_size - table_descriptor.dimension = table_config.dimension - - features_for_table = self._table_to_features_dict[table] - table_descriptor.num_features = len(features_for_table) - - table_descriptor.optimization_parameters.learning_rate.constant = ( - self._optimization_parameters.learning_rate) - table_descriptor.optimization_parameters.gradient_accumulation_status = ( - optimization_parameters_pb2.GradientAccumulationStatus.ENABLED - if self._optimization_parameters.use_gradient_accumulation else - optimization_parameters_pb2.GradientAccumulationStatus.DISABLED) - # For compatibility with old TPU workers. - table_descriptor.optimization_parameters.use_gradient_accumulation = ( - self._optimization_parameters.use_gradient_accumulation) - self._optimizer_handler.set_optimization_parameters(table_descriptor) - - config_proto.mode = self._mode - config_proto.batch_size_per_tensor_core = self._batch_size_per_core - config_proto.num_hosts = self._num_hosts - config_proto.num_tensor_cores = self._num_cores - config_proto.sharding_strategy = elc.TPUEmbeddingConfiguration.DIV_DEFAULT - config_proto.pipeline_execution_with_tensor_core = ( - self._optimization_parameters.pipeline_execution_with_tensor_core) - - return config_proto - - def create_variables_and_ops(self, embedding_variable_name_by_table=None, - slot_variable_names_by_table=None): - """Create embedding and slot variables, with ops to load and retrieve them. - - Args: - embedding_variable_name_by_table: A dictionary mapping from string of - table name to string of embedding variable name. If `None`, - defaults from `get_default_slot_variable_names()` will be used. - slot_variable_names_by_table: A dictionary mapping from string of table - name to `AdamSlotVariableNames`, `AdagradSlotVariableNames` etc. If - `None`, defaults from `get_default_slot_variable_names()` will be used. - - Returns: - `tpu_embedding.VariablesAndOps` with: - A dictionary mapping from string of table name to embedding variables, - A dictionary mapping from string of table name to AdagradSlotVariable, - AdamSlotVariables etc with slot variables, - A function which returns a list of ops to load embedding and slot - variables from TPU to CPU. - A function which returns a list of ops to retrieve embedding and slot - variables from TPU to CPU. - """ - embedding_variables_by_table = {} - slot_variables_by_table = {} - load_op_fns = [] - retrieve_op_fns = [] - for table in self._table_to_config_dict: - if embedding_variable_name_by_table: - embedding_variable_name = embedding_variable_name_by_table[table] - else: - embedding_variable_name = table - if slot_variable_names_by_table: - slot_variable_names = slot_variable_names_by_table[table] - else: - slot_variable_names = ( - self._optimizer_handler.get_default_slot_variable_names(table)) - - device_fn = _create_device_fn(self._hosts) - with ops.device(device_fn): - table_variables = _create_partitioned_variables( - name=embedding_variable_name, - num_hosts=self._num_hosts, - vocabulary_size=self._table_to_config_dict[table].vocabulary_size, - embedding_dimension=self._table_to_config_dict[table].dimension, - initializer=self._table_to_config_dict[table].initializer, - collections=[ops.GraphKeys.GLOBAL_VARIABLES]) - embedding_variables_by_table[table] = table_variables - - slot_variables_for_table, load_ops_fn, retrieve_ops_fn = ( - self._optimizer_handler.create_variables_and_ops( - table, slot_variable_names, self._num_hosts, - self._table_to_config_dict[table], table_variables) - ) - slot_variables_by_table[table] = slot_variables_for_table - load_op_fns.append(load_ops_fn) - retrieve_op_fns.append(retrieve_ops_fn) - - def load_ops(): - """Calls and returns the load ops for each embedding table. - - Returns: - A list of ops to load embedding and slot variables from CPU to TPU. - """ - load_ops_list = [] - for load_op_fn in load_op_fns: - load_ops_list.extend(load_op_fn()) - return load_ops_list - - def retrieve_ops(): - """Calls and returns the retrieve ops for each embedding table. - - Returns: - A list of ops to retrieve embedding and slot variables from TPU to CPU. - """ - retrieve_ops_list = [] - for retrieve_op_fn in retrieve_op_fns: - retrieve_ops_list.extend(retrieve_op_fn()) - return retrieve_ops_list - - return VariablesAndOps(embedding_variables_by_table, - slot_variables_by_table, - load_ops, retrieve_ops) - - def generate_enqueue_ops(self, sparse_features_list): - """Generate enqueue ops. - - Args: - sparse_features_list: a list of dictionary mapping from string - of feature names to sparse tensor. Each dictionary is for one - TPU core. Dictionaries for the same core should be contiguous - on the list. - - Returns: - Ops to enqueue to TPU for embedding. - """ - self._validate_generate_enqueue_ops_sparse_features_list( - sparse_features_list) - return [ - self._generate_enqueue_op( - sparse_features, device_ordinal=i % self._num_cores_per_host) - for i, sparse_features in enumerate(sparse_features_list) - ] - - def _validate_generate_enqueue_ops_sparse_features_list( - self, sparse_features_list): - """Validate `sparse_features_list`.""" - if len(sparse_features_list) != self._num_cores: - raise ValueError('Length of `sparse_features_list` should match the ' - 'number of cores; ' - '`len(sparse_features_list)` is {}, ' - 'number of cores is {}.'.format( - len(sparse_features_list), self._num_cores)) - - feature_set = set(self._feature_to_table_dict.keys()) - contiguous_device = None - for i, sparse_features in enumerate(sparse_features_list): - used_feature_set = set(sparse_features.keys()) - - # Check features are valid. - missing_feature_set = feature_set - used_feature_set - if missing_feature_set: - raise ValueError('`sparse_features_list[{}]` misses a feature that is ' - 'in `feature_to_config_dict`: {}.'.format( - i, missing_feature_set)) - - extra_feature_set = used_feature_set - feature_set - if extra_feature_set: - raise ValueError('`sparse_features_list[{}]` has a feature that is not ' - 'in `feature_to_config_dict`: {}.'.format( - i, extra_feature_set)) - - device = None - device_feature = None - for feature, tensor in six.iteritems(sparse_features): - if not isinstance(tensor, sparse_tensor.SparseTensor): - raise ValueError('`sparse_features_list[{}]` has a feature that is ' - 'not mapped to `SparseTensor`. ' - '`feature`: {}, type: {}'.format( - i, feature, type(tensor))) - - # Check all features are on the same device. - if device is None: - device = tensor.op.device - device_feature = feature - else: - if device != tensor.op.device: - raise ValueError('Devices are different between features in ' - '`sparse_features_list[{}]`; ' - 'devices: {}, {}; features: {}, {}.'.format( - i, device, tensor.op.device, feature, - device_feature)) - - if i % self._num_cores_per_host: - if device != contiguous_device: - raise ValueError('We expect the `sparse_features` which are on the ' - 'same host to be contiguous in ' - '`sparse_features_list`, ' - '`sparse_features_list[{}]` is on device {}, ' - 'but is expected to be on device {}.'.format( - i, device, contiguous_device)) - else: - contiguous_device = device - - def _generate_enqueue_op(self, sparse_features, device_ordinal): - with ops.colocate_with(list(sparse_features.values())[0]): - sample_idcs, embedding_idcs, aggregation_weights = ( - self._format_for_tpu_embedding_sparse_batch(sparse_features)) - return tpu_ops.enqueue_tpu_embedding_sparse_batch( - sample_idcs, - embedding_idcs, - aggregation_weights, - combiners=self._combiners, - device_ordinal=device_ordinal) - - def _format_for_tpu_embedding_sparse_batch(self, sparse_features): - """Format sparse features for `enqueue_tpu_embedding_sparse_batch()`. - - Args: - sparse_features: a `Dict` of `SparseTensor`s for embedding. - - Returns: - Arguments for `enqueue_tpu_embedding_sparse_batch()`. - """ - - sample_idcs, embedding_idcs, aggregation_weights = list(), list(), list() - for table in self._table_to_features_dict: - sample_t, indices_t, weights_t = list(), list(), list() - - features = self._table_to_features_dict[table] - for i, feature in enumerate(features): - tensor = sparse_features[feature] - sample_indices = tensor.indices[:, 0] - embedding_indices = tensor.values - weights = array_ops.ones_like(embedding_indices) - sample_t.append(i * self._batch_size_per_core + sample_indices) - indices_t.append(embedding_indices) - weights_t.append(weights) - - sample_idcs.append( - math_ops.cast(array_ops.concat(sample_t, axis=0), dtype=dtypes.int32)) - embedding_idcs.append( - math_ops.cast( - array_ops.concat(indices_t, axis=0), dtype=dtypes.int32)) - aggregation_weights.append( - math_ops.cast( - array_ops.concat(weights_t, axis=0), dtype=dtypes.float32)) - - return sample_idcs, embedding_idcs, aggregation_weights - - def get_activations(self): - """Get activations for features. - - This should be called within `computation` that is passed to - `tpu.replicate` and friends. - - Returns: - A dictionary mapping from `String` of feature name to `Tensor` - of activation. - """ - recv_activations = tpu_ops.recv_tpu_embedding_activations( - num_outputs=len(self._table_to_config_dict), - config=self._config_proto.SerializeToString()) - - activations = collections.OrderedDict() - for table_id, table in enumerate(self._table_to_features_dict): - features = self._table_to_features_dict[table] - for lookup_id, feature in enumerate(features): - start_row = lookup_id * self._batch_size_per_core - end_row = start_row + self._batch_size_per_core - activations[feature] = recv_activations[table_id][start_row:end_row, :] - return activations - - def generate_send_gradients_op(self, feature_to_gradient_dict): - """Send gradient to TPU embedding. - - Args: - feature_to_gradient_dict: dict mapping feature names to gradient wrt - activations. - - Returns: - SendTPUEmbeddingGradients Op. - - Raises: - RuntimeError: If `mode` is not `TRAINING`. - """ - if self._mode != TRAINING: - raise RuntimeError('Only in training mode gradients need to ' - 'be sent to TPU embedding; got mode {}.' - .format(self._mode)) - gradients = [] - for table in self._table_to_features_dict: - features = self._table_to_features_dict[table] - table_gradients = [ - feature_to_gradient_dict[feature] for feature in features - ] - concat_table_grads = array_ops.concat(table_gradients, axis=0) - gradients.append(concat_table_grads) - return tpu_ops.send_tpu_embedding_gradients( - inputs=gradients, config=self.config_proto.SerializeToString()) - - -def _validate_table_to_config_dict(table_to_config_dict): - """Validate `table_to_config_dict`.""" - for k, v in six.iteritems(table_to_config_dict): - if not isinstance(v, TableConfig): - raise ValueError('Value of `table_to_config_dict` must be of type ' - '`TableConfig`, got {} for {}.'.format(type(v), k)) - - -def _validate_feature_to_table_dict(table_to_config_dict, - feature_to_table_dict): - """Validate `feature_to_table_dict`.""" - used_table_set = set(feature_to_table_dict.values()) - table_set = set(table_to_config_dict.keys()) - - unused_table_set = table_set - used_table_set - if unused_table_set: - raise ValueError('`table_to_config_dict` specifies table that is not ' - 'used in `feature_to_table_dict`: {}.' - .format(unused_table_set)) - - extra_table_set = used_table_set - table_set - if extra_table_set: - raise ValueError('`feature_to_table_dict` refers to a table that is not ' - 'specified in `table_to_config_dict`: {}.' - .format(extra_table_set)) - - -def _validate_batch_size(batch_size, num_cores): - if batch_size % num_cores: - raise ValueError('`batch_size` is not a multiple of number of ' - 'cores. `batch_size`={}, `_num_cores`={}.'.format( - batch_size, num_cores)) - - -def _validate_optimization_parameters(optimization_parameters): - if not isinstance(optimization_parameters, _OptimizationParameters): - raise ValueError('`optimization_parameters` must inherit from ' - '`_OptimizationPramaters`. ' - '`type(optimization_parameters)`={}'.format( - type(optimization_parameters))) - - -class _OptimizerHandler(object): - """Interface class for handling optimizer specific logic.""" - - def __init__(self, optimization_parameters): - self._optimization_parameters = optimization_parameters - - def set_optimization_parameters(self, table_descriptor): - raise NotImplementedError() - - def get_default_slot_variable_names(self, table): - raise NotImplementedError() - - def create_variables_and_ops(self, table, slot_variable_names, num_hosts, - table_config, table_variables): - raise NotImplementedError() - - -class _AdagradHandler(_OptimizerHandler): - """Handles Adagrad specific logic.""" - - def __init__(self, optimization_parameters): - super(_AdagradHandler, self).__init__(optimization_parameters) - self._table_to_accumulator_variables_dict = {} - - def set_optimization_parameters(self, table_descriptor): - table_descriptor.optimization_parameters.adagrad.SetInParent() - - def get_default_slot_variable_names(self, table): - return AdagradSlotVariableName('{}/{}'.format(table, 'Adagrad')) - - def create_variables_and_ops(self, table, slot_variable_names, num_hosts, - table_config, table_variables): - accumulator_initializer = init_ops.constant_initializer( - self._optimization_parameters.initial_accumulator) - accumulator_variables = _create_partitioned_variables( - name=slot_variable_names.accumulator, - num_hosts=num_hosts, - vocabulary_size=table_config.vocabulary_size, - embedding_dimension=table_config.dimension, - collections=[ops.GraphKeys.GLOBAL_VARIABLES], - initializer=accumulator_initializer) - slot_variables = AdagradSlotVariable(accumulator_variables) - - def load_ops_fn(): - """Returns the retrieve ops for AdaGrad embedding tables. - - Returns: - A list of ops to load embedding and slot variables from CPU to TPU. - """ - load_op_list = [] - for host_id, table_variable, accumulator_variable in (zip( - range(num_hosts), table_variables, accumulator_variables)): - with ops.colocate_with(table_variable): - load_parameters_op = ( - tpu_ops.load_tpu_embedding_adagrad_parameters( - parameters=table_variable, - accumulators=accumulator_variable, - table_name=table, - num_shards=num_hosts, - shard_id=host_id)) - load_op_list.append(load_parameters_op) - return load_op_list - - def retrieve_ops_fn(): - """Returns the retrieve ops for AdaGrad embedding tables. - - Returns: - A list of ops to retrieve embedding and slot variables from TPU to CPU. - """ - retrieve_op_list = [] - for host_id, table_variable, accumulator_variable in (zip( - range(num_hosts), table_variables, accumulator_variables)): - with ops.colocate_with(table_variable): - retrieved_table, retrieved_accumulator = ( - tpu_ops.retrieve_tpu_embedding_adagrad_parameters( - table_name=table, - num_shards=num_hosts, - shard_id=host_id)) - retrieve_parameters_op = control_flow_ops.group( - state_ops.assign(table_variable, retrieved_table), - state_ops.assign(accumulator_variable, retrieved_accumulator)) - retrieve_op_list.append(retrieve_parameters_op) - return retrieve_op_list - - return slot_variables, load_ops_fn, retrieve_ops_fn - - -class _AdamHandler(_OptimizerHandler): - """Handles Adam specific logic.""" - - def __init__(self, optimization_parameters): - super(_AdamHandler, self).__init__(optimization_parameters) - self._table_to_m_variables_dict = {} - self._table_to_v_variables_dict = {} - - def set_optimization_parameters(self, table_descriptor): - table_descriptor.optimization_parameters.adam.beta1 = ( - self._optimization_parameters.beta1) - table_descriptor.optimization_parameters.adam.beta2 = ( - self._optimization_parameters.beta2) - table_descriptor.optimization_parameters.adam.epsilon = ( - self._optimization_parameters.epsilon) - table_descriptor.optimization_parameters.adam.use_non_lazy_adam = ( - not self._optimization_parameters.lazy_adam) - table_descriptor.optimization_parameters.adam.use_sum_inside_sqrt = ( - self._optimization_parameters.sum_inside_sqrt) - - def get_default_slot_variable_names(self, table): - return AdamSlotVariableNames('{}/{}/m'.format(table, 'Adam'), - '{}/{}/v'.format(table, 'Adam')) - - def create_variables_and_ops(self, table, slot_variable_names, num_hosts, - table_config, table_variables): - m_initializer = init_ops.zeros_initializer() - m_variables = _create_partitioned_variables( - name=slot_variable_names.m, - num_hosts=num_hosts, - vocabulary_size=table_config.vocabulary_size, - embedding_dimension=table_config.dimension, - collections=[ops.GraphKeys.GLOBAL_VARIABLES], - initializer=m_initializer) - v_initializer = init_ops.zeros_initializer() - v_variables = _create_partitioned_variables( - name=slot_variable_names.v, - num_hosts=num_hosts, - vocabulary_size=table_config.vocabulary_size, - embedding_dimension=table_config.dimension, - collections=[ops.GraphKeys.GLOBAL_VARIABLES], - initializer=v_initializer) - slot_variables = AdamSlotVariables(m_variables, v_variables) - - def load_ops_fn(): - """Returns the retrieve ops for AdaGrad embedding tables. - - Returns: - A list of ops to load embedding and slot variables from CPU to TPU. - """ - load_op_list = [] - for host_id, table_variable, m_variable, v_variable in (zip( - range(num_hosts), table_variables, - m_variables, v_variables)): - with ops.colocate_with(table_variable): - load_parameters_op = ( - tpu_ops.load_tpu_embedding_adam_parameters( - parameters=table_variable, - momenta=m_variable, - velocities=v_variable, - table_name=table, - num_shards=num_hosts, - shard_id=host_id)) - - load_op_list.append(load_parameters_op) - return load_op_list - - def retrieve_ops_fn(): - """Returns the retrieve ops for Adam embedding tables. - - Returns: - A list of ops to retrieve embedding and slot variables from TPU to CPU. - """ - - retrieve_op_list = [] - for host_id, table_variable, m_variable, v_variable in (zip( - range(num_hosts), table_variables, - m_variables, v_variables)): - with ops.colocate_with(table_variable): - retrieved_table, retrieved_m, retrieved_v = ( - tpu_ops.retrieve_tpu_embedding_adam_parameters( - table_name=table, - num_shards=num_hosts, - shard_id=host_id)) - retrieve_parameters_op = control_flow_ops.group( - state_ops.assign(table_variable, retrieved_table), - state_ops.assign(m_variable, retrieved_m), - state_ops.assign(v_variable, retrieved_v)) - - retrieve_op_list.append(retrieve_parameters_op) - return retrieve_op_list - - return slot_variables, load_ops_fn, retrieve_ops_fn - - -class _StochasticGradientDescentHandler(_OptimizerHandler): - """Handles stochastic gradient descent specific logic.""" - - def set_optimization_parameters(self, table_descriptor): - (table_descriptor.optimization_parameters.stochastic_gradient_descent - .SetInParent()) - - def get_default_slot_variable_names(self, table): - return None - - def create_variables_and_ops(self, table, slot_variable_names, num_hosts, - table_config, table_variables): - del table_config - - def load_ops_fn(): - """Returns the retrieve ops for AdaGrad embedding tables. - - Returns: - A list of ops to load embedding and slot variables from CPU to TPU. - """ - load_op_list = [] - for host_id, table_variable in (zip( - range(num_hosts), table_variables)): - with ops.colocate_with(table_variable): - load_parameters_op = ( - tpu_ops - .load_tpu_embedding_stochastic_gradient_descent_parameters( - parameters=table_variable, - table_name=table, - num_shards=num_hosts, - shard_id=host_id)) - - load_op_list.append(load_parameters_op) - return load_op_list - - def retrieve_ops_fn(): - """Returns the retrieve ops for SGD embedding tables. - - Returns: - A list of ops to retrieve embedding and slot variables from TPU to CPU. - """ - - retrieve_op_list = [] - for host_id, table_variable in (zip( - range(num_hosts), table_variables)): - with ops.colocate_with(table_variable): - retrieved_table = ( - tpu_ops - .retrieve_tpu_embedding_stochastic_gradient_descent_parameters( - table_name=table, - num_shards=num_hosts, - shard_id=host_id)) - retrieve_parameters_op = control_flow_ops.group( - state_ops.assign(table_variable, retrieved_table)) - - retrieve_op_list.append(retrieve_parameters_op) - return retrieve_op_list - - return None, load_ops_fn, retrieve_ops_fn - - -def _get_optimization_handler(optimization_parameters): - if isinstance(optimization_parameters, AdagradParameters): - return _AdagradHandler(optimization_parameters) - elif isinstance(optimization_parameters, AdamParameters): - return _AdamHandler(optimization_parameters) - elif isinstance(optimization_parameters, StochasticGradientDescentParameters): - return _StochasticGradientDescentHandler(optimization_parameters) - else: - return NotImplementedError() - - -def _create_ordered_dict(d): - """Create an OrderedDict from Dict.""" - return collections.OrderedDict((k, d[k]) for k in sorted(d)) - - -def _create_combiners(table_to_config_dict): - return [table_to_config_dict[t].combiner for t in table_to_config_dict] - - -def _create_table_to_features_dict(feature_to_table_dict): - """Create mapping from table to a list of its features.""" - table_to_features_dict_tmp = {} - for feature, table in six.iteritems(feature_to_table_dict): - if table in table_to_features_dict_tmp: - table_to_features_dict_tmp[table].append(feature) - else: - table_to_features_dict_tmp[table] = [feature] - - table_to_features_dict = collections.OrderedDict() - for table in sorted(table_to_features_dict_tmp): - table_to_features_dict[table] = sorted(table_to_features_dict_tmp[table]) - return table_to_features_dict - - -def _create_device_fn(hosts): - """Create device_fn() to use with _create_partitioned_variables().""" - - def device_fn(op): - """Returns the `device` for `op`.""" - part_match = re.match(r'.*/part_(\d+)(/|$)', op.name) - - if part_match: - idx = int(part_match.group(1)) - else: - raise RuntimeError('Internal Error: ' - 'Expected %s to contain /part_*.' % op.name) - - device = hosts[idx] - return device - - return device_fn - - -def _create_partitioned_variables(name, - num_hosts, - vocabulary_size, - embedding_dimension, - initializer, - collections=None): # pylint: disable=redefined-outer-name - """Creates ParitionedVariables based on `num_hosts` for `table`.""" - # TODO(shizhiw): automatically place embedding lookup elsewhere? - if vocabulary_size < num_hosts: - raise ValueError('`vocabulary_size`({}) is smaller than `num_hosts`({}). ' - 'As TPU embedding is not optimized for small tables, ' - 'please consider other ways for this embedding lookup.') - - return list(variable_scope.get_variable( - name, - shape=(vocabulary_size, embedding_dimension), - partitioner=partitioned_variables.fixed_size_partitioner(num_hosts), - dtype=dtypes.float32, - initializer=initializer, - collections=collections, - trainable=False)) +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.tpu_embedding import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_embedding_gradient.py b/tensorflow/contrib/tpu/python/tpu/tpu_embedding_gradient.py index dace0d801b3a91caae9cafea59366f4adc9325a7..308adc77e9ad2d912d0461512655b55faa53da60 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_embedding_gradient.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_embedding_gradient.py @@ -1,153 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# =================================================================== -"""Optional helper for gradient handling.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections - -from tensorflow.contrib.tpu.python.ops import tpu_ops -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops -from tensorflow.python.ops import variable_scope -from tensorflow.python.ops import variables - - -def get_gradients_through_compute_gradients(optimizer, loss, activations): - """Compute gradients to send to TPU embedding. - - Args: - optimizer: a subclass of optimizer.Optimizer, usually CrossShardOptimizer. - Used to call compute_gradients(). - loss: a Tensor to call optimizer.compute_gradients() on. - activations: an OrderedDict mapping feature_name to Tensors of activations. - - Returns: - An OrderedDict mapping from feature name Strings to Tensors of gradients of - the loss wrt the activations of the features. - """ - activation_list = activations.values() - grads_and_vars = optimizer.compute_gradients(loss, activation_list) - grads = [grad for grad, _ in grads_and_vars] - feature_to_gradient_dict = collections.OrderedDict( - zip(activations.keys(), grads)) - return feature_to_gradient_dict - - -def create_dummy_table_variables(tpu_embedding): - """Create dummy embedding table variables. - - The sole purpose of these dummy variables are to trigger gradient - calcuation wrt them so that the gradients wrt activation can be captured - and later sent to TPU embedding. - - Args: - tpu_embedding: TPUEmbedding, dummy table variables will be created for use - with tpu_embedding. - - Returns: - A tuple of dummy variables and their initializer. - - Raises: - RuntimeError: if collection to store gradients already exists and is not - empty. - """ - dummy_table_variables = collections.OrderedDict() - for table_id, table in enumerate(tpu_embedding.table_to_features_dict): - dummy_table_variables[table] = ( - # Explicitly specifying collections prevents this variable from - # being added to the GLOBAL_VARIABLES collection, so that Saver() - # ignores it. - # But Tensorflow optimizer creates slot variable for these dummy - # variable, e.g. tpu_embedding_dummy_table_variable_mlp_user/Adam{_1}, - # which will be in GLOBAL_VARIABLES collection, - variable_scope.get_variable( - 'tpu_embedding_dummy_table_variable_{}'.format(table), - dtype=dtypes.float32, - shape=[1], - use_resource=True, - trainable=True, - collections=['tpu_embedding_dummy_table_variables'])) - - g = ops.get_default_graph() - table_gradients = g.get_collection_ref( - 'tpu_embedding_gradients_table_{}'.format(table_id)) - if table_gradients: - raise RuntimeError( - 'tpu_embedding_gradients_table_{} is not empty.'.format(table_id)) - table_gradients.extend( - [None] * len(tpu_embedding.table_to_features_dict[table])) - - return (dummy_table_variables, - variables.variables_initializer( - dummy_table_variables.values(), - name='tpu_embedding_dummy_table_variables_init')) - - -def hook_dummy_table_variables_to_activations(tpu_embedding, activations, - dummy_table_variables): - """Have activations depend on dummy table variables for gradient intercept. - - Args: - tpu_embedding: TPUEmbedding, activations and dummy_table_variables are from - tpu_embedding. - activations: An OrderedDict of feature name String to activation tensors. - dummy_table_variables: An OrderedDict of table name String to dummy table - variables. - - Returns: - An OrderedDict of feature name String to activation tensors, which can be - used just as the activations input. - """ - new_activations = collections.OrderedDict() - for feature in activations: - table = tpu_embedding.feature_to_table_dict[feature] - new_activations[feature] = tpu_ops.tpu_embedding_activations( - dummy_table_variables[table], - activations[feature], - table_id=tpu_embedding.table_to_config_dict.keys().index(table), - lookup_id=tpu_embedding.table_to_features_dict[table].index(feature)) - return new_activations - - -def get_gradients_through_dummy_table_variables(tpu_embedding): - """Get gradients wrt the activations of each feature. - - Args: - tpu_embedding: TPUEmbedding, create dummy table variable to be used with - tpu_embedding. - - Returns: - An OrderedDict mapping feature name to gradient. - - Raises: - ValueError: if some gradients are not defined. - """ - g = ops.get_default_graph() - feature_to_gradient_dict = collections.OrderedDict() - for table_id, table in enumerate(tpu_embedding.table_to_config_dict): - table_gradients = g.get_collection( - 'tpu_embedding_gradients_table_{}'.format(table_id)) - if any(gradient is None for gradient in table_gradients): - raise ValueError( - 'Table {} with id {} has undefined gradients: this is probably ' - 'because the model asked TPUEmbedding to compute activations that ' - 'were not used.'.format(table, table_id)) - for feature, gradient in zip(tpu_embedding.table_to_features_dict[table], - table_gradients): - feature_to_gradient_dict[feature] = gradient - return feature_to_gradient_dict +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.tpu_embedding_gradient import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py index b2019c1083653f9af2c273cdf24ba5b0364bf478..893118412e1363ce50416e6ef36692bc23d04179 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py @@ -1,3761 +1,33 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# =================================================================== -"""TPUEstimator class.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections -import copy -import os -import signal -import sys -import threading -import time - -import numpy as np -import six -from six.moves import queue as Queue # pylint: disable=redefined-builtin -from six.moves import xrange # pylint: disable=redefined-builtin - -from tensorflow.contrib.tpu.python.ops import tpu_ops -from tensorflow.contrib.tpu.python.tpu import _tpu_estimator_embedding -from tensorflow.contrib.tpu.python.tpu import error_handling -from tensorflow.contrib.tpu.python.tpu import functional as tpu_functional -from tensorflow.contrib.tpu.python.tpu import session_support -from tensorflow.contrib.tpu.python.tpu import tensor_tracer -from tensorflow.contrib.tpu.python.tpu import tpu -from tensorflow.contrib.tpu.python.tpu import tpu_config -from tensorflow.contrib.tpu.python.tpu import tpu_context -from tensorflow.contrib.tpu.python.tpu import tpu_embedding_gradient -from tensorflow.contrib.tpu.python.tpu import tpu_feed -from tensorflow.contrib.tpu.python.tpu import tpu_function -from tensorflow.contrib.tpu.python.tpu import training_loop -from tensorflow.contrib.tpu.python.tpu import util as util_lib -from tensorflow.contrib.tpu.python.tpu._tpu_estimator_embedding import AdamParameters # pylint: disable=unused-import -from tensorflow.contrib.tpu.python.tpu._tpu_estimator_embedding import EmbeddingConfigSpec # pylint: disable=unused-import -from tensorflow.contrib.training.python.training import hparam -from tensorflow.core.framework import variable_pb2 -from tensorflow.core.framework.summary_pb2 import Summary -from tensorflow.core.protobuf import config_pb2 -from tensorflow.core.protobuf.tpu import compilation_result_pb2 as tpu_compilation_result -from tensorflow.python.client import session as tf_session -from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.data.util import nest as data_nest -from tensorflow.python.estimator import estimator as estimator_lib -from tensorflow.python.estimator import model_fn as model_fn_lib -from tensorflow.python.estimator.export import export_output as export_output_lib -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors -from tensorflow.python.framework import function -from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import check_ops -from tensorflow.python.ops import control_flow_ops -from tensorflow.python.ops import init_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import resource_variable_ops -from tensorflow.python.ops import state_ops -from tensorflow.python.ops import summary_ops_v2 as contrib_summary -from tensorflow.python.ops import variable_scope -from tensorflow.python.ops import variables -from tensorflow.python.platform import tf_logging as logging -from tensorflow.python.saved_model import tag_constants -from tensorflow.python.summary import summary -from tensorflow.python.training import basic_session_run_hooks -from tensorflow.python.training import evaluation -from tensorflow.python.training import session_run_hook -from tensorflow.python.training import training -from tensorflow.python.training import training_util -from tensorflow.python.util import function_utils -from tensorflow.python.util import nest -from tensorflow.python.util import tf_inspect - -_INITIAL_LOSS = 1e7 -_ZERO_LOSS = 0. -_TPU_ESTIMATOR = 'tpu_estimator' -_ITERATIONS_PER_LOOP_VAR = 'iterations_per_loop' -_BATCH_SIZE_KEY = 'batch_size' -_CTX_KEY = 'context' -_USE_TPU_KEY = 'use_tpu' -_CROSS_REPLICA_SUM_OP = 'CrossReplicaSum' -_ONE_GIGABYTE = 1024 * 1024 * 1024 -_TPU_ENQUEUE_OPS = '_tpu_enqueue_ops' -_TPU_TRAIN_OP = '_tpu_train_op' -_REWRITE_FOR_INFERENCE_MODE = '_rewrite_for_inference' -_KEY_WHEN_PREDICTIONS_IS_A_TENSOR = '_key_when_predictions_is_a_tensor' - -# Ideally _USE_TPU_KEY should be reserved as well. However there are already -# models that make use of this key, thus it can not be reserved now to prevent -# breakage. In the long run, we would like to mitigate this by migrating models -# off of using _USE_TPU_KEY. -_RESERVED_PARAMS_KEYS = [_BATCH_SIZE_KEY, _CTX_KEY] - -# TODO(b/65703635): Flip the value and remove all dead code. Currently, this is -# only used for per-core based deployments. For per-host based pipelines, if a -# user returns a Dataset instance it will be automatically wrapped in a -# tf.while_loop (This can be disabled by returning features and labels -# explicitly). -_WRAP_INPUT_FN_INTO_WHILE_LOOP = False - -ops.register_proto_function( - '{}_{}'.format(_TPU_ESTIMATOR, _ITERATIONS_PER_LOOP_VAR), - proto_type=variable_pb2.VariableDef, - to_proto=resource_variable_ops._to_proto_fn, # pylint: disable=protected-access - from_proto=resource_variable_ops._from_proto_fn) # pylint: disable=protected-access - - -def _is_iterable(obj): - """A Python 2 and 3 compatible util to check whether `obj` is iterable.""" - try: - iter(obj) - return True - except TypeError: - return False - - -class CatchInvalidHostcallFunctions(control_flow_ops.XLAControlFlowContext): - - def AddOp(self, op): - if op.type in [ - 'AudioSummary', 'AudioSummaryV2', 'HistogramSummary', 'ImageSummary', - 'MergeSummary', 'ScalarSummary', 'TensorSummary', 'TensorSummaryV2' - ]: - raise ValueError('Use tf.contrib.summary inside of host_calls.') - - -def _create_global_step(graph): - graph = graph or ops.get_default_graph() - if training.get_global_step(graph) is not None: - raise ValueError('"global_step" already exists.') - # Create in proper graph and base name_scope. - with graph.as_default() as g, g.name_scope(None): - return variable_scope.get_variable( - ops.GraphKeys.GLOBAL_STEP, - shape=[], - dtype=dtypes.int64, - initializer=init_ops.zeros_initializer(), - trainable=False, - use_resource=True, - collections=[ops.GraphKeys.GLOBAL_VARIABLES, ops.GraphKeys.GLOBAL_STEP]) - - -def _create_or_get_iterations_per_loop(): - """Creates or gets the iterations_per_loop variable. - - In TPUEstimator, the user provided computation, the model_fn, is wrapped - inside a tf.while_loop for peak performance. The iterations of the loop are - specified by this variable, which adjusts its value on the CPU after each TPU - program execution and before the next TPU execution. - - The purpose of using a variable, rather then a constant, is to allow - TPUEstimator adapt the TPU training iterations according to the final steps - specified by users. For example, if the user sets the iterations_per_loop as 4 - in TPUConfig and steps as 10 in TPUEstimator.train(), the iterations_per_loop - variable will have the following value before each TPU training. - - - 1-th TPU execution: iterations_per_loop = 4 - - 2-th TPU execution: iterations_per_loop = 4 - - 3-th TPU execution: iterations_per_loop = 2 - - As model_fn increases the global step once per train_op invocation, the global - step is 10 after all TPU executions, matching the steps=10 inputs passed in by - users. - - Returns: - A TF non-trainable resource variable. - - Raises: - RuntimeError: If multi iterations_per_loop variables were found. - """ - graph = ops.get_default_graph() - collection_name = '{}_{}'.format(_TPU_ESTIMATOR, _ITERATIONS_PER_LOOP_VAR) - iter_vars = graph.get_collection(collection_name) - if len(iter_vars) == 1: - return iter_vars[0] - elif len(iter_vars) > 1: - raise RuntimeError('Multiple iterations_per_loop_var in collection.') - - with ops.colocate_with(training_util.get_global_step()): - with variable_scope.variable_scope( - _TPU_ESTIMATOR, reuse=variable_scope.AUTO_REUSE): - return variable_scope.get_variable( - _ITERATIONS_PER_LOOP_VAR, - initializer=init_ops.zeros_initializer(), - shape=[], - dtype=dtypes.int32, - trainable=False, - collections=[collection_name, ops.GraphKeys.LOCAL_VARIABLES], - use_resource=True) - - -def _sync_variables_ops(ctx): - """Create varriables synchronization ops. - - Gets the variables back from TPU nodes. This means the variables updated - by TPU will now be *synced* to host memory. - In BROADCAST mode, we skip this sync since the variables are ususally too - big to transmit via RPC. - - Args: - ctx: A `_InternalTPUContext` instance with mode. - - Returns: - A list of sync ops. - """ - - if not ctx.is_input_broadcast_with_iterators(): - return [ - array_ops.check_numerics(v.read_value(), - 'Gradient for %s is NaN' % v.name).op - for v in variables.trainable_variables() - ] - else: - return [control_flow_ops.no_op()] - - -def _increase_eval_step_op(iterations_per_loop): - """Returns an op to increase the eval step for TPU evaluation. - - Args: - iterations_per_loop: Tensor. The number of eval steps running in TPU system - before returning to CPU host for each `Session.run`. - - Returns: - An operation - """ - eval_step = evaluation._get_or_create_eval_step() # pylint: disable=protected-access - # Estimator evaluate increases 1 by default. So, we increase the difference. - return state_ops.assign_add( - eval_step, - math_ops.cast(iterations_per_loop - 1, dtype=eval_step.dtype), - use_locking=True) - - -def _extract_key_names(tensor_or_dict): - if isinstance(tensor_or_dict, dict): - return sorted(tensor_or_dict.keys()) - return [] - - -class _SIGNAL(object): - """Signal used to control the thread of infeed/outfeed. - - All preserved signals must be negative numbers. Positive numbers are used to - indicate the number of iterations for next training/evaluation loop. - """ - NEXT_BATCH = -1 - STOP = -2 - - -class TPUEstimatorSpec(model_fn_lib._TPUEstimatorSpec): # pylint: disable=protected-access - """Ops and objects returned from a `model_fn` and passed to `TPUEstimator`. - - See `EstimatorSpec` for `mode`, `predictions`, `loss`, `train_op`, and - `export_outputs`. - - For evaluation, `eval_metrics `is a tuple of `metric_fn` and `tensors`, where - `metric_fn` runs on CPU to generate metrics and `tensors` represents the - `Tensor`s transferred from TPU system to CPU host and passed to `metric_fn`. - To be precise, TPU evaluation expects a slightly different signature from the - `tf.estimator.Estimator`. While `EstimatorSpec.eval_metric_ops` expects a - dict, `TPUEstimatorSpec.eval_metrics` is a tuple of `metric_fn` and `tensors`. - The `tensors` could be a list of `Tensor`s or dict of names to `Tensor`s. The - `tensors` usually specify the model logits, which are transferred back from - TPU system to CPU host. All tensors must have be batch-major, i.e., the batch - size is the first dimension. Once all tensors are available at CPU host from - all shards, they are concatenated (on CPU) and passed as positional arguments - to the `metric_fn` if `tensors` is list or keyword arguments if `tensors` is - a dict. `metric_fn` takes the `tensors` and returns a dict from metric string - name to the result of calling a metric function, namely a `(metric_tensor, - update_op)` tuple. See `TPUEstimator` for MNIST example how to specify the - `eval_metrics`. - - `scaffold_fn` is a function running on CPU to generate the `Scaffold`. This - function should not capture any Tensors in `model_fn`. - - `host_call` is a tuple of a `function` and a list or dictionary of `tensors` - to pass to that function and returns a list of Tensors. `host_call` currently - works for train() and evaluate(). The Tensors returned by the function is - executed on the CPU on every step, so there is communication overhead when - sending tensors from TPU to CPU. To reduce the overhead, try reducing the - size of the tensors. The `tensors` are concatenated along their major (batch) - dimension, and so must be >= rank 1. The `host_call` is useful for writing - summaries with `tf.contrib.summary.create_file_writer`. - """ - - def __new__(cls, - mode, - predictions=None, - loss=None, - train_op=None, - eval_metrics=None, - export_outputs=None, - scaffold_fn=None, - host_call=None, - training_hooks=None, - evaluation_hooks=None, - prediction_hooks=None): - """Creates a validated `TPUEstimatorSpec` instance.""" - host_calls = {} - if eval_metrics is not None: - host_calls['eval_metrics'] = eval_metrics - if host_call is not None: - host_calls['host_call'] = host_call - _OutfeedHostCall.validate(host_calls) - - training_hooks = tuple(training_hooks or []) - evaluation_hooks = tuple(evaluation_hooks or []) - prediction_hooks = tuple(prediction_hooks or []) - - for hook in training_hooks + evaluation_hooks + prediction_hooks: - if not isinstance(hook, session_run_hook.SessionRunHook): - raise TypeError('All hooks must be SessionRunHook instances, given: {}' - .format(hook)) - - return super(TPUEstimatorSpec, cls).__new__( - cls, - mode=mode, - predictions=predictions, - loss=loss, - train_op=train_op, - eval_metrics=eval_metrics, - export_outputs=export_outputs, - scaffold_fn=scaffold_fn, - host_call=host_call, - training_hooks=training_hooks, - evaluation_hooks=evaluation_hooks, - prediction_hooks=prediction_hooks) - - def as_estimator_spec(self): - """Creates an equivalent `EstimatorSpec` used by CPU train/eval.""" - host_calls = {} - if self.eval_metrics is not None: - host_calls['eval_metrics'] = self.eval_metrics - if self.host_call is not None: - host_calls['host_call'] = self.host_call - host_call_ret = _OutfeedHostCall.create_cpu_hostcall(host_calls) - eval_metric_ops = None - if self.eval_metrics is not None: - eval_metric_ops = host_call_ret['eval_metrics'] - hooks = None - if self.host_call is not None: - hooks = [_OutfeedHostCallHook(host_call_ret['host_call'])] - loss = self.loss - if tensor_tracer.TensorTracer.is_enabled() \ - and self.train_op is not None: - tt = tensor_tracer.TensorTracer() - loss = tt.trace_cpu(ops.get_default_graph(), loss, self.train_op) - - hooks = tuple(hooks or []) - scaffold = self.scaffold_fn() if self.scaffold_fn else None - return model_fn_lib.EstimatorSpec( - mode=self.mode, - predictions=self.predictions, - loss=loss, - train_op=self.train_op, - eval_metric_ops=eval_metric_ops, - export_outputs=self.export_outputs, - scaffold=scaffold, - training_hooks=self.training_hooks + hooks, - evaluation_hooks=self.evaluation_hooks + hooks, - prediction_hooks=self.prediction_hooks + hooks) - - -class _OpQueueContext(object): - """Manages work queue and thread for a infeed/outfeed thread.""" - - def __init__(self, name, target, args): - self._name = name - self._queue = Queue.Queue() - args = (self,) + args - self._thread = threading.Thread(name=name, target=target, args=args) - self._thread.daemon = True - self._thread.start() - - def stop(self): - self._queue.put(_SIGNAL.STOP) - - def send_next_batch_signal(self, iterations): - self._queue.put(iterations) - - def read_iteration_counts(self): - while True: - iterations = self._queue.get(block=True) - logging.debug('%s read iterations %s', self._name, iterations) - if iterations == _SIGNAL.STOP: - logging.info('%s received shutdown signal, stopping.', self._name) - return - yield iterations - - def join(self): - logging.info('Shutting down %s thread.', self._name) - self.stop() - self._thread.join() - - -class _OpSignalOnceQueueContext(_OpQueueContext): - """Manages work queue and thread for a infeed/outfeed thread. - - This subclass only signals once. - """ - - def __init__(self, name, target, args): - super(_OpSignalOnceQueueContext, self).__init__(name, target, args) - self._has_signaled = False - - def send_next_batch_signal(self, iterations): - if not self._has_signaled: - self._queue.put(iterations) - self._has_signaled = True - - -class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): - """A Session hook setting up the TPU initialization, infeed, and outfeed. - - This hook does two major things: - 1. initialize and shutdown TPU system. - 2. launch and join the threads for infeed enqueue and (optional) outfeed - dequeue. - """ - - def __init__(self, - ctx, - enqueue_ops, - dequeue_ops, - tpu_compile_op, - run_infeed_loop_on_coordinator=True, - rendezvous=None, - master=None, - session_config=None, - tpu_init_ops=None): - self._master_job = ctx.master_job - self._enqueue_ops = enqueue_ops - self._dequeue_ops = dequeue_ops - self._rendezvous = rendezvous - self._master = master - self._session_config = session_config - self._init_ops = list(tpu_init_ops or []) - if ctx.embedding_config is None: - self._embedding_layer_config = None - else: - self._embedding_layer_config = ( - ctx.embedding_config.tpu_embedding.config_proto) - self._run_infeed_loop_on_coordinator = run_infeed_loop_on_coordinator - self._initial_infeed_sleep_secs = ( - ctx.config.tpu_config.initial_infeed_sleep_secs) - - self._feed_error = None - self._finished = False - # When using model parallelism, the TPU is pre-initialized at startup to - # fetch mesh information. We skip re-initializing it here to avoid - # suspected issues due to the mesh layout changing on the second - # initialization. - self._should_initialize_tpu = not ctx.model_parallelism_enabled - self._tpu_compile_op = tpu_compile_op - - def begin(self): - logging.info('TPU job name %s', self._master_job) - self._iterations_per_loop_var = _create_or_get_iterations_per_loop() - if self._should_initialize_tpu: - self._finalize_ops = [tpu.shutdown_system(job=self._master_job)] - else: - self._finalize_ops = [] - - summary_writer_init_ops = contrib_summary.summary_writer_initializer_op() - self._init_ops.extend(summary_writer_init_ops) - # Get all the writer resources from the initializer, so we know what to - # flush. - for op in summary_writer_init_ops: - self._finalize_ops.append(contrib_summary.flush(writer=op.inputs[0])) - - def _run_infeed(self, queue_ctx, session): - logging.info('Starting infeed thread controller.') - if self._initial_infeed_sleep_secs: - logging.info('Infeed thread sleeping for %d seconds.', - self._initial_infeed_sleep_secs) - time.sleep(self._initial_infeed_sleep_secs) - logging.info('Infeed thread starting after sleep') - - with self._rendezvous.catch_errors(source='infeed', session=session): - if self._run_infeed_loop_on_coordinator: - for count, steps in enumerate(queue_ctx.read_iteration_counts()): - for i in xrange(steps): - logging.debug('Infeed enqueue for iteration (%d, %d)', count, i) - session.run(self._enqueue_ops) - else: - for _ in queue_ctx.read_iteration_counts(): - session.run(self._enqueue_ops) - logging.info('Infeed thread finished, shutting down.') - - def _run_outfeed(self, queue_ctx, session): - logging.info('Starting outfeed thread controller.') - with self._rendezvous.catch_errors(source='outfeed', session=session): - for count, steps in enumerate(queue_ctx.read_iteration_counts()): - for i in xrange(steps): - logging.debug('Outfeed dequeue for iteration (%d, %d)', count, i) - session.run(self._dequeue_ops) - logging.info('Outfeed thread finished, shutting down.') - - def _create_infeed_controller(self, name, target, args): - return _OpQueueContext(name=name, target=target, args=args) - - def _assertCompilationSucceeded(self, result, coord): - proto = tpu_compilation_result.CompilationResultProto() - proto.ParseFromString(result) - if proto.status_error_message: - logging.error('Compilation failed: {}'.format(proto.status_error_message)) - coord.request_stop() - else: - logging.info('Compilation succeeded') - - def after_create_session(self, session, coord): - if self._should_initialize_tpu: - logging.info('Init TPU system') - start = time.time() - with ops.Graph().as_default(): - with tf_session.Session( - self._master, config=self._session_config) as sess: - sess.run( - tpu.initialize_system( - job=self._master_job, - embedding_config=self._embedding_layer_config)) - logging.info('Initialized TPU in %d seconds', time.time() - start) - - session.run(self._init_ops, - options=config_pb2.RunOptions(timeout_in_ms=5 * 60 * 1000)) - - if os.environ.get('TPU_SPLIT_COMPILE_AND_EXECUTE', '') == '1': - logging.info('Compiling user program: this may take a while...') - self._assertCompilationSucceeded(session.run(self._tpu_compile_op), coord) - - self._infeed_controller = self._create_infeed_controller( - name='InfeedController', target=self._run_infeed, args=(session,)) - - self._outfeed_controller = _OpQueueContext( - name='OutfeedController', target=self._run_outfeed, args=(session,)) - - # Enable the worker watchdog to terminate workers on coordinator exit. - watchdog_timeout = int(os.environ.get('TF_TPU_WATCHDOG_TIMEOUT', '0')) - if watchdog_timeout > 0: - session_support.start_worker_watchdog(session, - shutdown_timeout=watchdog_timeout) - - def before_run(self, run_context): - self._feed_error = None - - iterations = run_context.session.run(self._iterations_per_loop_var) - - logging.info('Enqueue next (%d) batch(es) of data to infeed.', iterations) - self._infeed_controller.send_next_batch_signal(iterations) - - logging.info('Dequeue next (%d) batch(es) of data from outfeed.', - iterations) - self._outfeed_controller.send_next_batch_signal(iterations) - - def end(self, session): - self._finished = True - logging.info('Stop infeed thread controller') - self._infeed_controller.join() - self._rendezvous.record_done('infeed') - - logging.info('Stop output thread controller') - self._outfeed_controller.join() - self._rendezvous.record_done('outfeed') - - logging.info('Shutdown TPU system.') - session.run(self._finalize_ops) - - -class TPUInfeedOutfeedSessionHookForPrediction(TPUInfeedOutfeedSessionHook): - - def __init__(self, ctx, enqueue_ops, dequeue_ops, tpu_compile_op, - rendezvous=None, master=None, session_config=None): - super(TPUInfeedOutfeedSessionHookForPrediction, self).__init__( - ctx, - enqueue_ops, - dequeue_ops, - tpu_compile_op=tpu_compile_op, - run_infeed_loop_on_coordinator=False, - rendezvous=rendezvous, - master=master, - session_config=session_config) - - def _create_infeed_controller(self, name, target, args): - return _OpSignalOnceQueueContext(name=name, target=target, args=args) - - -class _TPUStopAtStepHook(session_run_hook.SessionRunHook): - """Hook that requests stop at a specified step. - - This hook is similar to the `session_run_hook._StopAfterNEvalsHook` with - following differences for TPU training: - - 1. This hook sets the variable for iterations_per_loop, which is used by - `TPUInfeedOutfeedSessionHook` to control the iterations for infeed/outfeed. - As the hook execution order is not guaranteed, the variable update is - handled in `after_create_session` and `after_run` as - `TPUInfeedOutfeedSessionHook` reads the variable value in `before_run`. - - 2. For each training loop (session.run), the global step could be increased - multiple times on TPU. The global step tensor value will be explicitly read - again in `after_run` to ensure the latest value is retrieved to avoid race - condition. - """ - - def __init__(self, iterations, num_steps=None, last_step=None): - """Initializes a `StopAtStepHook`. - - Args: - iterations: The number of iterations to run optimizer per training loop. - num_steps: Number of steps to execute. - last_step: Step after which to stop. - - Raises: - ValueError: If one of the arguments is invalid. - """ - if num_steps is None and last_step is None: - raise ValueError('One of num_steps or last_step must be specified.') - if num_steps is not None and last_step is not None: - raise ValueError('Only one of num_steps or last_step can be specified.') - self._num_steps = num_steps - self._last_step = last_step - self._iterations = iterations - - def _next_iterations(self, global_step, last_step): - gap = last_step - global_step - return min(gap, self._iterations) - - def begin(self): - self._global_step_tensor = training_util.get_global_step() - if self._global_step_tensor is None: - raise RuntimeError('Global step should be created.') - - self._iterations_per_loop_var = _create_or_get_iterations_per_loop() - - def after_create_session(self, session, coord): - global_step = session.run(self._global_step_tensor) - if self._last_step is None: - self._last_step = global_step + self._num_steps - - iterations = self._next_iterations(global_step, self._last_step) - - self._iterations_per_loop_var.load(iterations, session=session) - - def after_run(self, run_context, run_values): - # Global step cannot be retrieved via SessionRunArgs and before_run due to - # race condition. - global_step = run_context.session.run(self._global_step_tensor) - if global_step >= self._last_step: - run_context.request_stop() - else: - iterations = self._next_iterations(global_step, self._last_step) - self._iterations_per_loop_var.load( - iterations, session=run_context.session) - - -class _SetEvalIterationsHook(session_run_hook.SessionRunHook): - """Hook that requests stop at a specified step.""" - - def __init__(self, num_steps): - """Initializes a `_SetEvalIterationsHook`. - - Args: - num_steps: Number of steps to execute. - """ - self._num_steps = num_steps - - def begin(self): - self._iterations_per_loop_var = _create_or_get_iterations_per_loop() - - def after_create_session(self, session, coord): - self._iterations_per_loop_var.load(self._num_steps, session=session) - - -class _StoppingPredictHook(session_run_hook.SessionRunHook): - """Hook that requests stop according to the stopping signal in prediction.""" - - def __init__(self, scalar_stopping_signal): - self._scalar_stopping_signal = scalar_stopping_signal - - def begin(self): - self._iterations_per_loop_var = _create_or_get_iterations_per_loop() - - def after_create_session(self, session, coord): - # This is not necessary as we do not run infeed enqueue and outfeed dequeue - # in side threads for prediction model. But it makes the - # TPUInfeedOutfeedSessionHook prints nice message. - self._iterations_per_loop_var.load(1, session=session) - - def before_run(self, run_context): - return session_run_hook.SessionRunArgs(self._scalar_stopping_signal) - - def after_run(self, run_context, run_values): - _ = run_context - scalar_stopping_signal = run_values.results - if _StopSignals.should_stop(scalar_stopping_signal): - # NOTE(xiejw): In prediction, stopping signals are inserted for each - # batch. And we append one more batch to signal the system it should stop. - # The data flow might look like - # - # batch 0: images, labels, stop = 0 (user provided) - # batch 1: images, labels, stop = 0 (user provided) - # ... - # batch 99: images, labels, stop = 0 (user provided) - # batch 100: images, labels, stop = 1 (TPUEstimator appended) - # - # where the final batch (id = 100) is appended by TPUEstimator, so we - # should drop it before returning the predictions to user. - # To achieve that, we throw the OutOfRangeError in after_run. Once - # Monitored Session sees this error in SessionRunHook.after_run, the - # "current" prediction, i.e., batch with id=100, will be discarded - # immediately - raise errors.OutOfRangeError(None, None, 'Stopped by stopping signal.') - - -def generate_per_core_enqueue_ops_fn_for_host( - ctx, input_fn, inputs_structure_recorder, host_device, host_id): - """Generates infeed enqueue ops for per-core input_fn on a single host.""" - captured_infeed_queue = _CapturedObject() - tpu_ordinal_function_impl = ctx.tpu_ordinal_function(host_id) - - def enqueue_ops_fn(): - """A fn returns enqueue_ops.""" - num_cores_per_host = ctx.num_of_cores_per_host - per_host_sharded_inputs = [] - for core_ordinal in range(num_cores_per_host): - with ops.name_scope('ordinal_%d' % (core_ordinal)): - user_context = tpu_context.TPUContext( - internal_ctx=ctx, - input_device=host_device, - invocation_index=host_id * ctx.num_of_cores_per_host + core_ordinal) - inputs = _Inputs.from_input_fn(input_fn(user_context)) - if inputs.is_dataset: - raise TypeError( - '`input_fn` returning `Dataset` is not yet supported in ' - 'per-Core input pipeline deployment yet. Please set ' - 'TPUConfig.per_host_input_for_training to True or return ' - '`features` and `labels` from `input_fn`') - features, labels = inputs.features_and_labels() - - inputs_structure_recorder.validate_and_record_structure( - features, labels) - flattened_inputs = ( - inputs_structure_recorder.flatten_features_and_labels( - features, labels)) - per_host_sharded_inputs.append(flattened_inputs) - - infeed_queue = tpu_feed.InfeedQueue( - number_of_tuple_elements=len(per_host_sharded_inputs[0])) - captured_infeed_queue.capture(infeed_queue) - - per_host_enqueue_ops = infeed_queue.generate_enqueue_ops( - per_host_sharded_inputs, tpu_ordinal_function=tpu_ordinal_function_impl) - return per_host_enqueue_ops - - return enqueue_ops_fn, captured_infeed_queue - - -def generate_per_host_enqueue_ops_fn_for_host( - ctx, input_fn, inputs_structure_recorder, batch_axis, device, host_id): - """Generates infeed enqueue ops for per-host input_fn on a single host.""" - captured_infeed_queue = _CapturedObject() - - dataset_initializer = None - - with ops.device(device): - user_context = tpu_context.TPUContext( - internal_ctx=ctx, input_device=device, invocation_index=host_id) - inputs = _Inputs.from_input_fn(input_fn(user_context)) - - is_dataset = inputs.is_dataset - if ctx.mode == model_fn_lib.ModeKeys.PREDICT: - if not is_dataset: - raise TypeError( - 'For mode PREDICT, `input_fn` must return `Dataset` instead of ' - '`features` and `labels`.') - if batch_axis is not None: - raise TypeError('For mode PREDICT, batch_axis is not supported yet.') - inputs = _InputsWithStoppingSignals( - dataset=inputs.dataset, - batch_size=ctx.batch_size_for_input_fn, - add_padding=True) - - if is_dataset: - dataset_initializer = inputs.dataset_initializer() - - tpu_ordinal_function_impl = ctx.tpu_ordinal_function(host_id) - - def enqueue_ops_fn(): - """A Fn returning the TPU infeed enqueue ops. - - By providing as a Fn, it can be invoked inside the tf.while_loop such that - the input pipeline for multiple iterations can be executed by one - Session.run call. - - Returns: - list of dict of ops. - """ - with ops.device(device): - num_of_replicas_per_host = ctx.num_of_replicas_per_host - # Convert user input to features and labels. If the user returns a - # dataset, it is initialized and the features and labels extracted via - # `dataset.iterator.get_next()` - features, labels = inputs.features_and_labels() - signals = inputs.signals() - - inputs_structure_recorder.validate_and_record_structure(features, labels) - unsharded_tensor_list = ( - inputs_structure_recorder.flatten_features_and_labels( - features, labels, signals)) - - infeed_queue = tpu_feed.InfeedQueue( - tuple_types=[t.dtype for t in unsharded_tensor_list], - tuple_shapes=[t.shape for t in unsharded_tensor_list], - shard_dimensions=batch_axis) - captured_infeed_queue.capture(infeed_queue) - infeed_queue.set_number_of_shards(num_of_replicas_per_host) - per_host_enqueue_ops = ( - infeed_queue.split_inputs_and_generate_enqueue_ops( - unsharded_tensor_list, - placement_function=lambda x: device, - tpu_ordinal_function=tpu_ordinal_function_impl)) - if signals is None: - return per_host_enqueue_ops - else: - return { - 'ops': per_host_enqueue_ops, - 'signals': signals, - } - - return enqueue_ops_fn, captured_infeed_queue, dataset_initializer - - -def generate_per_host_v2_enqueue_ops_fn_for_host( - ctx, input_fn, inputs_structure_recorder, device, host_id): - """Generates infeed enqueue ops for per-host input_fn on a single host.""" - captured_infeed_queue = _CapturedObject() - dataset_initializer = None - - with ops.device(device): - user_context = tpu_context.TPUContext( - internal_ctx=ctx, input_device=device, invocation_index=host_id) - inputs = _Inputs.from_input_fn(input_fn(user_context)) - - is_dataset = inputs.is_dataset - if not is_dataset: - raise TypeError('`input_fn` must return a `Dataset` for the PER_HOST_V2 ' - 'input pipeline configuration.') - - if ctx.mode == model_fn_lib.ModeKeys.PREDICT: - inputs = _InputsWithStoppingSignals( - dataset=inputs.dataset, - batch_size=ctx.batch_size_for_input_fn, - add_padding=True, - num_invocations_per_step=ctx.num_of_replicas_per_host) - - dataset_initializer = inputs.dataset_initializer() - tpu_ordinal_function_impl = ctx.tpu_ordinal_function(host_id) - - def enqueue_ops_fn(): - """Generates the per_host enqueue ops.""" - control_deps = [] - per_host_sharded_inputs = [] - sparse_features_list = [] - num_replicas_per_host = ctx.num_of_replicas_per_host - cached_signals = None - with ops.device(device): - if not inputs.is_dataset: - raise TypeError('`input_fn` must return a `Dataset` for this mode.') - for _ in range(num_replicas_per_host): - # Use control dependencies to ensure a deterministic ordering. - with ops.control_dependencies(control_deps): - features, labels = inputs.features_and_labels() # Calls get_next() - signals = inputs.signals() - - # All the replicas share the replica 0's stopping singal. - # This avoids inconsistent state among different model replcias. - if cached_signals: - signals['stopping'] = cached_signals['stopping'] - else: - cached_signals = signals - - features, labels, sparse_features = ( - _tpu_estimator_embedding.split_inputs(ctx, features, labels)) - sparse_features_list.append(sparse_features) - - inputs_structure_recorder.validate_and_record_structure( - features, labels) - flattened_inputs = ( - inputs_structure_recorder.flatten_features_and_labels( - features, labels, signals)) - control_deps.extend(flattened_inputs) - per_host_sharded_inputs.append(flattened_inputs) - - if inputs_structure_recorder.flattened_input_dims: - input_partition_dims = inputs_structure_recorder.flattened_input_dims - if signals: - input_partition_dims += [None] * len(signals) - # pylint: disable=protected-access - infeed_queue = tpu_feed._PartitionedInfeedQueue( - number_of_tuple_elements=len(per_host_sharded_inputs[0]), - host_id=host_id, - input_partition_dims=input_partition_dims, - device_assignment=ctx.device_assignment) - per_host_enqueue_ops = infeed_queue.generate_enqueue_ops( - per_host_sharded_inputs) - else: - infeed_queue = tpu_feed.InfeedQueue( - number_of_tuple_elements=len(per_host_sharded_inputs[0])) - per_host_enqueue_ops = infeed_queue.generate_enqueue_ops( - per_host_sharded_inputs, - tpu_ordinal_function=tpu_ordinal_function_impl) - captured_infeed_queue.capture(infeed_queue) - - if ctx.embedding_config: - per_host_enqueue_ops.extend( - ctx.embedding_config.tpu_embedding.generate_enqueue_ops( - sparse_features_list)) - - if signals is None: - return per_host_enqueue_ops - else: - return { - 'ops': per_host_enqueue_ops, - 'signals': signals, - } - - return enqueue_ops_fn, captured_infeed_queue, dataset_initializer - - -def generate_broadcast_enqueue_ops_fn(ctx, input_fn, inputs_structure_recorder, - num_hosts): - """Generates infeed enqueue ops for one input_fn on all the hosts.""" - captured_infeed_queue = _CapturedObject() - dataset_initializer = None - device_0 = ctx.tpu_host_placement_function(host_id=0) - with ops.device(device_0): - user_context = tpu_context.TPUContext( - internal_ctx=ctx, input_device=device_0, invocation_index=0) - inputs = _Inputs.from_input_fn(input_fn(user_context)) - - is_dataset = inputs.is_dataset - if ctx.mode == model_fn_lib.ModeKeys.PREDICT: - if not is_dataset: - raise TypeError( - 'For mode PREDICT, `input_fn` must return `Dataset` instead of ' - '`features` and `labels`.') - - inputs = _InputsWithStoppingSignals( - dataset=inputs.dataset, - batch_size=ctx.batch_size_for_input_fn, - add_padding=True) - - if is_dataset: - dataset_initializer = inputs.dataset_initializer() - num_replicas_per_host = ctx.num_of_replicas_per_host - - def tpu_ordinal_function_impl(replica_id): - if ctx.device_assignment: - return ctx.device_assignment.tpu_ordinal(replica=replica_id) - else: - return replica_id % num_replicas_per_host - - def device_function_impl(replica_id): - return ctx.tpu_host_placement_function(replica_id=replica_id) - - def enqueue_ops_fn(): - """Generates enqueue ops for all the hosts.""" - broadcasted_inputs = [] - flattened_inputs = None # Cache result from input_fn. - signals = None - for host_id in xrange(num_hosts): - with ops.device(ctx.tpu_host_placement_function(host_id=host_id)): - for _ in xrange(ctx.num_of_replicas_per_host): - # Note: input_fn is only called once at host 0 for the first replica. - # The features and labels returned from that invocation are - # broadcasted to other replicas(including the replicas on other - # hosts). - if flattened_inputs is None: - features, labels = inputs.features_and_labels() # Calls get_next() - signals = inputs.signals() - - inputs_structure_recorder.validate_and_record_structure( - features, labels) - flattened_inputs = ( - inputs_structure_recorder.flatten_features_and_labels( - features, labels, signals)) - broadcasted_inputs.append(flattened_inputs) - - infeed_queue = tpu_feed.InfeedQueue( - number_of_tuple_elements=len(broadcasted_inputs[0])) - captured_infeed_queue.capture(infeed_queue) - enqueue_ops = infeed_queue.generate_enqueue_ops( - broadcasted_inputs, - tpu_ordinal_function=tpu_ordinal_function_impl, - placement_function=device_function_impl) - - if signals is None: - return enqueue_ops - else: - return { - 'ops': enqueue_ops, - 'signals': signals, - } - - return enqueue_ops_fn, captured_infeed_queue, dataset_initializer - - -class _InputPipeline(object): - """`_InputPipeline` handles invoking `input_fn` and piping to infeed queue. - - `_InputPipeline` abstracts the per-core/per-host `input_fn` invocation from - call site. To be precise, based on the configuration in - `_InternalTPUContext`, it invokes `input_fn` for all cores (usually - multi-host TPU training) or for one host (usually for single-host TPU - evaluation), and sends all `features` and `labels` returned by `input_fn` to - TPU infeed. For per-core invocation, `features` and `labels` are piped to - infeed directly, one tuple for each core. For per-host invocation, `features` - and `labels` are split at host (with respect to `batch_axis`) and piped to all - cores accordingly. - - In addition, flatten/unflatten are handled by `_InputPipeline` also. Model - inputs returned by the `input_fn` can have one of the following forms: - 1. features - 2. (features, labels) - 3. ((arbitrarily nested structure of features), labels) - - Internally, form 1 is reformed to `(features, None)` as features and labels - are passed separately to underlying methods. For TPU training, TPUEstimator - may expect multiple `features` and `labels` tuples one for each core. - - TPUEstimator allows various different structures for inputs (namely `features` - and `labels`). Both `features` and `labels` can be any nested sturcture - supported by TF nest (namely, dict, tuples, namedtuples or any nested - structure of such of Tensors). `labels` could be `None` as well. - - These are flattened before they are passed to the infeed/outfeed library - as that expectes flattend lists. - """ - - class InputsStructureRecorder(object): - """The recorder to record inputs structure.""" - - def __init__(self, input_partition_dims=None): - # Holds the structure of inputs - self._feature_structure = {} - self._flattened_input_dims = None - - if input_partition_dims: - # This should have been validated in TPUConfig. - assert len(input_partition_dims) <= 2, 'must have 1 or 2 elements.' - if len(input_partition_dims) == 2: - self._feature_dims, self._label_dims = input_partition_dims - else: - self._feature_dims = input_partition_dims[0] - self._label_dims = None - - assert self._feature_dims is not None, ('input_partition_dims[0] must ' - 'not be None') - else: - self._feature_dims = None - self._label_dims = None - - # Internal state. - self._initialized = False - - @property - def flattened_input_dims(self): - assert self._initialized, 'InputsStructureRecorder is not initialized.' - return self._flattened_input_dims - - def has_labels(self): - return 'labels' in self._feature_structure - - def _flatten_input_dims(self, feature_dims, feature_dims_names, label_dims, - label_dims_names, label_names, has_labels): - """Flatten input dims with the same order as flattened input tensors.""" - flattened_input_dims = [] - if feature_dims_names: - # We need a fixed ordering for matching the tensors in features. - flattened_input_dims.extend( - [feature_dims[name] for name in feature_dims_names]) - else: - flattened_input_dims.append(feature_dims) - - if label_dims_names: - # We need a fixed ordering for matching the tensors in labels. - flattened_input_dims.extend( - [label_dims[name] for name in label_dims_names]) - else: - if label_names: - num_tensors_in_label = len(label_names) - else: - num_tensors_in_label = int(has_labels) - # Setting `None` in input_partition_dims[1] will apply `None` to - # all the tensors in labels, regardless of internal structure. - flattened_input_dims.extend([label_dims] * num_tensors_in_label) - - return flattened_input_dims - - def validate_and_record_structure(self, features, labels): - """Validates and records the structure of `features` and `labels`.""" - # Extract structure. - has_labels = labels is not None - feature_names = _extract_key_names(features) - label_names = _extract_key_names(labels) - - if not self._initialized: - # Record structure. - self._initialized = True - if self._feature_dims is not None: - feature_dims_names = _extract_key_names(self._feature_dims) - if feature_dims_names != feature_names: - raise ValueError( - 'TPUConfig.input_partition_dims[0] mismatched feature' - ' keys. Expected {}, got {}'.format(feature_names, - feature_dims_names)) - - label_dims_names = _extract_key_names(self._label_dims) - if self._label_dims is not None and label_dims_names != label_names: - raise ValueError( - 'TPUConfig.input_partition_dims[1] mismatched label' - ' keys. Expected {}, got {}'.format(label_names, - label_dims_names)) - - self._flattened_input_dims = self._flatten_input_dims( - self._feature_dims, feature_dims_names, self._label_dims, - label_dims_names, label_names, has_labels) - - def flatten_features_and_labels(self, features, labels, signals=None): - """Flattens the `features` and `labels` to a single tensor list.""" - self._feature_structure['features'] = features - if labels is not None: - self._feature_structure['labels'] = labels - if signals is not None: - self._feature_structure['signals'] = signals - return data_nest.flatten(self._feature_structure) - - def unflatten_features_and_labels(self, flattened_inputs): - """Restores the flattened inputs to original features and labels form. - - Args: - flattened_inputs: Flattened inputs for each shard. - - Returns: - A tuple of (`features`, `labels`), where `labels` could be None. - Each one, if present, should have identical structure (single tensor vs - dict) as the one returned by input_fn. - - Raises: - ValueError: If the number of expected tensors from `flattened_inputs` - mismatches the recorded structure. - """ - - unflattened_inputs = data_nest.pack_sequence_as(self._feature_structure, - flattened_inputs) - return _Inputs( - unflattened_inputs['features'], - unflattened_inputs.get('labels'), - signals=unflattened_inputs.get('signals')) - - def __init__(self, input_fn, batch_axis, ctx): - """Constructor. - - Args: - input_fn: input fn for train or eval. - batch_axis: A python tuple of int values describing how each tensor - produced by the Estimator `input_fn` should be split across the TPU - compute shards. - ctx: A `_InternalTPUContext` instance with mode. - - Raises: - ValueError: If both `sharded_features` and `num_cores` are `None`. - """ - self._inputs_structure_recorder = _InputPipeline.InputsStructureRecorder( - ctx.input_partition_dims) - - self._sharded_per_core = ctx.is_input_sharded_per_core() - self._input_fn = input_fn - self._infeed_queue = None - self._ctx = ctx - self._batch_axis = batch_axis - - def generate_infeed_enqueue_ops_and_dequeue_fn(self): - """Generates infeed enqueue ops and dequeue_fn.""" - # While tf.while_loop is called, the body function, which invokes - # `enqueue_fn` passed in, is called to construct the graph. So, input_fn - # structure is recorded. - enqueue_ops, all_hooks, run_infeed_loop_on_coordinator = ( - self._invoke_input_fn_and_record_structure()) - - self._validate_input_pipeline() - - def dequeue_fn(): - """dequeue_fn is used by TPU to retrieve the tensors.""" - # In the model-parallel case, both the host-side and device-side - # computations must agree on the core on which infeed takes place. We - # choose to perform infeed on logical core 0 of each replica. - values = self._infeed_queue.generate_dequeue_op(tpu_device=0) - # The unflatten process uses the structure information recorded above. - return self._inputs_structure_recorder.unflatten_features_and_labels( - values) - - return (enqueue_ops, dequeue_fn, all_hooks, run_infeed_loop_on_coordinator) - - def _invoke_input_fn_and_record_structure(self): - """Deploys the input pipeline and record input structure.""" - enqueue_ops = [] - infeed_queues = [] - all_dataset_initializers = [] - num_hosts = self._ctx.num_hosts - tpu_host_placement_fn = self._ctx.tpu_host_placement_function - - run_infeed_loop_on_coordinator = True - - if self._sharded_per_core: - # Per-Core input pipeline deployment. - # Invoke input pipeline for each core and placed on the corresponding - # host. - for host_id in range(num_hosts): - host_device = tpu_host_placement_fn(host_id=host_id) - with ops.device(host_device): - with ops.name_scope('input_pipeline_task%d' % (host_id)): - enqueue_ops_fn, captured_infeed_queue = ( - generate_per_core_enqueue_ops_fn_for_host( - self._ctx, self._input_fn, self._inputs_structure_recorder, - host_device, host_id)) - - if _WRAP_INPUT_FN_INTO_WHILE_LOOP: - run_infeed_loop_on_coordinator = False - enqueue_ops.append( - _wrap_computation_in_while_loop( - device=host_device, op_fn=enqueue_ops_fn)) - else: - enqueue_ops.append(enqueue_ops_fn()) - # Infeed_queue_getter must be called after enqueue_ops_fn is called. - infeed_queues.append(captured_infeed_queue.get()) - - elif self._ctx.is_input_broadcast_with_iterators(): - # Only calls input_fn in host 0. - host_device = tpu_host_placement_fn(host_id=0) - enqueue_ops_fn, captured_infeed_queue, dataset_initializer = ( - generate_broadcast_enqueue_ops_fn(self._ctx, self._input_fn, - self._inputs_structure_recorder, - num_hosts)) - if dataset_initializer: - all_dataset_initializers.append(dataset_initializer) - run_infeed_loop_on_coordinator = False - wrap_fn = ( - _wrap_computation_in_while_loop - if self._ctx.mode != model_fn_lib.ModeKeys.PREDICT else - _wrap_computation_in_while_loop_with_stopping_signals) - enqueue_ops.append(wrap_fn(device=host_device, op_fn=enqueue_ops_fn)) - else: - enqueue_ops.append(enqueue_ops_fn()) - infeed_queues.append(captured_infeed_queue.get()) - else: - for host_id in range(num_hosts): - host_device = tpu_host_placement_fn(host_id=host_id) - with ops.device(host_device): - with ops.name_scope('input_pipeline_task%d' % (host_id)): - if self._ctx.is_input_per_host_with_iterators(): - enqueue_ops_fn, captured_infeed_queue, dataset_initializer = ( - generate_per_host_v2_enqueue_ops_fn_for_host( - self._ctx, self._input_fn, - self._inputs_structure_recorder, host_device, host_id)) - else: - enqueue_ops_fn, captured_infeed_queue, dataset_initializer = ( - generate_per_host_enqueue_ops_fn_for_host( - self._ctx, self._input_fn, - self._inputs_structure_recorder, self._batch_axis, - host_device, host_id)) - - # NOTE(xiejw): We dispatch here based on the return type of the - # users `input_fn`. - # - # 1. If input_fn returns a Dataset instance, we initialize the - # iterator outside of tf.while_loop, and call the iterator.get_next - # inside tf.while_loop. This should be always safe. - # - # 2. If input_fn returns (features, labels), it is too late to wrap - # them inside tf.while_loop, as resource initialization cannot be - # handled in TF control flow properly. In this case, we will use - # python loop to enqueue the data into TPU system. This may be - # slow compared to the previous case. - if dataset_initializer: - all_dataset_initializers.append(dataset_initializer) - run_infeed_loop_on_coordinator = False - wrap_fn = ( - _wrap_computation_in_while_loop - if self._ctx.mode != model_fn_lib.ModeKeys.PREDICT else - _wrap_computation_in_while_loop_with_stopping_signals) - enqueue_ops.append( - wrap_fn(device=host_device, op_fn=enqueue_ops_fn)) - else: - enqueue_ops.append(enqueue_ops_fn()) - infeed_queues.append(captured_infeed_queue.get()) - # infeed_queue is used to generate dequeue ops. The only thing it uses for - # dequeue is dtypes and types. So, any one can be used. Here, grab the - # first one. - self._infeed_queue = infeed_queues[0] - return enqueue_ops, [ - util_lib.MultiHostDatasetInitializerHook(all_dataset_initializers) - ], run_infeed_loop_on_coordinator - - def _validate_input_pipeline(self): - """Validates the input pipeline. - - Perform some sanity checks to log user friendly information. We should - error out to give users better error message. But, if - _WRAP_INPUT_FN_INTO_WHILE_LOOP is False (legacy behavior), we cannot break - user code, so, log a warning. - - Raises: - RuntimeError: If the validation failed. - """ - if ops.get_default_graph().get_collection(ops.GraphKeys.QUEUE_RUNNERS): - err_msg = ('Input pipeline contains one or more QueueRunners. ' - 'It could be slow and not scalable. Please consider ' - 'converting your input pipeline to use `tf.data` instead (see ' - 'https://www.tensorflow.org/guide/datasets for ' - 'instructions.') - if _WRAP_INPUT_FN_INTO_WHILE_LOOP: - raise RuntimeError(err_msg) - else: - logging.warn(err_msg) - - -def call_computation(computation, - experimental_exported_model_uses_all_cores=True): - """Call computation. - - computation uses a single-core for TPU inference. If - `experimental_exported_model_uses_all_cores` is `True`, this function will - round-robin - computation among all TPU cores visible to the host; otherwise, it will use - a single core. - - Args: - computation: A Python function that takes no inputs and builds computation - graph. If `computation` returns m outputs, this function will return a - list of m Tensors. - experimental_exported_model_uses_all_cores: Whether to round-robin among all - cores visible to the host, or to use a single core. - - Returns: - A list of output tensors. - """ - if experimental_exported_model_uses_all_cores: - # Using `TPUPartitionedCall` makes it possible to target a different - # TPU core with every `Session.run()` call. Note that the entire inference - # graph executes on a single core, and that invocations of this graph - # will round-robin among the cores attached to a host. - @function.Defun(capture_resource_var_by_value=False) - def tpu_subgraph(): - return computation() - - return tpu_functional.TPUPartitionedCall( - args=tpu_subgraph.captured_inputs, - device_ordinal=tpu_ops.tpu_ordinal_selector(), - Tout=[o.type for o in tpu_subgraph.definition.signature.output_arg], - f=tpu_subgraph) - else: - return computation() - - -class _ModelFnWrapper(object): - """A `model_fn` wrapper. - - This makes calling model_fn on CPU and TPU easier and more consistent and - performs necessary check and mutation required by TPU training and evaluation. - - In addition, this wrapper manages converting the `model_fn` to a single TPU - train and eval step. - """ - - def __init__(self, model_fn, config, params, ctx): - self._model_fn = model_fn - self._config = config - self._params = params - self._ctx = ctx - - def call_without_tpu(self, features, labels, is_export_mode): - return self._call_model_fn(features, labels, is_export_mode=is_export_mode) - - def _add_embedding_features(self, features, hook_dummy_table_variables): - """Add embedding features, optionally add hook to intercept gradient.""" - if self._ctx.embedding_config: - tpu_embedding_ = self._ctx.embedding_config.tpu_embedding - embedding_activations = tpu_embedding_.get_activations() - if hook_dummy_table_variables: - new_embedding_activations = ( - tpu_embedding_gradient.hook_dummy_table_variables_to_activations( - tpu_embedding_, embedding_activations, - self._ctx.embedding_config.dummy_table_variables)) - features.update(new_embedding_activations) - else: - features.update(embedding_activations) - - def convert_to_single_tpu_train_step(self, dequeue_fn): - """Converts user provided model_fn` as a single train step on TPU. - - The user provided `model_fn` takes input tuple - (features, labels) and produces the EstimatorSpec with train_op and loss for - train `mode`. This usually represents a single train computation on CPU. - - For TPU training, a train (computation) step is first wrapped in a - tf.while_loop control flow to repeat for many times and then replicated to - all TPU shards. Besides the input should be taken from TPU infeed rather - than input pipeline (input_fn) directly. To fit TPU loop and replicate - pattern, the original train computation should be reformed, which is the - returned `train_step`. - - Args: - dequeue_fn: The function to retrieve inputs, features and labels, from TPU - infeed dequeue channel. - - Returns: - A tuple of train_fn, host_calls, and captured scaffold_fn. The train_fn - representing the train step for TPU. - """ - - host_call = _OutfeedHostCall(self._ctx) - captured_scaffold_fn = _CapturedObject() - captured_training_hooks = _CapturedObject() - - def train_step(loss): - """Training step function for use inside a while loop.""" - del loss # unused; required in function signature. - inputs = dequeue_fn() - features, labels = inputs.features_and_labels() - self._add_embedding_features(features, True) - - estimator_spec = self._verify_estimator_spec( - self._call_model_fn(features, labels)) - loss, train_op = estimator_spec.loss, estimator_spec.train_op - - if isinstance(estimator_spec, model_fn_lib._TPUEstimatorSpec): # pylint: disable=protected-access - captured_scaffold_fn.capture(estimator_spec.scaffold_fn) - else: - captured_scaffold_fn.capture(None) - - captured_training_hooks.capture(estimator_spec.training_hooks) - - if self._ctx.embedding_config is None: - apply_sparse_grads = [] - else: - tpu_embedding_ = self._ctx.embedding_config.tpu_embedding - gradients = ( - tpu_embedding_gradient.get_gradients_through_dummy_table_variables( - tpu_embedding_) - ) - apply_sparse_grads = [ - tpu_embedding_.generate_send_gradients_op(gradients) - ] - - # We must run train_op to update the variables prior to running the - # outfeed. - with ops.control_dependencies([train_op] + apply_sparse_grads): - host_call_outfeed_ops = [] - if (isinstance(estimator_spec, model_fn_lib._TPUEstimatorSpec) # pylint: disable=protected-access - and estimator_spec.host_call is not None): - host_call.record({'host_call': estimator_spec.host_call}) - host_call_outfeed_ops = host_call.create_enqueue_op() - with ops.control_dependencies(host_call_outfeed_ops): - return array_ops.identity(loss) - - return (train_step, host_call, captured_scaffold_fn, - captured_training_hooks) - - def convert_to_single_tpu_eval_step(self, dequeue_fn): - """Converts user provided model_fn` as a single eval step on TPU. - - Similar to training, the user provided `model_fn` takes input tuple - (features, labels) and produces the TPUEstimatorSpec with eval_metrics for - eval `mode`. This usually represents a single evaluation computation on CPU. - - For TPU evaluation, a eval (computation) step is first wrapped in a - tf.while_loop control flow to repeat for many times and then replicated to - all TPU shards. Besides the input and output are slightly different. Input, - features and labels, should be taken from TPU infeed rather than input - pipeline (input_fn) directly. Output is managed in two stages. First, the - model outputs as the result of evaluation computation, usually model logits, - should be transferred from TPU system to CPU. Then, all model outputs are - concatenated first on CPU and sent to the metric_fn for metrics computation. - To fit TPU evaluation pattern, the original eval computation should be - reformed, which is the returned `eval_step`. - - Args: - dequeue_fn: The function to retrieve inputs, features and labels, from TPU - infeed dequeue channel. - - Returns: - A tuple of eval_fn, host_calls, and captured scaffold_fn. The eval_fn - representing the eval step for TPU. - """ - host_calls = _OutfeedHostCall(self._ctx) - captured_scaffold_fn = _CapturedObject() - captured_eval_hooks = _CapturedObject() - - def eval_step(total_loss): - """Evaluation step function for use inside a while loop.""" - inputs = dequeue_fn() - features, labels = inputs.features_and_labels() - self._add_embedding_features(features, False) - - tpu_estimator_spec = self._call_model_fn(features, labels) - if not isinstance(tpu_estimator_spec, model_fn_lib._TPUEstimatorSpec): # pylint: disable=protected-access - raise RuntimeError( - 'estimator_spec used by TPU evaluation must have type' - '`TPUEstimatorSpec`. Got {}'.format(type(tpu_estimator_spec))) - - loss = tpu_estimator_spec.loss - captured_scaffold_fn.capture(tpu_estimator_spec.scaffold_fn) - captured_eval_hooks.capture(tpu_estimator_spec.evaluation_hooks) - - to_record = {} - if tpu_estimator_spec.eval_metrics: - to_record['eval_metrics'] = tpu_estimator_spec.eval_metrics - if tpu_estimator_spec.host_call is not None: - # We assume that evaluate won't update global step, so we don't wrap - # this host_call. - to_record['host_call'] = tpu_estimator_spec.host_call - host_calls.record(to_record) - - with ops.control_dependencies(host_calls.create_enqueue_op()): - return math_ops.add(total_loss, loss) - - return eval_step, host_calls, captured_scaffold_fn, captured_eval_hooks - - def convert_to_single_tpu_predict_step(self, dequeue_fn): - """Converts user provided model_fn` as a single predict step on TPU. - - Args: - dequeue_fn: The function to retrieve inputs, features and labels, from TPU - infeed dequeue channel. - - Returns: - A tuple of predict_fn, host_calls, and captured scaffold_fn. The - predict_fn representing the predict step for TPU. - """ - host_calls = _OutfeedHostCall(self._ctx) - captured_scaffold_fn = _CapturedObject() - captured_predict_hooks = _CapturedObject() - - def predict_step(unused_scalar_stopping_signal): - """Evaluation step function for use inside a while loop.""" - inputs = dequeue_fn() - features, labels = inputs.features_and_labels() - stopping_signals = inputs.signals() - - assert stopping_signals is not None, ( - 'Internal Error: `signals` is missing.') - - tpu_estimator_spec = self._call_model_fn( - features, labels, is_export_mode=False) - if not isinstance(tpu_estimator_spec, model_fn_lib._TPUEstimatorSpec): # pylint: disable=protected-access - raise RuntimeError( - 'estimator_spec used by TPU prediction must have type' - '`TPUEstimatorSpec`. Got {}'.format(type(tpu_estimator_spec))) - - self._verify_tpu_spec_predictions(tpu_estimator_spec.predictions) - - captured_scaffold_fn.capture(tpu_estimator_spec.scaffold_fn) - captured_predict_hooks.capture(tpu_estimator_spec.prediction_hooks) - to_record = {} - identity_fn = lambda **kwargs: kwargs - to_record['predictions'] = [identity_fn, tpu_estimator_spec.predictions] - to_record['signals'] = [identity_fn, stopping_signals] - if tpu_estimator_spec.host_call is not None: - to_record['host_call'] = tpu_estimator_spec.host_call - host_calls.record(to_record) - - with ops.control_dependencies(host_calls.create_enqueue_op()): - return _StopSignals.as_scalar_stopping_signal(stopping_signals) - - return (predict_step, host_calls, captured_scaffold_fn, - captured_predict_hooks) - - def _verify_tpu_spec_predictions(self, predictions): - """Validates TPUEstimatorSpec.predictions dict.""" - # TODO(xiejw): Adds validation for prediction dictionrary. - # TODO(xiejw): Adds support for single tensor as predictions. - if not isinstance(predictions, dict): - raise TypeError('TPUEstimatorSpec.predictions must be dict of Tensors.') - - for (key, tensor) in predictions.items(): - if tensor.shape.dims[0].value is None: - raise ValueError( - 'The tensor with key ({}) in TPUEstimatorSpec.predictions has ' - 'dynamic shape (should be static). Tensor: {}'.format(key, tensor)) - return predictions - - def _validate_model_features_and_labels(self, features, labels, - is_export_mode): - """Validates that the features and labels for the model function are valid. - - A valid features/labels object is the one with: - - Type: A tensor or any nested structure of tensors supported by TF nest, - namely nested dictionary, tuple, namedtuple, or sequence of tensors. - - Static shape if is_export_mode is False. - - Args: - features: the features that would be input to the model function. - labels: the labels that would be input to the model function. - is_export_mode: boolean value specifying if in export mode. - - Raises: - TypeError: If features/labels are not of the correct type. - ValueError: If features/labels have dynamic shape. - """ - - def validate(obj, obj_name): - """Helper validate function.""" - if is_export_mode or self._ctx.is_running_on_cpu(is_export_mode): - return - if isinstance(obj, ops.Tensor): - if not obj.get_shape().is_fully_defined(): - raise ValueError( - 'The {} to the model returned by input_fn must have static shape.' - ' Tensor: {}'.format(obj_name, obj)) - else: - for tensor in data_nest.flatten(obj): - if not tensor.get_shape().is_fully_defined(): - raise ValueError( - ('The {} to the model returned by input_fn must have static ' - 'shape. Tensor: {}').format(obj_name, tensor)) - - validate(features, 'features') - if labels is not None: - validate(labels, 'labels') - - def _call_model_fn(self, features, labels, is_export_mode=False): - """Calls the model_fn with required parameters.""" - self._validate_model_features_and_labels(features, labels, is_export_mode) - model_fn_args = function_utils.fn_args(self._model_fn) - kwargs = {} - - # Makes deep copy with `config` and params` in case user mutates them. - config = copy.deepcopy(self._config) - params = copy.deepcopy(self._params) - - if 'labels' in model_fn_args: - kwargs['labels'] = labels - elif labels is not None: - raise ValueError( - 'model_fn does not take labels, but input_fn returns labels.') - if 'mode' in model_fn_args: - kwargs['mode'] = self._ctx.mode - if 'config' in model_fn_args: - kwargs['config'] = config - if 'params' in model_fn_args: - kwargs['params'] = params - - if 'params' not in model_fn_args: - raise ValueError('model_fn ({}) does not include params argument, ' - 'required by TPUEstimator to pass batch size as ' - 'params[\'batch_size\']'.format(self._model_fn)) - - if is_export_mode: - batch_size_for_model_fn = None - else: - batch_size_for_model_fn = self._ctx.batch_size_for_model_fn - - if batch_size_for_model_fn is not None: - _add_item_to_params(params, _BATCH_SIZE_KEY, batch_size_for_model_fn) - - running_on_cpu = self._ctx.is_running_on_cpu(is_export_mode) - _add_item_to_params(params, _USE_TPU_KEY, not running_on_cpu) - - if not running_on_cpu: - user_context = tpu_context.TPUContext( - internal_ctx=self._ctx, call_from_input_fn=False) - _add_item_to_params(params, _CTX_KEY, user_context) - - estimator_spec = self._model_fn(features=features, **kwargs) - if (running_on_cpu and - isinstance(estimator_spec, model_fn_lib._TPUEstimatorSpec)): # pylint: disable=protected-access - # The estimator_spec will be passed to `Estimator` directly, which expects - # type `EstimatorSpec`. - return estimator_spec.as_estimator_spec() - else: - return estimator_spec - - def _verify_estimator_spec(self, estimator_spec): - """Validates the estimator_spec.""" - if isinstance(estimator_spec, model_fn_lib._TPUEstimatorSpec): # pylint: disable=protected-access - return estimator_spec - - err_msg = '{} returned by EstimatorSpec is not supported in TPUEstimator.' - if estimator_spec.training_chief_hooks: - raise ValueError( - err_msg.format('training_chief_hooks') + 'If you want' + - ' to pass training hooks, please pass via training_hooks.') - - if estimator_spec.scaffold: - logging.warning('EstimatorSpec.Scaffold is ignored by TPU train/eval. ' - 'Please use TPUEstimatorSpec.') - return estimator_spec - - -class _OutfeedHostCall(object): - """Support for `eval_metrics` and `host_call` in TPUEstimatorSpec.""" - - def __init__(self, ctx): - self._ctx = ctx - self._names = [] - # All of these are dictionaries of lists keyed on the name. - self._host_fns = {} - self._tensor_keys = collections.defaultdict(list) - self._tensors = collections.defaultdict(list) - self._tensor_dtypes = collections.defaultdict(list) - self._tensor_shapes = collections.defaultdict(list) - - @staticmethod - def validate(host_calls): - """Validates the `eval_metrics` and `host_call` in `TPUEstimatorSpec`.""" - - for name, host_call in host_calls.items(): - if not isinstance(host_call, (tuple, list)): - raise ValueError('{} should be tuple or list'.format(name)) - if len(host_call) != 2: - raise ValueError('{} should have two elements.'.format(name)) - if not callable(host_call[0]): - raise TypeError('{}[0] should be callable.'.format(name)) - if not isinstance(host_call[1], (tuple, list, dict)): - raise ValueError('{}[1] should be tuple or list, or dict.'.format(name)) - - if isinstance(host_call[1], (tuple, list)): - fullargspec = tf_inspect.getfullargspec(host_call[0]) - fn_args = function_utils.fn_args(host_call[0]) - # wrapped_hostcall_with_global_step uses varargs, so we allow that. - if fullargspec.varargs is None and len(host_call[1]) != len(fn_args): - raise RuntimeError( - 'In TPUEstimatorSpec.{}, length of tensors {} does not match ' - 'method args of the function, which takes {}.'.format( - name, len(host_call[1]), len(fn_args))) - - @staticmethod - def create_cpu_hostcall(host_calls): - """Runs on the host_call on CPU instead of TPU when use_tpu=False.""" - - _OutfeedHostCall.validate(host_calls) - ret = {} - for name, host_call in host_calls.items(): - host_fn, tensors = host_call - if isinstance(tensors, (tuple, list)): - ret[name] = host_fn(*tensors) - else: - # Must be dict. - try: - ret[name] = host_fn(**tensors) - except TypeError as e: - logging.warning( - 'Exception while calling %s: %s. It is likely the tensors ' - '(%s[1]) do not match the ' - 'function\'s arguments', name, e, name) - raise - return ret - - def record(self, host_calls): - """Records the host_call structure.""" - - for name, host_call in host_calls.items(): - host_fn, tensor_list_or_dict = host_call - self._names.append(name) - self._host_fns[name] = host_fn - - if isinstance(tensor_list_or_dict, dict): - for (key, tensor) in six.iteritems(tensor_list_or_dict): - self._tensor_keys[name].append(key) - self._tensors[name].append(tensor) - self._tensor_dtypes[name].append(tensor.dtype) - self._tensor_shapes[name].append(tensor.shape) - else: - # List or tuple. - self._tensor_keys[name] = None - for tensor in tensor_list_or_dict: - self._tensors[name].append(tensor) - self._tensor_dtypes[name].append(tensor.dtype) - self._tensor_shapes[name].append(tensor.shape) - - def create_enqueue_op(self): - """Create the op to enqueue the recorded host_calls. - - Returns: - A list of enqueue ops, which is empty if there are no host calls. - """ - if not self._names: - return [] - - tensors = [] - # TODO(jhseu): Consider deduping tensors. - for name in self._names: - tensors.extend(self._tensors[name]) - - with ops.device(tpu.core(0)): - return [tpu_ops.outfeed_enqueue_tuple(tensors)] - - def create_tpu_hostcall(self): - """Sends the tensors through outfeed and runs the host_fn on CPU. - - The tensors are concatenated along dimension 0 to form a global tensor - across all shards. The concatenated function is passed to the host_fn and - executed on the first host. - - Returns: - A dictionary mapping name to the return type of the host_call by that - name. - - Raises: - RuntimeError: If outfeed tensor is scalar. - """ - if not self._names: - return {} - - ret = {} - # For each i, dequeue_ops[i] is a list containing the tensors from all - # shards. This list is concatenated later. - dequeue_ops = [] - tensor_dtypes = [] - tensor_shapes = [] - for name in self._names: - for _ in self._tensors[name]: - dequeue_ops.append([]) - for dtype in self._tensor_dtypes[name]: - tensor_dtypes.append(dtype) - for shape in self._tensor_shapes[name]: - tensor_shapes.append(shape) - - # Outfeed ops execute on each replica's first logical core. Note: we must - # constraint it such that we have at most one outfeed dequeue and enqueue - # per replica. - for i in xrange(self._ctx.num_replicas): - host_device, ordinal_id = self._ctx.device_for_replica(i) - with ops.device(host_device): - outfeed_tensors = tpu_ops.outfeed_dequeue_tuple( - dtypes=tensor_dtypes, - shapes=tensor_shapes, - device_ordinal=ordinal_id) - for j, item in enumerate(outfeed_tensors): - dequeue_ops[j].append(item) - - # Deconstruct dequeue ops. - flat_dequeue_ops = [] - for l in dequeue_ops: - flat_dequeue_ops.extend(l) - - dequeue_ops_by_name = {} - pos = 0 - for name in self._names: - dequeue_ops_by_name[name] = dequeue_ops[pos:pos + - len(self._tensors[name])] - pos += len(self._tensors[name]) - - def _call_host_fn(fn, *args, **kw): - context = CatchInvalidHostcallFunctions() - context.Enter() - result = fn(*args, **kw) - context.Exit() - context.ExitResult(result) - return result - - # It is assumed evaluation always happens on single host TPU system. So, - # place all ops on tpu host if possible. - # - # TODO(jhseu): Evaluate whether this is right for summaries. - with ops.device(self._ctx.tpu_host_placement_function(replica_id=0)): - for name in self._names: - dequeue_ops = dequeue_ops_by_name[name] - for i, item in enumerate(dequeue_ops): - if dequeue_ops[i][0].shape.ndims == 0: - raise RuntimeError( - 'All tensors outfed from TPU should preserve batch size ' - 'dimension, but got scalar {}'.format(dequeue_ops[i][0])) - # TODO(xiejw): Make the specification of the outfeed combinaton - # function more explicit and well-documented. We may want to give the - # user the option of concatenating along any axis. - if (self._ctx.config.tpu_config.per_host_input_for_training is - tpu_config.InputPipelineConfig.BROADCAST): - # If the infeed is in BROADCAST mode (each core recieving the same - # input), then we assume that the cores also produce identical - # copies of the same output, and we simply take the output from - # the first core. This mode is used by Mesh-TensorFlow. - with ops.control_dependencies(dequeue_ops[i]): - dequeue_ops[i] = array_ops.identity(dequeue_ops[i][0]) - else: - # Assume that the input has been batch-split and that axis 0 of the - # output tensors represents the batch size. Concatenate along - # the axis 0 to re-combine the batch. - dequeue_ops[i] = array_ops.concat(dequeue_ops[i], axis=0) - - if self._tensor_keys[name] is not None: - # The user-provided eval_metrics[1] is a dict. - dequeue_ops = dict(zip(self._tensor_keys[name], dequeue_ops)) - try: - ret[name] = _call_host_fn(self._host_fns[name], **dequeue_ops) - except TypeError as e: - logging.warning( - 'Exception while calling %s: %s. It is likely the tensors ' - '(%s[1]) do not match the ' - 'function\'s arguments', name, e, name) - raise - else: - ret[name] = _call_host_fn(self._host_fns[name], *dequeue_ops) - - # force all dequeue operations to be run if not consumed by the host calls - ret['__force_dequeue'] = control_flow_ops.group(*flat_dequeue_ops) - return ret - - -class _OutfeedHostCallHook(session_run_hook.SessionRunHook): - """Hook to run host calls when use_tpu=False.""" - - def __init__(self, tensors): - self._tensors = tensors - - def begin(self): - # We duplicate this code from the TPUInfeedOutfeedSessionHook rather than - # create a separate hook to guarantee execution order, because summaries - # need to be initialized before the outfeed thread starts. - # TODO(jhseu): Make a wrapper hook instead? - self._init_ops = contrib_summary.summary_writer_initializer_op() - # Get all the writer resources from the initializer, so we know what to - # flush. - self._finalize_ops = [] - for op in self._init_ops: - self._finalize_ops.append(contrib_summary.flush(writer=op.inputs[0])) - - def after_create_session(self, session, coord): - session.run(self._init_ops) - - def before_run(self, run_context): - return basic_session_run_hooks.SessionRunArgs(self._tensors) - - def end(self, session): - session.run(self._finalize_ops) - - -class ExamplesPerSecondHook(basic_session_run_hooks.StepCounterHook): - """Calculate and report global_step/sec and examples/sec during runtime.""" - - def __init__(self, - batch_size, - every_n_steps=100, - every_n_secs=None, - output_dir=None, - summary_writer=None): - self._batch_size = batch_size - super(ExamplesPerSecondHook, self).__init__( - every_n_steps=every_n_steps, - every_n_secs=every_n_secs, - output_dir=output_dir, - summary_writer=summary_writer) - - def _log_and_record(self, elapsed_steps, elapsed_time, global_step): - global_step_per_sec = elapsed_steps / elapsed_time - examples_per_sec = self._batch_size * global_step_per_sec - if self._summary_writer is not None: - global_step_summary = Summary(value=[ - Summary.Value(tag='global_step/sec', simple_value=global_step_per_sec) - ]) - example_summary = Summary(value=[ - Summary.Value(tag='examples/sec', simple_value=examples_per_sec) - ]) - self._summary_writer.add_summary(global_step_summary, global_step) - self._summary_writer.add_summary(example_summary, global_step) - logging.info('global_step/sec: %g', global_step_per_sec) - logging.info('examples/sec: %g', examples_per_sec) - - -class InstallSignalHandlerHook(session_run_hook.SessionRunHook): - """Change SIGINT (CTRL^C) handler to force quit the process. - - The default behavior often results in hanging processes. - The original handler is restored after training/evaluation. - """ - - def __init__(self): - self._signal_fn = signal.getsignal(signal.SIGINT) - - def before_run(self, run_context): - signal.signal(signal.SIGINT, signal.SIG_DFL) - - def end(self, session): - signal.signal(signal.SIGINT, self._signal_fn) - - -class TPUEstimator(estimator_lib.Estimator): - """Estimator with TPU support. - - TPUEstimator also supports training on CPU and GPU. You don't need to define - a separate `tf.estimator.Estimator`. - - TPUEstimator handles many of the details of running on TPU devices, such as - replicating inputs and models for each core, and returning to host - periodically to run hooks. - - TPUEstimator transforms a global batch size in params to a per-shard batch - size when calling the `input_fn` and `model_fn`. Users should specify - global batch size in constructor, and then get the batch size for each shard - in `input_fn` and `model_fn` by `params['batch_size']`. - - - For training, `model_fn` gets per-core batch size; `input_fn` may get - per-core or per-host batch size depending on `per_host_input_for_training` - in `TPUConfig` (See docstring for TPUConfig for details). - - - For evaluation and prediction, `model_fn` gets per-core batch size and - `input_fn` get per-host batch size. - - Evaluation - ========== - - `model_fn` should return `TPUEstimatorSpec`, which expects the `eval_metrics` - for TPU evaluation. If eval_on_tpu is False, the evaluation will execute on - CPU or GPU; in this case the following discussion on TPU evaluation does not - apply. - - `TPUEstimatorSpec.eval_metrics` is a tuple of `metric_fn` and `tensors`, where - `tensors` could be a list of any nested structure of `Tensor`s (See - `TPUEstimatorSpec` for details). `metric_fn` takes the `tensors` and returns - a dict from metric string name to the result of calling a metric function, - namely a `(metric_tensor, update_op)` tuple. - - One can set `use_tpu` to `False` for testing. All training, evaluation, and - predict will be executed on CPU. `input_fn` and `model_fn` will receive - `train_batch_size` or `eval_batch_size` unmodified as `params['batch_size']`. - - Current limitations: - -------------------- - - 1. TPU evaluation only works on a single host (one TPU worker) except - BROADCAST mode. - - 2. `input_fn` for evaluation should **NOT** raise an end-of-input exception - (`OutOfRangeError` or `StopIteration`). And all evaluation steps and all - batches should have the same size. - - Example (MNIST): - ---------------- - - ``` - # The metric Fn which runs on CPU. - def metric_fn(labels, logits): - predictions = tf.argmax(logits, 1) - return { - 'accuracy': tf.metrics.precision( - labels=labels, predictions=predictions), - } - - # Your model Fn which runs on TPU (eval_metrics is list in this example) - def model_fn(features, labels, mode, config, params): - ... - logits = ... - - if mode = tf.estimator.ModeKeys.EVAL: - return tpu_estimator.TPUEstimatorSpec( - mode=mode, - loss=loss, - eval_metrics=(metric_fn, [labels, logits])) - - # or specify the eval_metrics tensors as dict. - def model_fn(features, labels, mode, config, params): - ... - final_layer_output = ... - - if mode = tf.estimator.ModeKeys.EVAL: - return tpu_estimator.TPUEstimatorSpec( - mode=mode, - loss=loss, - eval_metrics=(metric_fn, { - 'labels': labels, - 'logits': final_layer_output, - })) - ``` - - Prediction - ========== - - Prediction on TPU is an experimental feature to support large batch inference. - It is not designed for latency-critical system. In addition, due to some - usability issues, for prediction with small dataset, CPU `.predict`, i.e., - creating a new `TPUEstimator` instance with `use_tpu=False`, might be more - convenient. - - Note: In contrast to TPU training/evaluation, the `input_fn` for prediction - *should* raise an end-of-input exception (`OutOfRangeError` or - `StopIteration`), which serves as the stopping signal to `TPUEstimator`. To be - precise, the ops created by `input_fn` produce one batch of the data. - The `predict()` API processes one batch at a time. When reaching the end of - the data source, an end-of-input exception should be raised by one of these - operations. The user usually does not need to do this manually. As long as the - dataset is not repeated forever, the `tf.data` API will raise an end-of-input - exception automatically after the last batch has been produced. - - Note: Estimator.predict returns a Python generator. Please consume all the - data from the generator so that TPUEstimator can shutdown the TPU system - properly for user. - - Current limitations: - -------------------- - 1. TPU prediction only works on a single host (one TPU worker). - - 2. `input_fn` must return a `Dataset` instance rather than `features`. In - fact, .train() and .evaluate() also support Dataset as return value. - - Example (MNIST): - ---------------- - ``` - height = 32 - width = 32 - total_examples = 100 - - def predict_input_fn(params): - batch_size = params['batch_size'] - - images = tf.random_uniform( - [total_examples, height, width, 3], minval=-1, maxval=1) - - dataset = tf.data.Dataset.from_tensor_slices(images) - dataset = dataset.map(lambda images: {'image': images}) - - dataset = dataset.batch(batch_size) - return dataset - - def model_fn(features, labels, params, mode): - # Generate predictions, called 'output', from features['image'] - - if mode == tf.estimator.ModeKeys.PREDICT: - return tf.contrib.tpu.TPUEstimatorSpec( - mode=mode, - predictions={ - 'predictions': output, - 'is_padding': features['is_padding'] - }) - - tpu_est = TPUEstimator( - model_fn=model_fn, - ..., - predict_batch_size=16) - - # Fully consume the generator so that TPUEstimator can shutdown the TPU - # system. - for item in tpu_est.predict(input_fn=input_fn): - # Filter out item if the `is_padding` is 1. - # Process the 'predictions' - ``` - - Exporting - ========= - - `export_savedmodel` exports 2 metagraphs, one with `tag_constants.SERVING`, - and another with `tag_constants.SERVING` and `tag_constants.TPU`. - At serving time, these tags are used to select metagraph to load. - - Before running the graph on TPU, TPU system needs to be initialized. If - TensorFlow Serving model-server is used, this is done automatically. If - not, please call `session.run(tpu.initialize_system())`. - - `tpu.outside_compilation` can be used to wrap TPU incompatible ops in - `model_fn`. - - Example: - ---------------- - - ``` - def model_fn(features, labels, mode, config, params): - ... - logits = ... - export_outputs = { - 'logits': export_output_lib.PredictOutput( - {'logits': logits}) - } - - def host_call(logits): - class_ids = math_ops.argmax(logits) - classes = string_ops.as_string(class_ids) - export_outputs['classes'] = - export_output_lib.ClassificationOutput(classes=classes) - - tpu.outside_compilation(host_call, logits) - - ... - ``` - - """ - - def __init__(self, - model_fn=None, - model_dir=None, - config=None, - params=None, - use_tpu=True, - train_batch_size=None, - eval_batch_size=None, - predict_batch_size=None, - batch_axis=None, - eval_on_tpu=True, - export_to_tpu=True, - export_to_cpu=True, - warm_start_from=None, - experimental_exported_model_uses_all_cores=False, - experimental_export_device_assignment=False, - experimental_embedding_config_spec=None): - """Constructs an `TPUEstimator` instance. - - Args: - model_fn: Model function as required by `Estimator` which returns - EstimatorSpec or TPUEstimatorSpec. `training_hooks`, 'evaluation_hooks', - and `prediction_hooks` must not capure any TPU Tensor inside the - model_fn. - model_dir: Directory to save model parameters, graph and etc. This can - also be used to load checkpoints from the directory into a estimator to - continue training a previously saved model. If `None`, the model_dir in - `config` will be used if set. If both are set, they must be same. If - both are `None`, a temporary directory will be used. - config: An `tpu_config.RunConfig` configuration object. Cannot be `None`. - params: An optional `dict` of hyper parameters that will be passed into - `input_fn` and `model_fn`. Keys are names of parameters, values are - basic python types. There are reserved keys for `TPUEstimator`, - including 'batch_size'. - use_tpu: A bool indicating whether TPU support is enabled. Currently, - - TPU training and evaluation respect this bit, but eval_on_tpu can - override execution of eval. See below. - Predict still happens on CPU. - train_batch_size: An int representing the global training batch size. - TPUEstimator transforms this global batch size to a per-shard batch - size, as params['batch_size'], when calling `input_fn` and `model_fn`. - Cannot be `None` if `use_tpu` is `True`. Must be divisible by total - number of replicas. - eval_batch_size: An int representing evaluation batch size. Must be - divisible by total number of replicas. - predict_batch_size: An int representing the prediction batch size. Must be - divisible by total number of replicas. - batch_axis: A python tuple of int values describing how each tensor - produced by the Estimator `input_fn` should be split across the TPU - compute shards. For example, if your input_fn produced (images, labels) - where the images tensor is in `HWCN` format, your shard dimensions would - be [3, 0], where 3 corresponds to the `N` dimension of your images - Tensor, and 0 corresponds to the dimension along which to split the - labels to match up with the corresponding images. If None is supplied, - and per_host_input_for_training is True, batches will be sharded based - on the major dimension. If tpu_config.per_host_input_for_training is - False or `PER_HOST_V2`, batch_axis is ignored. - eval_on_tpu: If False, evaluation runs on CPU or GPU. In this case, the - model_fn must return `EstimatorSpec` when called with `mode` as `EVAL`. - export_to_tpu: If True, `export_savedmodel()` exports a metagraph for - serving on TPU. Note that unsupported export modes such as EVAL will be - ignored. For those modes, only a CPU model will be exported. - Currently, export_to_tpu only supports PREDICT. - export_to_cpu: If True, `export_savedmodel()` exports a metagraph for - serving on CPU. - warm_start_from: Optional string filepath to a checkpoint or SavedModel to - warm-start from, or a `tf.estimator.WarmStartSettings` object to fully - configure warm-starting. If the string filepath is provided instead of - a `WarmStartSettings`, then all variables are warm-started, and it is - assumed that vocabularies and Tensor names are unchanged. - experimental_exported_model_uses_all_cores: Whether to round-robin among - all cores visible to the host which is serving the saved model, or to - use a single core. This is a temporary flag to enable using all TPU - cores for inference with TPUPartitionedCall(). Once outside compilation - is supported in TPUPartitionedCall(), this flag will be enabled by - default. - experimental_export_device_assignment: Whether to include the device - assignment in the exported model. Doing so is useful in case of model - parallel inference but will tie the exported model to the TPU topology - used to export the model. - experimental_embedding_config_spec: Optional EmbeddingConfigSpec instance - to support using TPU embedding. IT IS STILL WORK IN PROGRESS, SO PLEASE - DO NOT USE. - - Raises: - ValueError: `params` has reserved keys already. - """ - if config is None or not isinstance(config, tpu_config.RunConfig): - raise ValueError( - '`config` must be provided with type `tpu_config.RunConfig`') - - if params is not None and any(k in params for k in _RESERVED_PARAMS_KEYS): - raise ValueError('{} are reserved keys but existed in params {}.'.format( - _RESERVED_PARAMS_KEYS, params)) - - if use_tpu: - # Perform some very basic validations. More validations will be found in - # _InternalTPUContext. - if train_batch_size is None: - raise ValueError('`train_batch_size` cannot be `None`') - util_lib.check_positive_integer(train_batch_size, 'train_batch_size') - - if (config.tpu_config.per_host_input_for_training is - tpu_config.InputPipelineConfig.PER_SHARD_V1 and - config.tpu_config.num_cores_per_replica): - raise ValueError( - 'Model parallelism only supports per host input for training. ' - 'Please adjust TPURunconfig.per_host_input_for_training.') - - if eval_batch_size is not None: - util_lib.check_positive_integer(eval_batch_size, 'eval_batch_size') - - if predict_batch_size is not None: - util_lib.check_positive_integer(predict_batch_size, - 'predict_batch_size') - - # Verifies the model_fn signature according to Estimator framework. - estimator_lib._verify_model_fn_args(model_fn, params) # pylint: disable=protected-access - # We cannot store config and params in this constructor as parent - # constructor might change them, such as assigning a temp dir for - # config.model_dir. - model_function = self._augment_model_fn(model_fn, batch_axis) - - # Overwrite log_step_count_steps to disable TensorLoggingHook and - # StepCounterHook from being created in Estimator. TPUEstimator already - # added equivalent hooks in _augment_model_fn above. - self._log_every_n_steps = config.log_step_count_steps - config = config.replace(log_step_count_steps=None) - - # Passing non-None params as wrapped model_fn has it. - params = params or {} - super(TPUEstimator, self).__init__( - model_fn=model_function, - model_dir=model_dir, - config=config, - params=params, - warm_start_from=warm_start_from) - self._iterations_per_training_loop = ( - self._config.tpu_config.iterations_per_loop) - - # All properties passed to _InternalTPUContext are immutable. - # pylint: disable=protected-access - self._ctx = tpu_context._get_tpu_context( - self._config, train_batch_size, eval_batch_size, predict_batch_size, - use_tpu, eval_on_tpu, experimental_embedding_config_spec) - - self._export_to_cpu = export_to_cpu - self._export_to_tpu = export_to_tpu - self._experimental_exported_model_uses_all_cores = ( - experimental_exported_model_uses_all_cores) - self._experimental_export_device_assignment = ( - experimental_export_device_assignment) - if (experimental_exported_model_uses_all_cores and - experimental_export_device_assignment): - raise ValueError('experimental_exported_model_uses_all_cores and ' - 'experimental_export_device_assignment is not supported ' - 'at the same time.') - - self._is_input_fn_invoked = None - self._rendezvous = {} - - def _add_meta_graph_for_mode(self, - builder, - input_receiver_fn_map, - checkpoint_path, - save_variables=True, - mode=model_fn_lib.ModeKeys.PREDICT, - export_tags=None, - check_variables=True): - if self._export_to_tpu and mode != model_fn_lib.ModeKeys.PREDICT: - logging.warning('TPUEstimator only handles mode PREDICT for exporting ' - 'when `export_to_tpu` is `True`; Mode {} will be ignored ' - 'for TPU.'.format(mode)) - - if not self._export_to_cpu and not self._export_to_tpu: - raise ValueError('One of export_to_cpu and export_to_tpu must be true.') - - if self._export_to_cpu: - (super(TPUEstimator, self)._add_meta_graph_for_mode( - builder, - input_receiver_fn_map, - checkpoint_path, - save_variables, - mode=mode, - export_tags=export_tags, - check_variables=check_variables)) - - if self._export_to_tpu and mode == model_fn_lib.ModeKeys.PREDICT: - input_receiver_fn_map = { - _REWRITE_FOR_INFERENCE_MODE: input_receiver_fn_map[mode] - } - export_tags = [tag_constants.SERVING, tag_constants.TPU] - mode = _REWRITE_FOR_INFERENCE_MODE - - # See b/110052256 for why `check_variables` is `False`. - if not self._export_to_cpu: - check_variables = save_variables = True - else: - check_variables = save_variables = False - (super(TPUEstimator, self)._add_meta_graph_for_mode( - builder, - input_receiver_fn_map, - checkpoint_path, - save_variables=save_variables, - mode=mode, - export_tags=export_tags, - check_variables=check_variables)) - - def _call_model_fn(self, features, labels, mode, config): - if mode == _REWRITE_FOR_INFERENCE_MODE: - return self._call_model_fn_for_inference(features, labels, mode, config) - else: - return super(TPUEstimator, self)._call_model_fn(features, labels, mode, - config) - - def _call_model_fn_for_inference(self, features, labels, mode, config): - """Wraps `_call_model_fn` for `export_savedmodel`.""" - if mode != _REWRITE_FOR_INFERENCE_MODE: - raise ValueError('mode must be {}; ' - 'got {}.'.format(_REWRITE_FOR_INFERENCE_MODE, mode)) - - computation, capture = self._build_computation_for_inference( - features, labels, mode, config) - tensors = call_computation( - computation, - experimental_exported_model_uses_all_cores=self - ._experimental_exported_model_uses_all_cores) - estimator_spec, export_outputs_dict, predictions_dict, none_indices = ( - capture.get()) - predictions_list = tensors[:len(predictions_dict)] - export_outputs_list_without_none = tensors[len(predictions_dict):] - - # Reinsert `None`s which we've taken out in - # `_build_computation_for_inference()`. - export_outputs_list = [] - while none_indices or export_outputs_list_without_none: - if none_indices and none_indices[0] == len(export_outputs_list): - export_outputs_list.append(None) - none_indices.pop(0) - else: - export_outputs_list.append(export_outputs_list_without_none.pop(0)) - - # Reconstruct `export_outputs` with updated tensors. - new_export_outputs_dict = nest.pack_sequence_as(export_outputs_dict, - export_outputs_list) - export_outputs = estimator_spec.export_outputs - new_export_outputs = collections.OrderedDict( - (k, _clone_export_output_with_tensors(export_outputs[k], v)) - for k, v in six.iteritems(new_export_outputs_dict)) - # Reconstruct `predictions` with updated tensors. - new_predictions = nest.pack_sequence_as(predictions_dict, predictions_list) - if (len(new_predictions) == 1 and - _KEY_WHEN_PREDICTIONS_IS_A_TENSOR in new_predictions): - new_predictions = new_predictions[_KEY_WHEN_PREDICTIONS_IS_A_TENSOR] - - return estimator_spec._replace( - export_outputs=new_export_outputs, predictions=new_predictions) - - def _build_computation_for_inference(self, features, labels, mode, config): - capture = _CapturedObject() - - def computation(): - """Computation to be passed to `TPUPartitionedCall()`.""" - tpu_computation, tpu_capture = self._build_tpu_computation_for_inference( - features, labels, mode, config) - - if self._experimental_export_device_assignment: - # Export the device assignment as part of the model. This is useful for - # model parallel usecases where the model relies on the mapping between - # logical and physical devices. - with self._ctx.with_mode(mode) as ctx: - device_assignment = ctx.device_assignment - else: - device_assignment = None - - if self._experimental_exported_model_uses_all_cores: - tensors_on_cpu = tpu.rewrite( - tpu_computation, device_assignment=device_assignment) - else: - tensors_on_cpu = tpu.rewrite_for_inference( - tpu_computation, device_assignment=device_assignment) - - (estimator_spec, export_outputs_dict, export_outputs_list, - predictions_dict) = ( - tpu_capture.get()) - predictions_list = tensors_on_cpu[:len(predictions_dict)] - export_outputs_tpu_on_cpu_list = tensors_on_cpu[len(predictions_dict):] - - # Reconstruct tensors used in export_outputs, with TPU tensors replaced - # with their CPU counterpart returned from `rewrite_for_inference()`. - # `function.Defun()` does not like `None`s in return values, so we leave - # `None`s out but record their positions for later reconstruction. - export_outputs_list_without_none = [] - none_indices = [] - for i, t in enumerate(export_outputs_list): - if t is None: - none_indices.append(i) - else: - export_outputs_list_without_none.append( - export_outputs_tpu_on_cpu_list.pop(0)) - - capture.capture((estimator_spec, export_outputs_dict, predictions_dict, - none_indices)) - return predictions_list + export_outputs_list_without_none - - return computation, capture - - def _build_tpu_computation_for_inference(self, features, labels, mode, - config): - capture = _CapturedObject() - - def computation(): - """Compute tpu tensors used in export_outputs. - - Passed to rewrite_for_inference so that model_fn will be called under - the rewriting contexts. Only tpu tensors are returned, but export_outputs - and scaffold are captured. - - Returns: - A list of Tensors used in export_outputs and not marked for - outside_compilation. - """ - # We should only call model fn once and it should be inside `computation` - # so that building the graph will happen under `rewrite_for_inference`. - mode = model_fn_lib.ModeKeys.PREDICT - estimator_spec = self._call_model_fn(features, labels, mode, config) - - # We pick the TPU tensors out from `export_output` and later return them - # from `computation` for rewriting. - export_outputs_dict = collections.OrderedDict( - (k, _export_output_to_tensors(v)) - for k, v in six.iteritems(estimator_spec.export_outputs)) - export_outputs_list = nest.flatten(export_outputs_dict) - export_outputs_tpu_list = [ - t for t in export_outputs_list if t is not None - ] - - if isinstance(estimator_spec.predictions, dict): - predictions_dict = collections.OrderedDict( - (k, v) for k, v in six.iteritems(estimator_spec.predictions)) - else: - predictions_dict = { - _KEY_WHEN_PREDICTIONS_IS_A_TENSOR: estimator_spec.predictions - } - predictions_list = nest.flatten(predictions_dict) - - # We cannot return everything we want through the return values, so - # capture the rest here for later use. - capture.capture((estimator_spec, export_outputs_dict, export_outputs_list, - predictions_dict)) - return predictions_list + export_outputs_tpu_list - - return computation, capture - - def _create_global_step(self, graph): - """Creates a global step suitable for TPUs. - - Args: - graph: The graph in which to create the global step. - - Returns: - A global step `Tensor`. - - Raises: - ValueError: if the global step tensor is already defined. - """ - return _create_global_step(graph) - - def _convert_train_steps_to_hooks(self, steps, max_steps): - with self._ctx.with_mode(model_fn_lib.ModeKeys.TRAIN) as ctx: - if ctx.is_running_on_cpu(): - return super(TPUEstimator, self)._convert_train_steps_to_hooks( - steps, max_steps) - - # On TPU. - if steps is None and max_steps is None: - raise ValueError( - 'For TPU training, one of `steps` or `max_steps` must be set. ' - 'Cannot be both `None`.') - - # Estimator.train has explicit positiveness check. - if steps is not None: - util_lib.check_positive_integer(steps, 'Train steps') - if max_steps is not None: - util_lib.check_positive_integer(max_steps, 'Train max_steps') - - return [ - _TPUStopAtStepHook(self._iterations_per_training_loop, steps, max_steps) - ] - - def _convert_eval_steps_to_hooks(self, steps): - with self._ctx.with_mode(model_fn_lib.ModeKeys.EVAL) as ctx: - if ctx.is_running_on_cpu(): - return super(TPUEstimator, self)._convert_eval_steps_to_hooks(steps) - - if steps is None: - raise ValueError('Evaluate `steps` must be set on TPU. Cannot be `None`.') - - util_lib.check_positive_integer(steps, 'Eval steps') - - return [ - evaluation._StopAfterNEvalsHook( # pylint: disable=protected-access - num_evals=steps), - _SetEvalIterationsHook(steps) - ] - - def _call_input_fn(self, input_fn, mode): - """Calls the input function. - - Args: - input_fn: The input function. - mode: ModeKeys - - Returns: - In TPU mode, returns an input_fn to be called later in model_fn. - Otherwise, calls the input_fn and returns either fatures or - (features, labels). - - Raises: - ValueError: if input_fn takes invalid arguments or does not have `params`. - """ - input_fn_args = function_utils.fn_args(input_fn) - config = self.config # a deep copy. - kwargs = {} - if 'params' in input_fn_args: - kwargs['params'] = self.params # a deep copy. - else: - raise ValueError('input_fn ({}) does not include params argument, ' - 'required by TPUEstimator to pass batch size as ' - 'params["batch_size"]'.format(input_fn)) - if 'config' in input_fn_args: - kwargs['config'] = config - - if 'mode' in input_fn_args: - kwargs['mode'] = mode - - # Records the fact input_fn has been invoked. - self._is_input_fn_invoked = True - - with self._ctx.with_mode(mode) as ctx: - # Setting the batch size in params first. This helps user to have same - # input_fn for use_tpu=True/False. - batch_size_for_input_fn = ctx.batch_size_for_input_fn - if batch_size_for_input_fn is not None: - _add_item_to_params(kwargs['params'], _BATCH_SIZE_KEY, - batch_size_for_input_fn) - - # For export_savedmodel, input_fn is never passed to Estimator. So, - # `is_export_mode` must be False. - if ctx.is_running_on_cpu(is_export_mode=False): - with ops.device('/device:CPU:0'): - return input_fn(**kwargs) - - # For TPU computation, input_fn should be invoked in a tf.while_loop for - # performance. While constructing the tf.while_loop, the structure of - # inputs returned by the `input_fn` needs to be recorded. The structure - # includes whether features or labels is dict or single Tensor, dict keys, - # tensor shapes, and dtypes. The recorded structure is used to create the - # infeed dequeue ops, which must be wrapped and passed as a Fn, called - # inside the TPU computation, as the TPU computation is wrapped inside a - # tf.while_loop also. So, we either pass input_fn to model_fn or pass - # dequeue_fn to model_fn. Here, `input_fn` is passed directly as - # `features` in `model_fn` signature. - def _input_fn(ctx): - _add_item_to_params(kwargs['params'], _CTX_KEY, ctx) - return input_fn(**kwargs) - - return _input_fn - - def _validate_features_in_predict_input(self, result): - """Skip the validation. - - For TPUEstimator, we do not need to check the result type. `_InputPipeline` - has stronger check. Parent class's check generates confusing warning msg. - - Args: - result: `features` returned by input_fn. - """ - pass - - def train(self, - input_fn, - hooks=None, - steps=None, - max_steps=None, - saving_listeners=None): - rendezvous = error_handling.ErrorRendezvous(num_sources=3) - self._rendezvous[model_fn_lib.ModeKeys.TRAIN] = rendezvous - try: - return super(TPUEstimator, self).train( - input_fn=input_fn, - hooks=hooks, - steps=steps, - max_steps=max_steps, - saving_listeners=saving_listeners) - except Exception: # pylint: disable=broad-except - rendezvous.record_error('training_loop', sys.exc_info()) - finally: - rendezvous.record_done('training_loop') - rendezvous.raise_errors() - - def evaluate(self, - input_fn, - steps=None, - hooks=None, - checkpoint_path=None, - name=None): - rendezvous = error_handling.ErrorRendezvous(num_sources=3) - self._rendezvous[model_fn_lib.ModeKeys.EVAL] = rendezvous - try: - return super(TPUEstimator, self).evaluate( - input_fn, - steps=steps, - hooks=hooks, - checkpoint_path=checkpoint_path, - name=name) - except Exception: # pylint: disable=broad-except - rendezvous.record_error('evaluation_loop', sys.exc_info()) - finally: - rendezvous.record_done('evaluation_loop') - rendezvous.raise_errors() - - def predict(self, - input_fn, - predict_keys=None, - hooks=None, - checkpoint_path=None, - yield_single_examples=True): - rendezvous = error_handling.ErrorRendezvous(num_sources=3) - self._rendezvous[model_fn_lib.ModeKeys.PREDICT] = rendezvous - try: - for result in super(TPUEstimator, self).predict( - input_fn=input_fn, - predict_keys=predict_keys, - hooks=hooks, - checkpoint_path=checkpoint_path, - yield_single_examples=yield_single_examples): - yield result - except Exception: # pylint: disable=broad-except - rendezvous.record_error('prediction_loop', sys.exc_info()) - finally: - rendezvous.record_done('prediction_loop') - rendezvous.raise_errors() - - rendezvous.record_done('prediction_loop') - rendezvous.raise_errors() - - def _augment_model_fn(self, model_fn, batch_axis): - """Returns a new model_fn, which wraps the TPU support.""" - - def _model_fn(features, labels, mode, config, params): - """A Estimator `model_fn` for TPUEstimator.""" - with self._ctx.with_mode(mode) as ctx: - model_fn_wrapper = _ModelFnWrapper(model_fn, config, params, ctx) - - # `input_fn` is called in `train()`, `evaluate()`, and `predict()`, - # but not in `export_savedmodel()`. - if self._is_input_fn_invoked: - is_export_mode = False - else: - is_export_mode = True - - # Clear the bit. - self._is_input_fn_invoked = None - - # examples_hook is added to training_hooks for both CPU and TPU - # execution. - if self._log_every_n_steps is not None: - examples_hook = ExamplesPerSecondHook( - ctx.global_batch_size, - # pylint:disable=g-long-ternary - output_dir=(self.model_dir - if not config or config.save_summary_steps - else None), - # pylint:enable=g-long-ternary - every_n_steps=self._log_every_n_steps) - - if ctx.is_running_on_cpu(is_export_mode=is_export_mode): - logging.info('Running %s on CPU', mode) - estimator_spec = model_fn_wrapper.call_without_tpu( - features, labels, is_export_mode=is_export_mode) - if self._log_every_n_steps is not None: - estimator_spec = estimator_spec._replace( - training_hooks=estimator_spec.training_hooks + (examples_hook,)) - return estimator_spec - - assert labels is None, '`labels` passed to `model_fn` must be `None`.' - # TPUEstimator._call_input_fn passes `input_fn` as features to here. - assert callable(features), '`input_fn` is not callable.' - input_fn = features - - tpu_init_ops = [] - if ctx.embedding_config and mode == model_fn_lib.ModeKeys.TRAIN: - dummy_table_variables, dummy_table_variables_init = ( - tpu_embedding_gradient.create_dummy_table_variables( - ctx.embedding_config.tpu_embedding)) - ctx.embedding_config.dummy_table_variables = dummy_table_variables - tpu_init_ops.append(dummy_table_variables_init) - - input_holders = _InputPipeline(input_fn, batch_axis, ctx) - enqueue_ops, dequeue_fn, input_hooks, run_infeed_loop_on_coordinator = ( - input_holders.generate_infeed_enqueue_ops_and_dequeue_fn()) - - graph = ops.get_default_graph() - for enqueue_op in enqueue_ops: - if isinstance(enqueue_op, list): - graph.get_collection_ref(_TPU_ENQUEUE_OPS).extend(enqueue_op) - else: - graph.add_to_collection(_TPU_ENQUEUE_OPS, enqueue_op) - - if mode == model_fn_lib.ModeKeys.TRAIN: - compile_op, loss, host_call, scaffold, training_hooks = ( - _train_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn)) - if ctx.embedding_config: - g = ops.get_default_graph() - table_to_config_dict = ( - ctx.embedding_config.tpu_embedding.table_to_config_dict) - optimization_parameters = ( - ctx.embedding_config.tpu_embedding.optimization_parameters) - embedding_variable_name_by_table, slot_variable_names_by_table = ( - _tpu_estimator_embedding.get_full_variable_names( - g, table_to_config_dict, optimization_parameters - ) - ) - embedding_variables_and_ops = ( - ctx.embedding_config.tpu_embedding.create_variables_and_ops( - embedding_variable_name_by_table, - slot_variable_names_by_table - )) - tpu_init_ops.extend(embedding_variables_and_ops.load_ops()) - - host_ops = host_call.create_tpu_hostcall() - if host_ops is None: - host_ops = [] - - shutdown_hooks = [] - shutdown_mode = os.environ.get('TF_TPU_GRACEFUL_SHUTDOWN_MODE', - 'shutdown_worker') - if shutdown_mode: - if shutdown_mode == 'shutdown_worker': - finalizer_hooks = [ - session_support.ShutdownLameWorkers(timeout_ms=60 * 1000), - ] - elif shutdown_mode == 'shutdown_computation': - finalizer_hooks = [ - session_support.RestartComputation(timeout_ms=60 * 1000), - ] - else: - raise ValueError( - 'Unknown TF_TPU_GRACEFUL_SHUTDOWN_MODE "%s"' % shutdown_mode) - - shutdown_hooks.append( - session_support.GracefulShutdownHook( - checkpoint_prefix=self.model_dir + '/model.ckpt', - on_shutdown_hooks=finalizer_hooks)) - - with ops.control_dependencies([loss]): - global_step = array_ops.identity(training.get_global_step()) - hooks = input_hooks + shutdown_hooks - hooks.extend([ - TPUInfeedOutfeedSessionHook( - ctx, - enqueue_ops, - host_ops, - tpu_compile_op=compile_op, - run_infeed_loop_on_coordinator=( - run_infeed_loop_on_coordinator), - rendezvous=self._rendezvous[mode], - master=self._config.master, - session_config=self._session_config, - tpu_init_ops=tpu_init_ops), - InstallSignalHandlerHook() - ]) - if self._log_every_n_steps is not None: - logging_hook_frequency = ( # Divide and round up - (self._log_every_n_steps + - self._config.tpu_config.iterations_per_loop - 1) // - self._config.tpu_config.iterations_per_loop) - hooks.append( - training.LoggingTensorHook({ - 'loss': array_ops.identity(loss), - 'step': global_step, - }, - every_n_iter=logging_hook_frequency)) - examples_hook._set_steps_per_run( # pylint: disable=protected-access - self._config.tpu_config.iterations_per_loop) - hooks.append(examples_hook) - - if training_hooks: - hooks.extend(training_hooks) - - chief_hooks = [] - if (self._config.save_checkpoints_secs or - self._config.save_checkpoints_steps): - checkpoint_hook = training.CheckpointSaverHook( - self.model_dir, - save_secs=self._config.save_checkpoints_secs, - save_steps=self._config.save_checkpoints_steps, - scaffold=scaffold) - checkpoint_hook._set_steps_per_run( # pylint: disable=protected-access - self._config.tpu_config.iterations_per_loop) - chief_hooks.append(checkpoint_hook) - - summary.scalar(model_fn_lib.LOSS_METRIC_KEY, loss) - with ops.control_dependencies([loss]): - update_ops = _sync_variables_ops(ctx) - if ctx.embedding_config: - update_ops.extend(embedding_variables_and_ops.retrieve_ops()) - - # Validate the TPU training graph to catch basic errors - _validate_tpu_training_graph() - - train_op = control_flow_ops.group(*update_ops) - graph.add_to_collection(_TPU_TRAIN_OP, train_op) - - return model_fn_lib.EstimatorSpec( - mode, - loss=loss, - training_chief_hooks=chief_hooks, - training_hooks=hooks, - train_op=train_op, - scaffold=scaffold) - - if mode == model_fn_lib.ModeKeys.EVAL: - compile_op, total_loss, host_calls, scaffold, eval_hooks = ( - _eval_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn)) - iterations_per_loop_var = _create_or_get_iterations_per_loop() - mean_loss = math_ops.div( - total_loss, - math_ops.cast(iterations_per_loop_var, dtype=total_loss.dtype)) - - with ops.control_dependencies([mean_loss]): - # After TPU evaluation computation is done (the mean_loss tensor), - # reads all variables back from TPU and updates the eval step - # counter properly - internal_ops_to_run = _sync_variables_ops(ctx) - internal_ops_to_run.append( - _increase_eval_step_op(iterations_per_loop_var)) - - host_call_ret = host_calls.create_tpu_hostcall() - eval_metric_ops = {} - eval_update_ops = [] - - eval_metrics = host_call_ret.get('eval_metrics', {}) - if eval_metrics: - # Creates a dummy metric update_op for all metrics. Estimator - # expects all metrics in `eval_metric_ops` have update_op and calls - # them one by one. The real metric update_ops are invoked in a - # separated thread. So, here give Estimator the dummy op for all - # metrics. - with ops.control_dependencies(internal_ops_to_run): - dummy_update_op = control_flow_ops.no_op() - - for k, v in eval_metrics.items(): - eval_metric_ops[k] = (v[0], dummy_update_op) - eval_update_ops.append(v[1]) - else: - # If no eval metrics are passed, create an identity node for the - # loss and add `internal_ops_to_run` to its dependencies. So - # `internal_ops_to_run` can be executed. - with ops.control_dependencies(internal_ops_to_run): - mean_loss = array_ops.identity(mean_loss) - - if 'host_call' not in host_call_ret: - host_ops = [] - else: - host_ops = host_call_ret['host_call'] - hooks = [ - TPUInfeedOutfeedSessionHook( - ctx, - enqueue_ops, - eval_update_ops + host_ops, - tpu_compile_op=compile_op, - run_infeed_loop_on_coordinator=( - run_infeed_loop_on_coordinator), - rendezvous=self._rendezvous[mode], - master=self._config.evaluation_master, - session_config=self._session_config, - tpu_init_ops=tpu_init_ops) - ] + input_hooks - - if eval_hooks: - hooks.extend(eval_hooks) - - return model_fn_lib.EstimatorSpec( - mode, - loss=mean_loss, - evaluation_hooks=hooks, - eval_metric_ops=eval_metric_ops, - scaffold=scaffold) - - # Predict - assert mode == model_fn_lib.ModeKeys.PREDICT - - (compile_op, dummy_predict_op, host_calls, - scaffold, prediction_hooks) = _predict_on_tpu_system( - ctx, model_fn_wrapper, dequeue_fn) - with ops.control_dependencies([dummy_predict_op]): - internal_ops_to_run = _sync_variables_ops(ctx) - with ops.control_dependencies(internal_ops_to_run): - dummy_predict_op = control_flow_ops.no_op() - - # In train and evaluation, the main TPU program is passed to monitored - # training session to run. Infeed enqueue and outfeed dequeue are - # executed in side threads. This is not the configuration for - # prediction mode. - # - # For prediction, the Estimator executes the EstimatorSpec.predictions - # directly and yield the element (via generator) to call site. So, the - # outfeed based prediction must be passed to MonitoredSession directly. - # Other parts of the TPU execution are organized as follows. - # - # 1. All outfeed based Tensors must be grouped with predictions Tensors - # to form a single invocation. This avoid the issue we might trigger - # multiple outfeeds incorrectly. To achieve this, `host_call` is - # placed in control_dependencies of `stopping_signals`, and - # `stopping_signals` is passed into _StoppingPredictHook, which sets - # the `stopping_signals` as SessionRunArgs. MonitoredSession merges - # all SessionRunArgs with the fetch in session.run together. - # - # 2. The TPU program (dummy_predict_op) and enqueue_ops (infeed Enqueue) - # are grouped together. They will be launched once and only once in - # side threads and they quit naturally according to the SAME stopping - # condition. - enqueue_ops.append(dummy_predict_op) - - host_call_ret = host_calls.create_tpu_hostcall() - if 'host_call' not in host_call_ret: - host_ops = [] - else: - host_ops = host_call_ret['host_call'] - - predictions = host_call_ret['predictions'] - _verify_cross_hosts_transfer_size( - predictions, - message=( - 'The estimated size for TPUEstimatorSpec.predictions is too ' - 'large.')) - signals = host_call_ret['signals'] - - with ops.control_dependencies(host_ops): - host_ops = [] # Empty, we do do not need it anymore. - scalar_stopping_signal = _StopSignals.as_scalar_stopping_signal( - signals) - predictions = _PaddingSignals.slice_tensor_or_dict( - predictions, signals) - - hooks = [ - _StoppingPredictHook(scalar_stopping_signal), - TPUInfeedOutfeedSessionHookForPrediction( - ctx, enqueue_ops, host_ops, rendezvous=self._rendezvous[mode], - tpu_compile_op=compile_op, - master=self._config.master, - session_config=self._session_config), - ] + input_hooks - - if prediction_hooks: - hooks.extend(prediction_hooks) - - return model_fn_lib.EstimatorSpec( - mode, - prediction_hooks=hooks, - predictions=predictions, - scaffold=scaffold) - - return _model_fn - - -def _export_output_to_tensors(export_output): - """Get a list of `Tensors` used in `export_output`. - - Args: - export_output: an `ExportOutput` object such as `ClassificationOutput`, - `RegressionOutput`, or `PredictOutput`. - - Returns: - a list of tensors used in export_output. - - Raises: - ValueError: if `export_output` is not one of `ClassificationOutput`, - `RegressionOutput`, or `PredictOutput`. - """ - if isinstance(export_output, export_output_lib.ClassificationOutput): - return [export_output.scores, export_output.classes] - elif isinstance(export_output, export_output_lib.RegressionOutput): - return [export_output.value] - elif isinstance(export_output, export_output_lib.PredictOutput): - return list(export_output.outputs.values()) - else: - raise ValueError( - '`export_output` must be have type `ClassificationOutput`, ' - '`RegressionOutput`, or `PredictOutput`; got {}.'.format(export_output)) - - -def _clone_export_output_with_tensors(export_output, tensors): - """Clones `export_output` but with new `tensors`. - - Args: - export_output: an `ExportOutput` object such as `ClassificationOutput`, - `RegressionOutput`, or `PredictOutput`. - tensors: a list of `Tensors` used to construct a new `export_output`. - - Returns: - A dict similar to `export_output` but with `tensors`. - - Raises: - ValueError: if `export_output` is not one of `ClassificationOutput`, - `RegressionOutput`, or `PredictOutput`. - """ - if isinstance(export_output, export_output_lib.ClassificationOutput): - if len(tensors) != 2: - raise ValueError('tensors must be of length 2; ' - 'got {}.'.format(len(tensors))) - return export_output_lib.ClassificationOutput(*tensors) - elif isinstance(export_output, export_output_lib.RegressionOutput): - if len(tensors) != 1: - raise ValueError('tensors must be of length 1; ' - 'got {}'.format(len(tensors))) - return export_output_lib.RegressionOutput(*tensors) - elif isinstance(export_output, export_output_lib.PredictOutput): - return export_output_lib.PredictOutput( - dict(zip(export_output.outputs.keys(), tensors))) - else: - raise ValueError( - '`export_output` must be have type `ClassificationOutput`, ' - '`RegressionOutput`, or `PredictOutput`; got {}.'.format(export_output)) - - -def _eval_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): - """Executes `model_fn_wrapper` multiple times on all TPU shards.""" - iterations_per_loop_var = _create_or_get_iterations_per_loop() - - (single_tpu_eval_step, host_calls, captured_scaffold_fn, captured_eval_hooks - ) = model_fn_wrapper.convert_to_single_tpu_eval_step(dequeue_fn) - - def multi_tpu_eval_steps_on_single_shard(): - return training_loop.repeat(iterations_per_loop_var, single_tpu_eval_step, - [_ZERO_LOSS]) - - (compile_op, loss,) = tpu.split_compile_and_shard( - multi_tpu_eval_steps_on_single_shard, - inputs=[], - num_shards=ctx.num_replicas, - outputs_from_all_shards=False, - device_assignment=ctx.device_assignment) - - loss = loss[0] - scaffold = _get_scaffold(captured_scaffold_fn) - return compile_op, loss, host_calls, scaffold, captured_eval_hooks.get() - - -def _train_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): - """Executes `model_fn_wrapper` multiple times on all TPU shards.""" - iterations_per_loop_var = _create_or_get_iterations_per_loop() - - (single_tpu_train_step, host_call, captured_scaffold_fn, - captured_training_hooks) = ( - model_fn_wrapper.convert_to_single_tpu_train_step(dequeue_fn)) - - @tpu_function.on_device_training_loop - def multi_tpu_train_steps_on_single_shard(): - return training_loop.repeat(iterations_per_loop_var, single_tpu_train_step, - [_INITIAL_LOSS]) - - (compile_op, loss,) = tpu.split_compile_and_shard( - multi_tpu_train_steps_on_single_shard, - inputs=[], - num_shards=ctx.num_replicas, - outputs_from_all_shards=False, - device_assignment=ctx.device_assignment) - - loss = loss[0] - scaffold = _get_scaffold(captured_scaffold_fn) - return compile_op, loss, host_call, scaffold, captured_training_hooks.get() - - -def _predict_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): - """Executes `model_fn_wrapper` multiple times on all TPU shards.""" - (single_tpu_predict_step, host_calls, captured_scaffold_fn, - captured_predict_hooks - ) = model_fn_wrapper.convert_to_single_tpu_predict_step(dequeue_fn) - - @tpu_function.on_device_training_loop - def multi_tpu_predict_steps_on_single_shard(): - - def cond(scalar_stopping_signal): - return math_ops.logical_not( - _StopSignals.should_stop(scalar_stopping_signal)) - - inputs = [_StopSignals.NON_STOPPING_SIGNAL] - outputs = training_loop.while_loop( - cond, single_tpu_predict_step, inputs=inputs, name=b'loop') - return outputs - - (compile_op, dummy_predict_op,) = tpu.split_compile_and_shard( - multi_tpu_predict_steps_on_single_shard, - inputs=[], - num_shards=ctx.num_replicas, - outputs_from_all_shards=False, - device_assignment=ctx.device_assignment) - - dummy_predict_op = dummy_predict_op[0] - scaffold = _get_scaffold(captured_scaffold_fn) - return (compile_op, dummy_predict_op, host_calls, scaffold, - captured_predict_hooks.get()) - - -def _wrap_computation_in_while_loop(device, op_fn): - """Wraps the ops generated by `op_fn` in tf.while_loop.""" - - def computation(i): - with ops.control_dependencies(op_fn()): - return i + 1 - - iterations_per_loop_var = _create_or_get_iterations_per_loop() - # By setting parallel_iterations=1, the parallel execution in while_loop is - # basically turned off. - with ops.device(device): - iterations = array_ops.identity(iterations_per_loop_var) - return control_flow_ops.while_loop( - lambda i: i < iterations, - computation, [constant_op.constant(0)], - parallel_iterations=1) - - -def _wrap_computation_in_while_loop_with_stopping_signals(device, op_fn): - """Wraps the ops generated by `op_fn` in tf.while_loop.""" - - def cond(scalar_stopping_signal): - return math_ops.logical_not( - _StopSignals.should_stop(scalar_stopping_signal)) - - def computation(unused_scalar_stopping_signal): - return_value = op_fn() - execute_ops = return_value['ops'] - signals = return_value['signals'] - with ops.control_dependencies(execute_ops): - return _StopSignals.as_scalar_stopping_signal(signals) - - # By setting parallel_iterations=1, the parallel execution in while_loop is - # basically turned off. - with ops.device(device): - return control_flow_ops.while_loop( - cond, - computation, [_StopSignals.NON_STOPPING_SIGNAL], - parallel_iterations=1) - - -def _validate_tpu_training_graph(): - """Validate graph before running distributed training. - - Raises: - ValueError: If the graph seems invalid for running on device - """ - operations = ops.get_default_graph().get_operations() - - # Check if there is atleast one CrossReplicaSum operation in the graph - # This should be introduced by using the CrossShardOptimizer wrapper - cross_replica_sum_ops = [ - o for o in operations if o.type == _CROSS_REPLICA_SUM_OP - ] - if not cross_replica_sum_ops: - raise ValueError( - 'CrossShardOptimizer must be used for model training on TPUs.') - - -class _CapturedObject(object): - """A placeholder to capture an object. - - This is useful when we need to capture a Python object in the Tensorflow - control flow body function and use it outside the control flow. - """ - - def __init__(self): - self._object = None - self._captured = False - - def capture(self, o): - if self._captured: - raise RuntimeError( - 'InternalError: Object can capture only once. Please file bug.') - - self._captured = True - self._object = o - - def get(self): - if not self._captured: - raise RuntimeError( - 'InternalError: Object is not captured properly before `get`. ' - 'Please file bug.') - return self._object - - -def _get_scaffold(captured_scaffold_fn): - """Retrieves the Scaffold from `captured_scaffold_fn`.""" - with _CapturingContext(message='Inside scaffold_fn'): - scaffold_fn = captured_scaffold_fn.get() - if scaffold_fn: - scaffold = scaffold_fn() - if scaffold is None: - raise ValueError( - 'TPUEstimatorSpec.scaffold_fn returns None, which is not allowed') - else: - scaffold = None - - if scaffold: - wrapped_finalize = scaffold.finalize - - def _finalize(): - with _CapturingContext('Inside Scaffold.finalize'): - wrapped_finalize() - - scaffold.finalize = _finalize - return scaffold - - -class _CapturingContext(control_flow_ops.ControlFlowContext): - """Tracks references to Tensors defined in TPU replication.""" - - def __init__(self, message): - control_flow_ops.ControlFlowContext.__init__(self) - self._message = message - - def to_control_flow_context_def(self, context_def, export_scope=None): - # pylint: disable=useless-super-delegation - # NOTE(slebedev): the method is required by `ControlFlowContext`. - super(_CapturingContext, self).to_control_flow_context_def( - context_def, export_scope) - - def AddOp(self, op): # pylint: disable=invalid-name - for c in op.inputs: - if tpu._TPU_REPLICATE_ATTR in c.op.node_def.attr: # pylint: disable=protected-access - raise ValueError('{}: Op {} depends on TPU computation {}, ' - 'which is not allowed.'.format(self._message, op, c)) - - def __enter__(self): - # pylint: disable=protected-access - self._g = ops.get_default_graph() - self._old = self._g._get_control_flow_context() - self._g._set_control_flow_context(self) - # pylint: enable=protected-access - - def __exit__(self, _, __, ___): # pylint: disable=invalid-name - self._g._set_control_flow_context(self._old) # pylint: disable=protected-access - - -class _Inputs(object): - """A data structure representing the input_fn returned values. - - This also supports the returned value from input_fn as `Dataset`. - """ - - def __init__(self, features=None, labels=None, dataset=None, signals=None): - if dataset is not None and (features is not None or labels is not None or - signals is not None): - raise RuntimeError('Internal Error: Either (features and labels) or ' - 'dataset should be provided, not both. Please file ' - 'bug') - - self._features = features - self._labels = labels - self._signals = signals - - self._dataset = dataset - self._iterator = None - - @staticmethod - def from_input_fn(return_values): - """Returns an `_Inputs` instance according to `input_fn` return value.""" - if isinstance(return_values, dataset_ops.DatasetV2): - dataset = return_values - return _Inputs(dataset=dataset) - - features, labels = _Inputs._parse_inputs(return_values) - return _Inputs(features, labels) - - @staticmethod - def _parse_inputs(return_values): - if isinstance(return_values, tuple): - features, labels = return_values - else: - features, labels = return_values, None - return features, labels - - @property - def is_dataset(self): - """Returns True if the return value from input_fn is Dataset.""" - return self._dataset is not None - - def dataset_initializer(self): - """Returns the dataset's initializer. - - The initializer must be run before calling `features_and_labels`. - """ - self._iterator = dataset_ops.make_initializable_iterator(self._dataset) - return self._iterator.initializer - - def features_and_labels(self): - """Gets `features` and `labels`.""" - if self.is_dataset: - if self._iterator is None: - raise RuntimeError('Internal error: Must run dataset_initializer ' - 'before calling features_and_labels(). Please file ' - 'a bug!') - return _Inputs._parse_inputs(self._iterator.get_next()) - - return (self._features, self._labels) - - def signals(self): - return self._signals - - @property - def dataset(self): - return self._dataset - - -class _InputsWithStoppingSignals(_Inputs): - """Inputs with `_StopSignals` inserted into the dataset.""" - - def __init__(self, - dataset, - batch_size, - add_padding=False, - num_invocations_per_step=1): - - assert dataset is not None - user_provided_dataset = dataset.map( - _InputsWithStoppingSignals.insert_stopping_signal( - stop=False, batch_size=batch_size, add_padding=add_padding)) - if num_invocations_per_step == 1: - final_batch_dataset = dataset.take(1).map( - _InputsWithStoppingSignals.insert_stopping_signal( - stop=True, batch_size=batch_size, add_padding=add_padding)) - else: - # We append (2 * num_invocations_per_step - 1) batches for exhausting the - # user_provided_dataset and stop properly. - # For example, if num_invocations_per_step is 2, we append 3 additional - # padding batches: b1, b2, b3. - # If user_provided_dataset contains two batches: a1, a2 - # Step 1: [a1, a2] - # Step 2: [b1, b2] -> STOP - # If user_provided_dataset contains three batches: a1, a2, a3. - # The training loops: - # Step 1: [a1, a2] - # Step 2: [a3, b1] - # Step 3: [b2, b3] -> STOP. - final_batch_dataset = dataset.take(1).map( - _InputsWithStoppingSignals.insert_stopping_signal( - stop=True, batch_size=batch_size, add_padding=add_padding)) - final_batch_dataset = final_batch_dataset.repeat( - 2 * num_invocations_per_step - 1) - - def _set_mask(data_dict): - signals = data_dict['signals'] - signals['padding_mask'] = array_ops.ones_like(signals['padding_mask']) - data_dict['signals'] = signals - return data_dict - - # Mask out the extra batch. - final_batch_dataset = final_batch_dataset.map(_set_mask) - - dataset = user_provided_dataset.concatenate(final_batch_dataset).prefetch(2) - - super(_InputsWithStoppingSignals, self).__init__(dataset=dataset) - self._current_inputs = None - - def features_and_labels(self): - if self._current_inputs is not None: - raise RuntimeError( - 'Internal Error: The previous inputs have not been properly ' - 'consumed. First call features_and_labels, then call signals.') - - inputs_with_signals = self._iterator.get_next() - features = inputs_with_signals['features'] - labels = inputs_with_signals.get('labels') - - self._current_inputs = inputs_with_signals - return features, labels - - def signals(self): - """Returns the `Signals` from `_Inputs`.""" - if self._current_inputs is None: - raise RuntimeError( - 'Internal Error: The current inputs have not been properly ' - 'generated. First call features_and_labels, then call signals.') - signals = self._current_inputs['signals'] - self._current_inputs = None - return signals - - @staticmethod - def insert_stopping_signal(stop, batch_size, add_padding=False): - """Inserts stopping_signal into dataset via _map_fn. - - Here we change the data structure in the dataset, such that the return value - is a dictionary now and `features`, `labels`, and `signals` are three - distinguished keys in that dict. This provides a better structure, which - eases the process to decompose the inputs (see `features_and_labels`). - - Args: - stop: bool, state of current stopping signals. - batch_size: int, batch size. - add_padding: bool, whether to pad the tensor to full batch size. - - Returns: - A map_fn passed to dataset.map API. - """ - - def _map_fn(*args): - """The map fn to insert signals.""" - if len(args) == 1: - # Unpack the single Tensor/dict argument as features. This is required - # for the input_fn returns no labels. - args = args[0] - features, labels = _Inputs._parse_inputs(args) - new_input_dict = {} - - if add_padding: - padding_mask, features, labels = ( - _PaddingSignals.pad_features_and_labels(features, labels, - batch_size)) - - new_input_dict['features'] = features - if labels is not None: - new_input_dict['labels'] = labels - - else: - new_input_dict['features'] = features - if labels is not None: - new_input_dict['labels'] = labels - padding_mask = None - - new_input_dict['signals'] = _StopSignals( - stop=stop, batch_size=batch_size, - padding_mask=padding_mask).as_dict() - - return new_input_dict - - return _map_fn - - -class _StopSignals(object): - """Signals class holding all logic to handle TPU stopping condition.""" - - NON_STOPPING_SIGNAL = False - STOPPING_SIGNAL = True - - def __init__(self, stop, batch_size, padding_mask=None): - self._stop = stop - self._batch_size = batch_size - self._padding_mask = padding_mask - - def as_dict(self): - """Returns the signals as Python dict.""" - shape = [self._batch_size, 1] - dtype = dtypes.bool - - if self._stop: - stopping = array_ops.ones(shape=shape, dtype=dtype) - else: - stopping = array_ops.zeros(shape=shape, dtype=dtype) - - signals = {'stopping': stopping} - if self._padding_mask is not None: - signals['padding_mask'] = self._padding_mask - return signals - - @staticmethod - def as_scalar_stopping_signal(signals): - return array_ops.identity(signals['stopping'][0][0]) - - @staticmethod - def should_stop(scalar_stopping_signal): - """Detects whether scalar_stopping_signal indicates stopping.""" - if isinstance(scalar_stopping_signal, ops.Tensor): - # STOPPING_SIGNAL is a constant True. Here, the logical_and is just the TF - # way to express the bool check whether scalar_stopping_signal is True. - return math_ops.logical_and(scalar_stopping_signal, - _StopSignals.STOPPING_SIGNAL) - else: - # For non Tensor case, it is used in SessionRunHook. So, we cannot modify - # the graph anymore. Here, we use pure Python. - return bool(scalar_stopping_signal) - - -class _PaddingSignals(object): - """Signals class holding all logic to handle padding.""" - - @staticmethod - def pad_features_and_labels(features, labels, batch_size): - """Pads out the batch dimension of features and labels.""" - real_batch_size = array_ops.shape( - _PaddingSignals._find_any_tensor(features))[0] - - batch_size_tensor = constant_op.constant(batch_size, dtypes.int32) - - check_greater = check_ops.assert_greater_equal( - batch_size_tensor, - real_batch_size, - data=(batch_size_tensor, real_batch_size), - message='The real batch size should not be greater than batch_size.') - - with ops.control_dependencies([check_greater]): - missing_count = batch_size_tensor - real_batch_size - - def pad_single_tensor(tensor): - """Pads out the batch dimension of a tensor to the complete batch_size.""" - rank = len(tensor.shape) - assert rank > 0 - padding = array_ops.stack([[0, missing_count]] + [[0, 0]] * (rank - 1)) - padded_shape = (batch_size,) + tuple(tensor.shape[1:]) - padded_tensor = array_ops.pad(tensor, padding) - padded_tensor.set_shape(padded_shape) - return padded_tensor - - def nest_pad(tensor_or_dict): - return nest.map_structure(pad_single_tensor, tensor_or_dict) - - features = nest_pad(features) - if labels is not None: - labels = nest_pad(labels) - - padding_mask = _PaddingSignals._padding_mask(real_batch_size, missing_count, - batch_size) - - return padding_mask, features, labels - - @staticmethod - def slice_tensor_or_dict(tensor_or_dict, signals): - """Slice the real Tensors according to padding mask in signals.""" - - padding_mask = signals['padding_mask'] - batch_size = array_ops.shape(padding_mask)[0] - - def verify_batch_size(tensor): - check_batch_size = math_ops.equal(batch_size, tensor.shape[0]) - with ops.control_dependencies([check_batch_size]): - return array_ops.identity(tensor) - - def slice_single_tensor(tensor): - rank = len(tensor.shape) - assert rank > 0 - real_batch_size = batch_size - math_ops.reduce_sum(padding_mask) - return verify_batch_size(tensor)[0:real_batch_size] - - # As we split the Tensors to all TPU cores and concat them back, it is - # important to ensure the real data is placed before padded ones, i.e., - # order is preserved. By that, the sliced padding mask should have all 0's. - # If this assertion failed, # the slice logic here would not hold. - sliced_padding_mask = slice_single_tensor(padding_mask) - assert_padding_mask = math_ops.equal( - math_ops.reduce_sum(sliced_padding_mask), 0) - - with ops.control_dependencies([assert_padding_mask]): - should_stop = _StopSignals.should_stop( - _StopSignals.as_scalar_stopping_signal(signals)) - - is_full_batch = math_ops.equal(math_ops.reduce_sum(padding_mask), 0) - - def slice_fn(tensor): - # If the current batch is full batch or part of stopping signals, we do - # not need to slice to save performance. - return control_flow_ops.cond( - math_ops.logical_or(should_stop, is_full_batch), - (lambda: verify_batch_size(tensor)), - (lambda: slice_single_tensor(tensor))) - - return nest.map_structure(slice_fn, tensor_or_dict) - - @staticmethod - def _find_any_tensor(batch_features): - tensors = [ - x for x in nest.flatten(batch_features) if isinstance(x, ops.Tensor) - ] - if not tensors: - raise ValueError('Cannot find any Tensor in features dict.') - return tensors[0] - - @staticmethod - def _padding_mask(real_batch_size, missing_count, batch_size): - padding_mask = array_ops.concat([ - array_ops.zeros((real_batch_size,), dtype=dtypes.int32), - array_ops.ones((missing_count,), dtype=dtypes.int32) - ], - axis=0) - padding_mask.set_shape((batch_size,)) - return padding_mask - - -def _verify_cross_hosts_transfer_size(tensor_dict, message): - total_size = 0 - tensor_structure = {} - for key, tensor in tensor_dict.items(): - shape = tensor.shape - size = np.product(shape) * tensor.dtype.size - tensor_structure[key] = shape - total_size += size - if total_size >= _ONE_GIGABYTE: - raise ValueError( - '{} The transfer size is larger than the protobuf limit. Please ' - 'consider to use Tensors with smaller shapes or reduce batch ' - 'size. Given:\n' - '{}'.format( - message, '\n'.join([ - ' -- Key: {}, Shape: {}'.format(k, v) - for k, v in tensor_structure.items() - ]))) - - -def _add_item_to_params(params, key, value): - """Adds a new item into `params`.""" - if isinstance(params, hparam.HParams): - # For HParams, we need to use special API. - if key in params: - params.set_hparam(key, value) - else: - params.add_hparam(key, value) - else: - # Now params is Python dict. - params[key] = value - - -def export_estimator_savedmodel(estimator, - export_dir_base, - serving_input_receiver_fn, - assets_extra=None, - as_text=False, - checkpoint_path=None, - strip_default_attrs=False): - """Export `Estimator` trained model for TPU inference. - - Args: - estimator: `Estimator` with which model has been trained. - export_dir_base: A string containing a directory in which to create - timestamped subdirectories containing exported SavedModels. - serving_input_receiver_fn: A function that takes no argument and returns a - `ServingInputReceiver` or `TensorServingInputReceiver`. - assets_extra: A dict specifying how to populate the assets.extra directory - within the exported SavedModel, or `None` if no extra assets are needed. - as_text: whether to write the SavedModel proto in text format. - checkpoint_path: The checkpoint path to export. If `None` (the default), - the most recent checkpoint found within the model directory is chosen. - strip_default_attrs: Boolean. If `True`, default-valued attributes will be - removed from the NodeDefs. - - Returns: - The string path to the exported directory. - """ - # `TPUEstimator` requires `tpu_config.RunConfig`, so we cannot use - # `estimator.config`. - config = tpu_config.RunConfig(model_dir=estimator.model_dir) - est = TPUEstimator( - estimator._model_fn, # pylint: disable=protected-access - config=config, - params=estimator.params, - use_tpu=True, - train_batch_size=2048, # Does not matter. - eval_batch_size=2048, # Does not matter. - ) - return est.export_savedmodel(export_dir_base, serving_input_receiver_fn, - assets_extra, as_text, checkpoint_path, - strip_default_attrs) +# pylint: disable=wildcard-import,unused-import,redefined-builtin +from tensorflow.python.tpu.tpu_estimator import * +# used by tests +from tensorflow.python.tpu.tpu_estimator import _clone_export_output_with_tensors +from tensorflow.python.tpu.tpu_estimator import _create_global_step +from tensorflow.python.tpu.tpu_estimator import _export_output_to_tensors +from tensorflow.python.tpu.tpu_estimator import _get_scaffold +from tensorflow.python.tpu.tpu_estimator import _Inputs +from tensorflow.python.tpu.tpu_estimator import _ITERATIONS_PER_LOOP_VAR +from tensorflow.python.tpu.tpu_estimator import _TPU_ENQUEUE_OPS +from tensorflow.python.tpu.tpu_estimator import _TPU_ESTIMATOR +from tensorflow.python.tpu.tpu_estimator import _TPU_TRAIN_OP +# pylint: enable=wildcard-import,unused-import,redefined-builtin diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_feed.py b/tensorflow/contrib/tpu/python/tpu/tpu_feed.py index 97fddbc2adb688b3e5ec8c3f39adcebd8db6cbc7..af2542ea85290170ce6a38223188c4f9b871f032 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_feed.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_feed.py @@ -1,919 +1,25 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# =================================================================== - -"""Helper library for handling infeed between hosts and TPUs. -""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import itertools - -import numpy as np -from six.moves import xrange # pylint: disable=redefined-builtin - -from tensorflow.compiler.xla.experimental.xla_sharding import xla_sharding -from tensorflow.contrib.tpu.python.ops import tpu_ops -from tensorflow.contrib.tpu.python.tpu import tpu -from tensorflow.contrib.tpu.python.tpu import tpu_sharding - -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_shape -from tensorflow.python.ops import array_ops -from tensorflow.python.util import nest - - -def partition_or_replicate_on_host(tensor, dims): - """Partitions or replicates the input tensor. - - The ops inside this function are placed on the host side. - - Args: - tensor: The input tensor which will be partioned or replicated. - dims: A list of integer describes how to partition the input tensor. - - Returns: - An iterator of `Tensor`s or a list of partioned tensors. - """ - if dims is None: - return itertools.repeat(tensor) - dims = np.array(dims) - output = [tensor] - shape_list = np.array(tensor.shape.as_list()) - quotients, remainders = np.divmod(shape_list, dims) - for axis, (quotient, remainder, dim, original_size) in enumerate( - zip(quotients, remainders, dims, shape_list)): - if dim <= 1: - continue - if remainder > 0: - # For each dimension, when it cannot be evenly partitioned, XLA assumes - # tensors are partitioned in a greedy manner by using - # ceil_ratio(size/dim) first. E.g. 2D tensor with shape (5, 14) and dims - # are (2, 4). Since 5 % 2 = 1 and 14 % 4 = 2, [5, 14] => - # [[(3, 4), (3, 4), (2, 4), (2, 2)], - # [(2, 4), (2, 4), (2, 4), (2, 2)]] - ceil_ratio = quotient + 1 - num_full_slots, left_over = np.divmod(original_size, ceil_ratio) - num_or_size_splits = [ceil_ratio] * num_full_slots + [left_over] - if len(num_or_size_splits) < dim: - num_or_size_splits += [0] * (dim - len(num_or_size_splits)) - new_output = [] - for x in output: - new_output.append( - array_ops.split( - x, num_or_size_splits=num_or_size_splits, axis=axis)) - output = new_output - else: - output = [array_ops.split(x, dim, axis=axis) for x in output] - output = nest.flatten(output) - return output - - -def _tag_sharding_attribute_for_dequeued_tensor(tensor, dims): - """Tags appropriate XLA sharding attribute to the dequeued tensor. - - Args: - tensor: The dequeued tensor on TPU. - dims: A list of integer describes how the tensor is partitioned. - - Returns: - The same tensor with the xla_sharding attribute. - """ - if dims is None: - return xla_sharding.replicate(tensor) - elif np.prod(dims) == 1: - return xla_sharding.assign_device(tensor, 0) - else: - tile_assignment = np.arange(np.prod(dims)).reshape(dims) - return xla_sharding.tile(tensor=tensor, tile_assignment=tile_assignment) - - -def tag_sharding_attribute_for_dequeued_tensors(dequeues, dims): - """Tags appropriate XLA sharding attribute to the dequeued tensors. - - Args: - dequeues: A list of dequeued tensors on TPU. - dims: A list of integer describes how the tensor is partitioned. - - Returns: - The same dequeues with appropriate xla_sharding attribute. - """ - nest.assert_shallow_structure(dequeues, dims) - return nest.map_structure_up_to( - dequeues, _tag_sharding_attribute_for_dequeued_tensor, dequeues, dims) - - -class InfeedQueue(object): - """A helper object to build a device infeed queue. - - The InfeedQueue builds the host-side and device-side Ops to enqueue and - dequeue elements, respectively, and ensures that their types and - shapes match. - """ - - def __init__(self, - number_of_tuple_elements=None, - tuple_types=None, - tuple_shapes=None, - shard_dimensions=None, - name=None): - """Creates a new InfeedQueue with the given configuration. - - The configuration need not be fully specified at creation since it - can be modified subsequently by methods that set the values - explicitly or infer them from the shapes of inputs. - - Args: - number_of_tuple_elements: the number of Tensors fed atomically through the - queue, must be present unless it can be inferred from other arguments. - tuple_types: if not None, a list of types of the elements of the queue. - tuple_shapes: if not None, a list of shapes of the elements of the queue. - shard_dimensions: if not None, a list of dimensions on which the - elements of the queue should be sharded during automatic - parallelization. - name: the name of the queue. - - Raises: - ValueError: if number_of_tuple_elements <= 0; or - number_of_tuple_arguments, tuple_types, tuple_shapes, and - shard_dimensions are all None; or the length of tuple_types, - tuple_shapes, or shard_dimensions is not equal to - number_of_tuple_elements; or any element of shard_dimensions - can't be converted to a Dimension. - TypeError: if any element of tuple_types or tuple_shapes can't - be converted to a dtype or TensorShape, respectively. - """ - self._frozen = False - self._generated_enqueue_ops = False - self._generated_dequeue_op = False - self._name = "InfeedQueue" if name is None else name - if number_of_tuple_elements is None: - if tuple_types is not None: - number_of_tuple_elements = len(tuple_types) - elif tuple_shapes is not None: - number_of_tuple_elements = len(tuple_shapes) - elif shard_dimensions is not None: - number_of_tuple_elements = len(shard_dimensions) - else: - raise ValueError( - "number of tuple elements cannot be inferred from InfeedQueue " - "constructor") - if number_of_tuple_elements <= 0: - raise ValueError("number_of_tuple_elements %d must be > 0" % - number_of_tuple_elements) - # Make an empty sharding policy for each tuple element. - self._sharding_policies = [ - tpu_sharding.ShardingPolicy() - for _ in xrange(number_of_tuple_elements) - ] - if tuple_types is not None: - self.set_tuple_types(tuple_types) - else: - self._tuple_types = None - if tuple_shapes is not None: - self.set_tuple_shapes(tuple_shapes) - else: - self._tuple_shapes = None - if shard_dimensions is not None: - self.set_shard_dimensions(shard_dimensions) - self._validate() - - def _validate(self): - """Checks that the configuration is self-consistent. - - Raises: - ValueError: if the shapes and sharding policies don't match. - """ - if self.tuple_shapes is not None: - for (policy, shape) in zip(self._sharding_policies, self._tuple_shapes): - # Raise an error if the policy is incompatible with the shape. - _ = policy.get_sharded_shape(shape) - - @property - def number_of_tuple_elements(self): - """Returns the number of InfeedQueue tuple elements.""" - return len(self._sharding_policies) - - @property - def tuple_types(self): - """Returns the types of the InfeedQueue tuple elements.""" - return self._tuple_types - - def set_tuple_types(self, tuple_types): - """Sets the type of each element of the queue. - - tuple_types must be a list of length - self.number_of_tuple_elements, and each element must be - convertible to a dtype. - - Args: - tuple_types: the types of each queue element. - - Raises: - ValueError: if tuple_types is not of length - self.number_of_tuple_elements. - TypeError: if an element of tuple_types cannot be converted to a - dtype. - """ - if len(tuple_types) != self.number_of_tuple_elements: - raise ValueError("tuple_types is %s, but must be a list of length %d" % - (str(tuple_types), self.number_of_tuple_elements)) - if self._frozen: - for (frozen, updated) in zip(self._tuple_types, tuple_types): - if frozen != updated: - raise ValueError( - "Trying to update InfeedQueue with frozen configuration with an " - "incompatible type. Frozen types are %s, updated types are %s" % ( - str(self._tuple_types), str(tuple_types))) - else: - try: - self._tuple_types = [dtypes.as_dtype(t) for t in tuple_types] - except (TypeError) as e: - raise TypeError( - "tuple_types is %s, but must be a list of elements each " - "convertible to dtype: got error %s" % (str(tuple_types), str(e))) - - @property - def tuple_shapes(self): - """Returns the shapes of the InfeedQueue tuple elements.""" - return self._tuple_shapes - - def set_tuple_shapes(self, tuple_shapes): - """Sets the shape of each element of the queue. - - tuple_shapes must be a list of length - self.number_of_tuple_elements, and each element must be - convertible to a TensorShape. - - Args: - tuple_shapes: the shapes of each queue element. - - Raises: - ValueError: if tuple_shapes is not of length - self.number_of_tuple_elements. - TypeError: if an element of tuple_shapes cannot be converted to - a TensorShape. - """ - if len(tuple_shapes) != self.number_of_tuple_elements: - raise ValueError("tuple_shapes is %s, but must be a list of length %d" % - (str(tuple_shapes), self.number_of_tuple_elements)) - try: - tuple_shapes = [tensor_shape.as_shape(shape) for shape in tuple_shapes] - except (ValueError, TypeError) as e: - raise TypeError( - "tuple_shapes is %s, but must be a list of elements each " - "convertible to TensorShape: got error %s" % (str(tuple_shapes), - str(e))) - if self._frozen: - for (frozen, updated) in zip(self._tuple_shapes, tuple_shapes): - if frozen != updated: - raise ValueError( - "Trying to update InfeedQueue with frozen configuration with an " - "incompatible shape. Frozen shapes are %s, updated shapes are %s" - % (str(self._tuple_shapes), str(tuple_shapes))) - else: - self._tuple_shapes = tuple_shapes - self._validate() - - @property - def sharding_policies(self): - """Returns the sharding policies of the InfeedQueue tuple elements.""" - return self._sharding_policies - - @property - def shard_dimensions(self): - """Gets the shard dimension of each tuple element. - - Returns: - A list of length number_of_tuple_elements, where each list entry - is the shard dimension of that tuple element or None if the - shard dimension has not been set. - """ - # The number of shards is always the same for all the policies. - return [policy.shard_dimension for policy in self._sharding_policies] - - def set_shard_dimensions(self, shard_dimensions): - """Sets the shard_dimension of each element of the queue. - - shard_dimensions must be a list of length - self.number_of_tuple_elements, and each element must be - convertible to a Dimension compatible with self.tuple_shapes. - - Args: - shard_dimensions: the dimensions of each queue element. - - Raises: - ValueError: if shard_dimensions is not of length - self.number_of_tuple_elements; or an element of - shard_dimensions cannot be converted to a Dimension; or an - element of shard_dimensions is a Dimension that is out of - range for the corresponding tuple element shape. - """ - if len(shard_dimensions) != self.number_of_tuple_elements: - raise ValueError("shard_dimensions is %s, but must be a list of length %d" - % (str(shard_dimensions), - self.number_of_tuple_elements)) - for (policy, dimension) in zip(self._sharding_policies, shard_dimensions): - policy.set_shard_dimension(dimension) - self._validate() - - @property - def number_of_shards(self): - """Gets the number of shards to use for the InfeedQueue. - - Returns: - Number of shards or None if the number of shards has not been set. - """ - # The number of shards is always the same for all the policies. - return self._sharding_policies[0].number_of_shards - - def set_number_of_shards(self, number_of_shards): - """Sets the number of shards to use for the InfeedQueue. - - Args: - number_of_shards: number of ways to shard the InfeedQueue. - - Raises: - ValueError: if number_of_shards is not > 0; or the policies have - been frozen and number_of_shards was already set to something - else. - """ - for policy in self._sharding_policies: - policy.set_number_of_shards(number_of_shards) - self._validate() - - def set_configuration_from_input_tensors(self, input_tensors): - """Sets the shapes and types of the queue tuple elements. - - input_tensors is a list of Tensors whose types and shapes are used - to set the queue configuration. - - Args: - input_tensors: list of Tensors of the same types and shapes as - the desired queue Tuple. - - Raises: - ValueError: if input_tensors is not a list of length - self.number_of_tuple_elements - """ - if len(input_tensors) != self.number_of_tuple_elements: - raise ValueError("input_tensors is %s, but should be a list of %d Tensors" - % (str(input_tensors), self.number_of_tuple_elements)) - self.set_tuple_shapes([t.shape for t in input_tensors]) - self.set_tuple_types([t.dtype for t in input_tensors]) - - def set_configuration_from_sharded_input_tensors(self, input_tensors): - """Sets the shapes and types of the queue tuple elements. - - input_tensors is a list of lists of Tensors whose types and shapes are used - to set the queue configuration. The length of the outer list is the number - of shards required, and each inner list is the tuple of Tensors to use to - determine the types and shapes of the corresponding shard. This method - depends on the shard dimension, and calling it freezes the shard policy. - - Args: - input_tensors: list of lists of Tensors. The outer list length corresponds - to the desired number of shards, and each inner list is the size - and shape of the desired configuration of the corresponding shard. - - Raises: - ValueError: if any inner list is not a list of length - self.number_of_tuple_elements; or the inner lists do not combine to - form a consistent unsharded shape. - TypeError: if the types of the Tensors in the inner lists do not match. - """ - if not self._frozen: - # Unset the tuple shapes in case the configuration becomes - # transiently inconsistent. - self._tuple_shapes = None - number_of_shards = len(input_tensors) - self.set_number_of_shards(number_of_shards) - for t in input_tensors: - if len(t) != self.number_of_tuple_elements: - raise ValueError( - "input_tensors is %s but must be a list of lists, where each inner" - " list has length number_of_tuple_elements=%d" % ( - str(input_tensors), self.number_of_tuple_elements)) - # Transpose the inputs to make a list of shard shapes for each tuple - # element. - sharded_shapes = [[t[i].shape for t in input_tensors] - for i in xrange(self.number_of_tuple_elements)] - # For each tuple, get the unsharded shape using that tuple's policy. - unsharded_shapes = [ - policy.get_unsharded_shape(s) - for (policy, s) in zip(self._sharding_policies, sharded_shapes) - ] - self.set_tuple_shapes(unsharded_shapes) - for i in xrange(1, self.number_of_shards): - for (t1, t2) in zip(input_tensors[0], input_tensors[i]): - if t1.dtype != t2.dtype: - raise TypeError( - "types of the tuple elements of input_tensors %s are not " - "consistent" % str(input_tensors)) - self.set_tuple_types([t.dtype for t in input_tensors[0]]) - - def freeze(self): - """Freezes the InfeedQueue so it can no longer be modified. - - The configuration is implicitly frozen before any host-side or - device-side Ops are generated. The configuration cannot be frozen - until the types and shapes of the tuple elements have been set. - - Raises: - ValueError: if the types or shapes of the tuple elements have not been - set. - """ - self._frozen = True - if self._tuple_types is None: - raise ValueError( - "Can't freeze an InfeedQueue without setting all tuple types.") - if self._tuple_shapes is None: - raise ValueError( - "Can't freeze an InfeedQueue without setting all tuple shapes.") - for shape in self._tuple_shapes: - if shape.dims is None: - raise ValueError( - "Can't freeze an InfeedQueue without setting all tuple shapes.") - for policy in self._sharding_policies: - policy.freeze() - self._validate() - - def generate_dequeue_op(self, tpu_device=0): - """Generates the device-side Op to dequeue a tuple from the queue. - - Implicitly freezes the queue configuration if it is not already - frozen, which will raise errors if the shapes and types have not - been fully specified. - - Args: - tpu_device: The TPU device ordinal where the infeed instruction should be - placed. If None, no explicit placement will be performed, and it is up - to the user to call this API from within a proper TPU device scope. - The XLA code will fail if the TPU dequeue instruction is not bound to - any device. - - Returns: - A list of Outputs corresponding to a shard of infeed dequeued - into XLA, suitable for use within a replicated block. - - Raises: - ValueError: if the types or shapes of the tuple elements have not been - set; or if a dequeue op has already been generated. - """ - self.freeze() - if self._generated_dequeue_op: - raise ValueError("Can't generate two dequeue Ops from the same queue") - self._generated_dequeue_op = True - full_name = "%s/dequeue" % self._name - sharded_shapes = [ - policy.get_sharded_shape(shape) - for (shape, policy) in zip(self._tuple_shapes, self._sharding_policies) - ] - if tpu_device is not None: - with ops.device(tpu.core(tpu_device)): - return tpu_ops.infeed_dequeue_tuple( - dtypes=self._tuple_types, shapes=sharded_shapes, name=full_name) - else: - return tpu_ops.infeed_dequeue_tuple( - dtypes=self._tuple_types, shapes=sharded_shapes, name=full_name) - - def _generate_enqueue_op(self, - inputs, - name_prefix, - index, - device=None, - tpu_ordinal=-1): - """Generate a host-side Op to enqueue a tuple to the queue. - - If device is None the inputs are all required to have the same - device specification, and the enqueue Op is colocated with - inputs[0]. Otherwise the enqueue Op is placed on 'device'. - - Args: - inputs: a list of Tensors with the types and shapes of the tuple elements. - name_prefix: the base name for the Op. - index: the shard index, used to uniquify the Op name. - device: device to place the Op on, or None if it should be - colocated with the inputs. - tpu_ordinal: ordinal of the TPU device on the host to use for - infeed if device is a CPU device. Should be set to -1 if device - is a TPU device. - - Returns: - An Op corresponding to a shard of infeed enqueued at the host, - suitable for use within a replicated block. - - Raises: - ValueError: if device is None and inputs do not all have the - same device specification. - """ - full_name = "%s/%d" % (name_prefix, index) - shapes = [t.shape for t in inputs] - if device is None: - devices = [t.device for t in inputs] - for i in xrange(1, self.number_of_tuple_elements): - if devices[0] != devices[i]: - raise ValueError( - "input devices for shard %d are %s, but should all be the same" % - (index, str(devices))) - with ops.colocate_with(inputs[0]): - return tpu_ops.infeed_enqueue_tuple( - inputs=inputs, - shapes=shapes, - name=full_name, - device_ordinal=tpu_ordinal) - else: - with ops.device(device): - return tpu_ops.infeed_enqueue_tuple( - inputs=inputs, - shapes=shapes, - name=full_name, - device_ordinal=tpu_ordinal) - - def generate_enqueue_ops(self, - sharded_inputs, - tpu_ordinal_function=None, - placement_function=None): - """Generates the host-side Ops to enqueue the shards of a tuple. - - sharded_inputs is a list, one for each shard, of lists of - Tensors. sharded_inputs[0] is the tuple of Tensors to use to feed - shard 0 if the queue. Returns the host-side Ops that must be run to - enqueue the sharded tuple. The Op for shard i is colocated with the inputs - for shard i. - - Implicitly freezes the queue configuration if it is not already - frozen. If the configuration has already been frozen, and is not - compatible with the types and shapes of sharded_inputs, an error - will be raised. - - Args: - sharded_inputs: a list of lists of Tensors. The length of the outer list - determines the number of shards. Each inner list indicates the types - and shapes of the tuples in the corresponding shard. - tpu_ordinal_function: if not None, a function that takes the - shard index as input and returns the ordinal of the TPU device - the shard's infeed should be placed on. tpu_ordinal_function must be - set if the inputs are placed on CPU devices. - placement_function: if not None, a function that takes the shard index as - input and returns the host device where the enqueue op should be placed - on. - - Returns: - A list of host-side Ops, one for each shard, that when executed together - will enqueue a full-size element of infeed. - - Raises: - ValueError: if the queue configuration has previously been frozen and the - shapes of the elements of sharded_inputs are not compatible with the - frozen configuration; or if the shapes of the elements of sharded_inputs - don't form a consistent unsharded tuple; or if the elements of a tuple - have different device constraints. - TypeError: if the queue configuration has previously been frozen and the - types of the elements of sharded_inputs are not compatible with the - frozen configuration; or if the types of the elements of sharded_inputs - don't form a consistent unsharded tuple. - """ - self.set_configuration_from_sharded_input_tensors(sharded_inputs) - self.freeze() - if self._generated_enqueue_ops: - raise ValueError("Can't generate two enqueue Ops from the same queue") - self._generated_enqueue_ops = True - if tpu_ordinal_function is None: - tpu_ordinal_function = lambda index: -1 - name_prefix = "%s/enqueue" % self._name - return [ - self._generate_enqueue_op( - shard, - name_prefix, - index, - tpu_ordinal=tpu_ordinal_function(index), - device=placement_function(index) if placement_function else None) - for (shard, index) in zip(sharded_inputs, xrange(self.number_of_shards)) - ] - - # TODO(misard) Generalize this to the case of systems that don't - # have 8 devices per host, and figure out what to do with - # model-parallelism. - def _default_placement_function(self, index): - return "/task:%d/device:CPU:0" % (index / 8) - - def _default_ordinal_function(self, index): - return index % 8 - - # TODO(b/36470756) remove this from tutorials once we have a better story - # for automatic placement of input pipelines. - def split_inputs_and_generate_enqueue_ops(self, - inputs, - device_assignment=None, - placement_function=None, - tpu_ordinal_function=None): - """POORLY-PERFORMING ON MULTI-HOST SYSTEMS. - - Generates the host-side Ops to enqueue a tuple. - - This method performs poorly because it takes an entire input on a single - host, splits it, and distributes it to all of the cores. It is present only - to simplify tutorial examples. - - inputs is a list of Tensors to use to feed the queue. Each input is split - into self.number_of_shards shards. Returns an Op for each shard to enqueue - the shard. The Op for shard i is placed on device placement_function(i). - - Implicitly freezes the queue configuration if it is not already - frozen. If the configuration has already been frozen, and is not - compatible with the types and shapes of inputs, an error - will be raised. - - Args: - inputs: a list of Tensors which indicates the types and shapes of the - queue tuple. - device_assignment: if not `None`, a TPU `DeviceAssignment`. If - device_assignment is not `None`, but `placement_function` and - `ordinal_function` are None, then `device_assignment` will be used to - place infeeds on the first k TPU shards, where k is the number of shards - in the queue. If all three are `None`, then default placement and - ordinal functions are used. - placement_function: if not None, a function that takes the shard - index as input and returns a device string indicating which - device the shard's infeed should be placed on. If placement_function - and tpu_ordinal_function are None, inputs are sharded round-robin - across the devices in the system. - tpu_ordinal_function: if not None, a function that takes the - shard index as input and returns the ordinal of the TPU device - the shard's infeed should be placed on. If placement_function - and tpu_ordinal_function are None, inputs are sharded round-robin - across the devices in the system. - - Returns: - A list of host-side Ops, one for each shard, that when executed together - will enqueue a full-size element of infeed. - - Raises: - ValueError: if the queue configuration has previously been frozen and the - shapes of the elements of inputs are not compatible with the frozen - configuration. - TypeError: if the queue configuration has previously been frozen and the - types of the elements of inputs are not compatible with the frozen - configuration. - """ - if device_assignment is None: - if placement_function is None: - placement_function = self._default_placement_function - if tpu_ordinal_function is None: - tpu_ordinal_function = self._default_ordinal_function - else: - - def _placement_function_from_map(index): - return device_assignment.host_device(replica=index) - - def _ordinal_function_from_map(index): - return device_assignment.tpu_ordinal(replica=index) - - if placement_function is None: - placement_function = _placement_function_from_map - if tpu_ordinal_function is None: - tpu_ordinal_function = _ordinal_function_from_map - self.set_configuration_from_input_tensors(inputs) - self.freeze() - if self._generated_enqueue_ops: - raise ValueError("Can't generate two enqueue Ops from the same queue") - self._generated_enqueue_ops = True - split_name_prefix = "%s/split" % self._name - if self.number_of_shards == 1: - transposed_sharded_inputs = [[inp] for inp in inputs] - else: - - def split_fn(inp, num_shards, axis, name): - with ops.colocate_with(inp): - return array_ops.split(inp, num_shards, axis=axis, name=name) - - transposed_sharded_inputs = [ - split_fn( - inp, - self.number_of_shards, - axis=policy.shard_dimension, - name="%s/%d" % (split_name_prefix, index)) - for (inp, policy, index) in zip(inputs, self._sharding_policies, - xrange(self.number_of_tuple_elements)) - ] - sharded_inputs = [[shard[i] for shard in transposed_sharded_inputs] - for i in xrange(self.number_of_shards)] - name_prefix = "%s/enqueue" % self._name - return [ - self._generate_enqueue_op( - shard, - name_prefix, - index, - device=placement_function(index), - tpu_ordinal=tpu_ordinal_function(index)) - for (shard, index) in zip(sharded_inputs, xrange(self.number_of_shards)) - ] - - -class _PartitionedInfeedQueue(InfeedQueue): - """A helper object to build a device infeed queue with input partition. - - Args: - number_of_tuple_elements: the number of Tensors fed atomically through the - queue, must be present unless it can be inferred from other arguments. - device_assignment: A TPU `DeviceAssignment` which is used to place all the - partitions to different TPU infeed queues. - host_id: The id of the host machine. - input_partition_dims: A nested list/tuple of integers. Each inner - list/tuple describes how to partition the corresponding input tensor. - tuple_types: If not None, a list of types of the elements of the queue. - tuple_shapes: If not None, a list of shapes of the elements of the queue. - name: The name of the queue. - """ - - def __init__(self, - number_of_tuple_elements, - device_assignment, - host_id, - input_partition_dims=None, - tuple_types=None, - tuple_shapes=None, - name=None): - super(_PartitionedInfeedQueue, self).__init__( - number_of_tuple_elements=number_of_tuple_elements, - tuple_types=tuple_types, - tuple_shapes=None, - shard_dimensions=None, - name="PartitionedInfeedQueue" if name is None else name) - self._input_partition_dims = input_partition_dims - self._host_id = host_id - self._device_assignment = device_assignment - - def generate_dequeue_op(self, tpu_device=0): - """Generate TPU dequeue ops. - - Args: - tpu_device: The TPU device ordinal where the infeed instruction should be - placed. - - Returns: - A list of Outputs corresponding to a partition of infeed dequeued - into XLA, suitable for use within a replicated block. - - Raises: - ValueError: if the types or shapes of the tuple elements have not been - set; or if a dequeue op has already been generated. - """ - self.freeze() - if self._generated_dequeue_op: - raise ValueError("Can't generate two dequeue Ops from the same queue") - self._generated_dequeue_op = True - full_name = "%s/dequeue" % self._name - sharded_shapes = [ - policy.get_sharded_shape(shape) - for (shape, policy) in zip(self._tuple_shapes, self._sharding_policies) - ] - with ops.device(tpu.core(tpu_device)): - values = tpu_ops.infeed_dequeue_tuple( - dtypes=self._tuple_types, shapes=sharded_shapes, name=full_name) - return tag_sharding_attribute_for_dequeued_tensors( - values, self._input_partition_dims) - - def generate_enqueue_ops(self, per_host_sharded_inputs): - """Generates the host-side Ops to enqueue the partitioned inputs. - - per_host_sharded_inputs is a list, one for each replica, of lists of - Tensors. sharded_inputs[i] is the tuple of Tensors to use to feed - replica i. - sharded_inputs[i][j] is partitioned by self._input_partition_dims[j]. - - For example, if sharded_inputs[i][j] is a 2-D Tensor: - [[A, B, C, D], - [E ,F, G, H]] - self._input_partition_dims[j] is [2, 4]. - - sharded_inputs[i][j] will be partitioned and flattened into: - [A, B, C, D, E, F, G, H] and fed into the logical core ids: - [0, 1, 2, 3, 4, 5, 6, 7] respectively. - - Args: - per_host_sharded_inputs: a list of lists of Tensors. The length of the - outer list determines the number of shards. Each inner list indicates - the types and shapes of the tuples in the corresponding shard. - - Returns: - A list of host-side Ops, one for each shard, that when executed together - will enqueue a full-size element of infeed. - - Raises: - ValueError: if the queue configuration has previously been frozen and the - shapes of the elements of sharded_inputs are not compatible with the - frozen configuration; or if the shapes of the elements of sharded_inputs - don't form a consistent unsharded tuple; or if the elements of a tuple - have different device constraints; or if the partition dims are invalid. - TypeError: if the queue configuration has previously been frozen and the - types of the elements of sharded_inputs are not compatible with the - frozen configuration; or if the types of the elements of sharded_inputs - don't form a consistent unsharded tuple. - """ - self.set_configuration_from_sharded_input_tensors(per_host_sharded_inputs) - number_of_replicas_per_host = len(per_host_sharded_inputs) - number_of_tuple_elements = len(per_host_sharded_inputs[0]) - - assert len(self._input_partition_dims) == number_of_tuple_elements - per_host_enqueue_ops = [] - - for replica_index in range(number_of_replicas_per_host): - flattened_inputs = per_host_sharded_inputs[replica_index] - inputs_part_dims_flat = nest.flatten_up_to(flattened_inputs, - self._input_partition_dims) - inputs_parted_iters = [ - iter(self._check_dims_and_partition_or_replicate_on_host(x, dims)) - for x, dims in zip(per_host_sharded_inputs[replica_index], - inputs_part_dims_flat) - ] - - for logical_core in xrange(self._device_assignment.num_cores_per_replica): - # Places different partitions to different logic cores. - replica_id = self._device_assignment.lookup_replicas( - self._host_id, logical_core)[replica_index] - ordinal = self._device_assignment.tpu_ordinal( - replica=replica_id, logical_core=logical_core) - infeed_inputs = [] - for it in inputs_parted_iters: - input_for_device = next(it, None) - if input_for_device is not None: - infeed_inputs.append(input_for_device) - - if infeed_inputs: - per_host_enqueue_ops.append( - tpu_ops.infeed_enqueue_tuple( - inputs=infeed_inputs, - shapes=[x.shape for x in infeed_inputs], - name="enqueue/replica_{0}/input_{1}".format( - replica_index, logical_core), - device_ordinal=ordinal)) - return per_host_enqueue_ops - - def _check_input_partition_dims(self, tensor, dims): - """Checks that input partition dims are valid for the `Tensor`. - - Args: - tensor: Input tensor for partitioning. - dims: A list of integer describes how to partition the input tensor. - - Raises: - ValueError: If the tensor can't be partitioned by dims or the - num_cores_per_replica doesn't match the number of - partitions(dims.prod()). - """ - # No partitioning specified, so don't perform further checks. - if dims is None: - return - - dims = np.array(dims) - - if (dims < 1).any(): - raise ValueError("All input partition dims must be >= 1.") - - # No partitioning, so don't perform further checks. - if dims.prod() == 1: - return - - if dims.prod() != self._device_assignment.num_cores_per_replica: - raise ValueError( - "The product of each input parition dim should equal to " - "num_cores_per_replica. (dim = {}, num_cores_per_replica " - "= {})".format(dims, self._device_assignment.num_cores_per_replica)) - if dims.shape[0] != tensor.shape.ndims: - raise ValueError( - "Input partition dims must have the same number of dimensions " - "as the `Tensor` to be partitioned. (tensor shape = {}, input " - "partition dims = {}).".format(tensor.shape.as_list(), dims)) - - tensor.shape.assert_is_fully_defined() - - def _check_dims_and_partition_or_replicate_on_host(self, tensor, dims): - """Checks dims and partitions or replicates the input tensor. - - The ops inside this function are placed on the host side. - - Args: - tensor: The input tensor which will be partioned or replicated. - dims: A list of integer describes how to partition the input tensor. - - Returns: - An iterator of `Tensor`s or a list of partioned tensors. - """ - self._check_input_partition_dims(tensor, dims) - return partition_or_replicate_on_host(tensor, dims) +# pylint: disable=wildcard-import,unused-import,redefined-builtin +from tensorflow.python.tpu.tpu_feed import * +# used by tests +from tensorflow.python.tpu.tpu_feed import _PartitionedInfeedQueue +# pylint: enable=wildcard-import,unused-import,redefined-builtin diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_function.py b/tensorflow/contrib/tpu/python/tpu/tpu_function.py index 422c7d3b26ffb4ad1b72450c4803ac2eb87cea3b..f2755c6979c2e49dbc19b6800462949601811496 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_function.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_function.py @@ -1,66 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ============================================================================= - -"""Helper library for functions used during TPU compilation.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import contextlib - - -class TpuContext(object): - """A context object holding state about the TPU computation being built.""" - - def __init__(self): - """Creates a new TpuContext.""" - self._number_of_shards = None - - @property - def number_of_shards(self): - return self._number_of_shards - - def set_number_of_shards(self, number_of_shards): - self._number_of_shards = number_of_shards - - -# The Tpu context holds the number of shards when a sharded computation is -# being built, or None if no computation is being built. -_current_tpu_context = TpuContext() - - -@contextlib.contextmanager -def tpu_shard_context(number_of_shards): - if _current_tpu_context.number_of_shards is not None: - raise NotImplementedError("tpu_shard_context cannot be nested.") - try: - _current_tpu_context.set_number_of_shards(number_of_shards) - yield - finally: - _current_tpu_context.set_number_of_shards(None) - - -def get_tpu_context(): - return _current_tpu_context - - -# Decorator function for tpu computation func that was passed to tpu.rewrite() -# if there is an embedded training loop in this func, trace tools will generate -# step markers for each iteration. -def on_device_training_loop(func): - # Value for this attribute is from xla.DebugOptions.StepMarkerLocation. - setattr(func, "step_marker_location", "STEP_MARK_AT_TOP_LEVEL_WHILE_LOOP") - return func +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.tpu_function import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_optimizer.py b/tensorflow/contrib/tpu/python/tpu/tpu_optimizer.py index 1e11de6421e360faf0b9ad573a84f9aecdf9c98f..ca58e78d7b342c7ca70400652d99092ccbecbbde 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_optimizer.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_optimizer.py @@ -1,203 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ============================================================================= - -"""Optimizer that implements cross-shard gradient reduction for TPU.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function - -from tensorflow.contrib.tpu.python.ops import tpu_ops -from tensorflow.contrib.tpu.python.tpu import tpu_function -from tensorflow.python.framework import ops -from tensorflow.python.ops.losses import losses -from tensorflow.python.platform import tf_logging as logging -from tensorflow.python.training import optimizer - - -class CrossShardOptimizer(optimizer.Optimizer): - """An optimizer that averages gradients across TPU shards.""" - - def __init__(self, - opt, - reduction=losses.Reduction.MEAN, - name="CrossShardOptimizer", - group_assignment=None): - """Construct a new cross-shard optimizer. - - Args: - opt: An existing `Optimizer` to encapsulate. - reduction: The reduction to apply to the shard losses. - name: Optional name prefix for the operations created when applying - gradients. Defaults to "CrossShardOptimizer". - group_assignment: Optional 2d int32 lists with shape - [num_groups, num_replicas_per_group] which describles how to apply - optimizer to subgroups. - - Raises: - ValueError: If reduction is not a valid cross-shard reduction. - """ - if reduction not in (losses.Reduction.SUM, losses.Reduction.MEAN): - raise ValueError("Unsupported reduction: %s." % reduction) - - super(CrossShardOptimizer, self).__init__(False, name) - self._opt = opt - self._reduction = reduction - self._group_assignment = group_assignment - - def _verify_and_get_subgroup_size(self, group_assignment, num_shards): - """Verify group_assignment and get the subgroup size". - - Args: - group_assignment: list of group ids for applying the optimizer - to subgroups. - num_shards: The number of TPU shards. - - Returns: - The size of one subgroup in group_assignment. - - Raises: - ValueError: If group_assignment is invalid. - """ - if not group_assignment: - return None - if not (isinstance(group_assignment, list) and - all(isinstance(i, list) for i in group_assignment)): - raise ValueError("group_assignment must be a list of list. Got {}".format( - group_assignment)) - - replica_ids = set() - for g in group_assignment: - for i in g: - replica_ids.add(i) - - if set(range(num_shards)) != replica_ids: - raise ValueError("group_assignment must be a permutation of range({0})." - " Got group_assignment={1}".format( - num_shards, group_assignment)) - - subgroup_size_list = [len(group) for group in group_assignment] - if all(subgroup_size_list[0] == size for size in subgroup_size_list): - return subgroup_size_list[0] - else: - raise ValueError("The size of each subgroup in group_assignment must " - "be equal. Got group_assignment={}".format( - self._group_assignment)) - - def compute_gradients(self, loss, var_list=None, **kwargs): - """Compute gradients of "loss" for the variables in "var_list". - - This simply wraps the compute_gradients() from the real optimizer. The - gradients will be aggregated in the apply_gradients() so that user can - modify the gradients like clipping with per replica global norm if needed. - The global norm with aggregated gradients can be bad as one replica's huge - gradients can hurt the gradients from other replicas. - - Args: - loss: A Tensor containing the value to minimize. - var_list: Optional list or tuple of `tf.Variable` to update to minimize - `loss`. Defaults to the list of variables collected in the graph - under the key `GraphKey.TRAINABLE_VARIABLES`. - **kwargs: Keyword arguments for compute_gradients(). - - Returns: - A list of (gradient, variable) pairs. - - Raises: - ValueError: If not within a tpu_shard_context or group_assignment is - invalid. - """ - num_shards = tpu_function.get_tpu_context().number_of_shards - if num_shards is None: - logging.warning( - "CrossShardOptimizer should be used within a tpu_shard_context, but " - "got unset number_of_shards. Assuming 1.") - num_shards = 1 - - subgroup_size = self._verify_and_get_subgroup_size(self._group_assignment, - num_shards) - - if num_shards > 1 and self._reduction == losses.Reduction.MEAN: - if self._group_assignment: - scale = 1.0 / subgroup_size - else: - scale = 1.0 / num_shards - loss *= scale - - return self._opt.compute_gradients(loss, var_list=var_list, **kwargs) - - def apply_gradients(self, grads_and_vars, global_step=None, name=None): - """Apply gradients to variables. - - Calls tpu_ops.cross_replica_sum() to sum gradient contributions across - replicas, and then applies the real optimizer. - - Args: - grads_and_vars: List of (gradient, variable) pairs as returned by - compute_gradients(). - global_step: Optional Variable to increment by one after the - variables have been updated. - name: Optional name for the returned operation. Default to the - name passed to the Optimizer constructor. - - Returns: - An `Operation` that applies the gradients. If `global_step` was not None, - that operation also increments `global_step`. - - Raises: - ValueError: If the grads_and_vars is malformed. - """ - summed_grads_and_vars = [] - for (grad, var) in grads_and_vars: - if grad is None: - summed_grads_and_vars.append((grad, var)) - else: - with ops.colocate_with(grad): - summed_grads_and_vars.append((tpu_ops.cross_replica_sum( - grad, self._group_assignment), var)) - return self._opt.apply_gradients(summed_grads_and_vars, global_step, name) - - def get_slot(self, *args, **kwargs): - """Return a slot named "name" created for "var" by the Optimizer. - - This simply wraps the get_slot() from the actual optimizer. - - Args: - *args: Arguments for get_slot(). - **kwargs: Keyword arguments for get_slot(). - - Returns: - The `Variable` for the slot if it was created, `None` otherwise. - """ - return self._opt.get_slot(*args, **kwargs) - - def get_slot_names(self, *args, **kwargs): - """Return a list of the names of slots created by the `Optimizer`. - - This simply wraps the get_slot_names() from the actual optimizer. - - Args: - *args: Arguments for get_slot(). - **kwargs: Keyword arguments for get_slot(). - - Returns: - A list of strings. - """ - return self._opt.get_slot_names(*args, **kwargs) - - def variables(self): - """Forwarding the variables from the underlying optimizer.""" - return self._opt.variables() +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.tpu_optimizer import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_sharding.py b/tensorflow/contrib/tpu/python/tpu/tpu_sharding.py index f5af03f33ca8f13af517007672e9ce0e12be6205..93c52335a582e5fa83092f78212ca268079b7c12 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_sharding.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_sharding.py @@ -1,253 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ============================================================================= - -"""Helper library for sharding during TPU compilation.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -from six.moves import xrange # pylint: disable=redefined-builtin - -from tensorflow.python.framework import tensor_shape - -_DEFAULT_NUMBER_OF_SHARDS = 1 -_DEFAULT_SHARD_DIMENSION = 0 - - -# TODO(b/36777903) change other parts of tpu.py to use this class. -class ShardingPolicy(object): - """An object use to hold the sharding policy for a Tensor. - """ - - def __init__(self): - self._number_of_shards = None - self._shard_dimension = None - self._frozen = False - - def __str__(self): - if self.number_of_shards is None or self.shard_dimension is None: - return "ShardingPolicy(unset)" - else: - return ("ShardingPolicy(%d shards dimension %d)" % - (self.number_of_shards, self.shard_dimension)) - - def _fill_default_values(self): - if self._number_of_shards is None: - self._number_of_shards = _DEFAULT_NUMBER_OF_SHARDS - if self._shard_dimension is None: - self._shard_dimension = tensor_shape.as_dimension( - _DEFAULT_SHARD_DIMENSION) - - def freeze(self): - """Prevents further modification to the sharding policy. - - Any values that have not been set when freeze is called are set to - defaults. If the ShardingPolicy is already frozen, this is a NoOp. - """ - if not self._frozen: - self._fill_default_values() - self._frozen = True - - @property - def number_of_shards(self): - """Returns the number of shards in the policy or None if unspecified.""" - return self._number_of_shards - - def set_number_of_shards(self, number_of_shards): - """Sets the number of shards for the current policy. - - If the policy has been frozen then number_of_shards must match the - existing setting. - - Args: - number_of_shards: The number of shards to use in the policy. - - Raises: - ValueError: If the policy has been frozen and number_of_shards - differs from the frozen value; or number_of_shards <= 0. - """ - if self._frozen: - if self._number_of_shards != number_of_shards: - raise ValueError( - "Can't set sharding policy to use %d shards since it has been " - "frozen to use %d." % (number_of_shards, self._number_of_shards)) - else: - if number_of_shards > 0: - self._number_of_shards = number_of_shards - else: - raise ValueError( - "Can't set sharding policy to use %s shards; value must be >0", - str(number_of_shards)) - - @property - def shard_dimension(self): - """Returns the shard dimension of the policy or None if unspecified.""" - return self._shard_dimension - - def set_shard_dimension(self, shard_dimension): - """Sets the shard dimension for the current policy. - - If the policy has been frozen then shard_dimension must match the - existing setting. - - Args: - shard_dimension: The shard dimension to use in the policy. - - Raises: - ValueError: If the policy has been frozen and shard_dimension - differs from the frozen value, or shard_dimension can't be - interpreted as a Dimension. - """ - if self._frozen: - if self._shard_dimension != shard_dimension: - raise ValueError( - "Can't set shard dimension to %d since it has been frozen to " - "use %d." % (shard_dimension, self._shard_dimension)) - else: - self._shard_dimension = tensor_shape.as_dimension(shard_dimension) - - def merge(self, other): - """Merges the policy of another policy into the current policy. - - Args: - other: The policy to merge into this one. - - Raises: - ValueError: If this policy has been frozen and the merge conflicts with - the frozen policy. - """ - if other.number_of_shards is not None: - self.set_number_of_shards(other.number_of_shards) - if other.shard_dimension is not None: - self.set_shard_dimension(other.shard_dimension) - - def get_sharded_shape(self, shape, shard_index=None): - """Returns the shape of a shard of a full Tensor. - - When given the shape of a 'full-size' Tensor, returns the shape of - the sub-Tensor after it has been sharded. Freezes the policy if it - has not yet been frozen. - - Args: - shape: The shape of the full-size Tensor to be sharded. - shard_index: The index of the shard whose shape should be returned. - shard_index can be None for sharding policies that use the same - shape for every shard. - freeze_config: - - Returns: - The shape of the sharded version of the Tensor. - - Raises: - ValueError: If shard_index is None when shards are of different - shapes; or shard_index is not None and - !(0<=shard_index= self.number_of_shards: - raise ValueError("shard_index %d, but must be in [0,%d)." % - (shard_index, self._number_of_shards)) - shape = tensor_shape.as_shape(shape) - if self._number_of_shards == 1: - # Don't do anything when there's only one shard. - return shape - ndims = shape.ndims - if ndims is None: - raise ValueError("shape must be a specified shape not Unknown") - if ndims <= self._shard_dimension: - raise ValueError("shape %s does not contain shard_dimension %d" % - (shape.as_list(), self._shard_dimension)) - dims = shape.as_list() - if dims[self._shard_dimension] is None: - raise ValueError("shape %s must have a fixed size for dimension %d " - "that is known at graph construction time." % - (shape.as_list(), self._shard_dimension)) - if (dims[self._shard_dimension] % self._number_of_shards) != 0: - raise ValueError("shape %s cannot be sharded %d ways along dimension %d" % - (shape.as_list(), self._number_of_shards, - self._shard_dimension)) - dims[self._shard_dimension] /= self._number_of_shards - return tensor_shape.as_shape(dims) - - def _unshard_shape(self, shape): - """Return the unsharded shape that would generate a given sharded shape. - - Args: - shape: the sharded shape to unshard - - Returns: - The unsharded shape. - - Raises: - ValueError: if shape is unknown or does not contain - self.shard_dimension - TypeError: if shape is not convertible to a TensorShape - """ - shape = tensor_shape.as_shape(shape) - if self._number_of_shards == 1: - # Don't do anything when there's only one shard. - return shape - ndims = shape.ndims - if ndims is None: - raise ValueError("shape must be a specified shape not Unknown") - if ndims <= self._shard_dimension: - raise ValueError("shape %s does not contain shard_dimension %d" % - (shape.as_list(), self._shard_dimension)) - dims = shape.as_list() - dims[self._shard_dimension] *= self._number_of_shards - return tensor_shape.as_shape(dims) - - def get_unsharded_shape(self, shapes): - """Returns the shape of an unsharded Tensor given a list of shards. - - When given a list of shapes of shards, returns the shape of the - unsharded Tensor that would generate the shards. Sets defaults for the - policy if number_of_shards or shard_dimension is None. - - Args: - shapes: The shapes of the Tensor shards to be combined. - - Returns: - The shape of the unsharded version of the Tensor. - - Raises: - ValueError: if shapes is not a list of length - self.number_of_shards; or any element of shapes is not a valid - shape consistent with the sharding policy; or the list of - shapes is not a valid sharding of a full shape. - TypeError: if an element of shapes is not convertible to a - TensorShape - """ - self._fill_default_values() - if len(shapes) != self.number_of_shards: - raise ValueError( - "shapes is %s but must be a list of length number_of_shards=%d" % ( - str(shapes), self.number_of_shards)) - unsharded_shapes = [self._unshard_shape(s) for s in shapes] - for i in xrange(self.number_of_shards - 1): - if not unsharded_shapes[i].is_compatible_with( - unsharded_shapes[self.number_of_shards - 1]): - raise ValueError( - "sharded shapes %s are not consistent shards of a full shape " - "sharded %d ways along dimension %d" % ( - str(shapes), self.number_of_shards, self.shard_dimension)) - return unsharded_shapes[0] +# pylint: disable=wildcard-import,unused-import,redefined-builtin +from tensorflow.python.tpu.tpu_sharding import * +# pylint: enable=wildcard-import,unused-import,redefined-builtin diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py b/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py index d66ecfcf4a56b8da1c2d2f518bebe4baa76b315e..258d34ddaf5250e49c5a354caf018e4b64abae62 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py @@ -1,156 +1,25 @@ -# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# =================================================================== -"""TPU system metadata and associated tooling.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections -import re - -from tensorflow.contrib.tpu.python.tpu import tpu -from tensorflow.core.protobuf import config_pb2 -from tensorflow.python.client import session as session_lib -from tensorflow.python.framework import errors -from tensorflow.python.framework import ops -from tensorflow.python.platform import tf_logging as logging - -_PINGING_MASTER_TIMEOUT_IN_MS = 60 * 1000 # 1 min -_RETRY_TIMES = 120 -_INITIAL_TPU_SYSTEM_TIMEOUT_IN_MS = 300 * 1000 # 5 mins - -_TPU_DEVICE_REG = re.compile(r'.*task:(\d+)/.*device:TPU:(\d+)$') - -# _TPUSystemMetadata is used by TPUEstimator to hold TPU configuration, -# including num_cores and num_hosts. -_TPUSystemMetadata = collections.namedtuple('_TPUSystemMetadata', [ - 'num_cores', - 'num_hosts', - 'num_of_cores_per_host', - 'topology', - 'devices', -]) - - -def _query_tpu_system_metadata(master_address, cluster_def=None, - query_topology=False): - """Automatically detects the TPU system metadata in the system.""" - tpu_core_count = 0 - devices = [] - device_dict = collections.defaultdict(list) - - # TODO(b/120564445): Replace with standard library for retries. - retry_count = 1 - while True: - logging.info('Querying Tensorflow master (%s) for TPU system metadata.', - master_address) - try: - with ops.Graph().as_default(): - with session_lib.Session( - master_address, - config=get_session_config_with_timeout( - _PINGING_MASTER_TIMEOUT_IN_MS, - cluster_def)) as sess: - devices = sess.list_devices() - for device in devices: - match = _TPU_DEVICE_REG.match(device.name) - if match: - host_id = match.group(1) - core_id = match.group(2) - device_dict[host_id].append(core_id) - tpu_core_count += 1 - break - except errors.DeadlineExceededError: - msg = ('Failed to connect to the Tensorflow master. The TPU worker may ' - 'not be ready (still scheduling) or the Tensorflow master address ' - 'is incorrect: got (%s).' % - (master_address)) - - # TODO(xiejw): For local or grpc master we might not need retry logic - # here. - if retry_count <= _RETRY_TIMES: - logging.warning('%s', msg) - logging.warning('Retrying (%d/%d).', retry_count, _RETRY_TIMES) - retry_count += 1 - else: - raise ValueError(msg) - - num_of_cores_per_host = 0 - if tpu_core_count: - num_cores_per_host_set = set( - [len(core_ids) for core_ids in device_dict.values()]) - if len(num_cores_per_host_set) != 1: - raise RuntimeError( - 'TPU cores on each host is not same. This should not happen!. ' - 'devices: {}'.format(devices)) - num_of_cores_per_host = num_cores_per_host_set.pop() - - topology = None - if query_topology: - if not tpu_core_count: - raise RuntimeError( - 'Cannot find any TPU cores in the system (master address {}). ' - 'This usually means the master address is incorrect or the ' - 'TPU worker has some problems. Available devices: {}'.format( - master_address, devices)) - - topology = _obtain_topology(master_address, cluster_def) - - metadata = _TPUSystemMetadata( - num_cores=tpu_core_count, - num_hosts=len(device_dict), - num_of_cores_per_host=num_of_cores_per_host, - topology=topology, - devices=devices) - - if tpu_core_count: - logging.info('Found TPU system:') - logging.info('*** Num TPU Cores: %d', metadata.num_cores) - logging.info('*** Num TPU Workers: %d', metadata.num_hosts) - logging.info('*** Num TPU Cores Per Worker: %d', - metadata.num_of_cores_per_host) - for device in metadata.devices: - logging.info('*** Available Device: %s', device) - else: - logging.info('Failed to find TPU: %s', metadata) - return metadata - - -def _obtain_topology(master_address, cluster_def): - """Obtains TPU fabric topology.""" - try: - logging.info('Initializing TPU system (master: %s) to fetch topology ' - 'for model parallelism. This might take a while.', - master_address) - with ops.Graph().as_default(): - session_config = get_session_config_with_timeout( - _INITIAL_TPU_SYSTEM_TIMEOUT_IN_MS, cluster_def) - with session_lib.Session( - master_address, config=session_config) as sess: - topology = sess.run(tpu.initialize_system()) - return topology - except errors.DeadlineExceededError: - raise ValueError( - 'Fail to initialize TPU system with master (%s). ' - 'Please double check the TPU system is functional.' % ( - master_address)) - - -def get_session_config_with_timeout(timeout_in_secs, cluster_def): - """Returns a session given a timeout and a cluster configuration.""" - config = config_pb2.ConfigProto( - operation_timeout_in_ms=timeout_in_secs, cluster_def=cluster_def) - return config +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.tpu_system_metadata import * +# used by tests +from tensorflow.python.tpu.tpu_system_metadata import _query_tpu_system_metadata +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_test.py b/tensorflow/contrib/tpu/python/tpu/tpu_test.py deleted file mode 100644 index 6bdaa528f9f946ae4b9813d554409da2406b1f8d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tpu/python/tpu/tpu_test.py +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================= - -"""Tests for tpu_function helpers.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.tpu.python.tpu import tpu -from tensorflow.contrib.tpu.python.tpu import tpu_feed -from tensorflow.contrib.tpu.python.tpu import training_loop - -from tensorflow.python.framework import dtypes -from tensorflow.python.layers import convolutional -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import control_flow_ops -from tensorflow.python.ops import control_flow_util -from tensorflow.python.ops import math_ops - -from tensorflow.python.platform import test - - -class TPUContextTest(test.TestCase): - - def testIsInContext(self): - """Test that control_flow_util can check that we're in a TPU context.""" - z1 = array_ops.identity(1) - pivot = control_flow_ops.no_op() - context = tpu.TPUReplicateContext(b"context", 1, pivot=pivot) - context.Enter() - z2 = array_ops.identity(1) - context.Exit() - self.assertFalse(control_flow_util.IsInXLAContext(z1.op)) - self.assertTrue(control_flow_util.IsInXLAContext(z2.op)) - - -class TPULayerRewriteTest(test.TestCase): - - def testUsingInfeedQueueWithRegularizer(self): - """Test that Layer regularizers can reference data created in loops.""" - - def make_regularizer(scale): - return lambda inputs: scale * math_ops.reduce_sum(math_ops.square(inputs)) - - def training_step(inputs, scale): - outputs = convolutional.conv2d( - inputs, - filters=16, - kernel_size=(3, 3), - data_format="channels_first", - kernel_regularizer=make_regularizer(scale)) - loss = math_ops.reduce_mean(math_ops.square(outputs)) - return loss.op - - inputs = array_ops.zeros(shape=(128, 32, 32, 16)) - scale = array_ops.ones(shape=()) - infeed = tpu_feed.InfeedQueue( - tuple_types=[dtypes.float32, dtypes.float32], - tuple_shapes=[inputs.shape, scale.shape]) - - def loop(): - return training_loop.repeat(5, training_step, infeed_queue=infeed) - - # This should not throw an error. - tpu.rewrite(loop) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/tpu/python/tpu/training_loop.py b/tensorflow/contrib/tpu/python/tpu/training_loop.py index 50848e83f0ef8d999206909ebfe1b0bbc78d1e5b..673359b232d6857d468723873c449cb3e48168c7 100644 --- a/tensorflow/contrib/tpu/python/tpu/training_loop.py +++ b/tensorflow/contrib/tpu/python/tpu/training_loop.py @@ -1,223 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# ============================================================================= - -"""Library for constructing a training loop, suitable for TPUs.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.compiler import xla -from tensorflow.contrib.tpu.python.tpu import tensor_tracer -from tensorflow.contrib.tpu.python.tpu import tpu_function - -from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import control_flow_ops - - -def while_loop(condition, body, inputs=None, infeed_queue=None, name=None): - """Builds a training loop for TPUs. - - The set of loop-carried tensors corresponds to `inputs`. Both - `condition` and `body` take the current value of the loop-carried - tensors. 'body' additionally takes a tuple of infeed from - infeed_queue if infeed_queue is not None. `condition` must return a - single boolean value that determines whether iteration - continues. `body` must return an updated list of values for the - loop-carried tensors. - - Args: - condition: a Python function that builds the loop condition. - body: a Python function that builds the loop body. - inputs: a list of initial values passed into the training loop, or - None (equivalent to an empty list). - infeed_queue: if not None, the infeed queue from which to append a tuple - of arguments as inputs to condition. - name: (Deprecated) Does nothing. - - Returns: - The final values of the loop-carried tensors. - - Raises: - TypeError: if body or condition has the wrong signature. - """ - del name - # Converts inputs to Tensors. - inputs = [] if inputs is None else [ops.convert_to_tensor(x) for - x in inputs] - input_types = [x.dtype for x in inputs] - input_arity = len(inputs) - - body_arg_error = xla.check_function_argument_count( - body, input_arity, infeed_queue) - if body_arg_error is not None: - if infeed_queue is None: - raise TypeError( - "Supplied loop body function cannot be called with the specified " - "inputs. You specified %d inputs: %s, but the loop body needs %s" % ( - input_arity, str([i.name for i in inputs]), body_arg_error)) - else: - raise TypeError( - "Supplied loop body function cannot be called with the specified " - "inputs. You specified %d inputs: %s and %d additional inputs from " - "infeed, but the computation needs %s" % (input_arity, str( - [i.name for i in inputs]), infeed_queue.number_of_tuple_elements, - body_arg_error)) - condition_arg_error = xla.check_function_argument_count( - condition, input_arity, None) - if condition_arg_error is not None: - if infeed_queue is None: - raise TypeError( - "Supplied loop condition function cannot be called with the " - "specified inputs. You specified %d inputs: %s, but the loop " - "condition needs %s" % (input_arity, str([i.name for i in inputs]), - condition_arg_error)) - else: - raise TypeError( - "Supplied loop condition function cannot be called with the " - "specified inputs. You specified %d inputs: %s, but the loop " - "condition needs %s. Note that infeed is not passed to the loop " - "condition." % (input_arity, str([i.name for i in inputs]), - condition_arg_error)) - - def condition_wrapper(*inputs): - # Discards the dummy output added for arity-0 loops. - if input_arity == 0: - inputs = [] - return condition(*inputs) - - def body_wrapper(*inputs): - """Wrapper around `body` that handles infeed queues and control deps.""" - inputs = list(inputs) - - # Discards the dummy output added for arity-0 loops. - if input_arity == 0: - inputs = [] - - # Runs `body` with the dequeue_ops appended. - if infeed_queue: - number_of_shards = tpu_function.get_tpu_context().number_of_shards - if number_of_shards is None: - raise ValueError("Can't build training loop with infeed when there is " - "no tpu_shard_context. Are you building a loop or " - "graph directly rather than from inside tpu.rewrite, " - "tpu.batch_parallel, tpu.shard, or tpu.replicate?") - infeed_queue.set_number_of_shards(number_of_shards) - dequeue_ops = [d for d in infeed_queue.generate_dequeue_op()] - else: - dequeue_ops = [] - outputs = body(*(inputs + dequeue_ops)) - - # If the computation only returned one value, make it a tuple. - if not isinstance(outputs, (list, tuple)): - outputs = (outputs,) - - outputs = [ - o if isinstance(o, ops.Operation) else ops.convert_to_tensor(o) - for o in outputs - ] - - # Separates the returned Operations and Tensors. - output_operations = [o for o in outputs if isinstance(o, ops.Operation)] - output_tensors = [o for o in outputs - if not isinstance(o, ops.Operation)] - - if outputs != output_tensors + output_operations: - raise ValueError( - "TPU training loop body must return zero or more Tensor values " - "followed by zero or more Operations.") - - output_types = [op.dtype for op in output_tensors] - if input_types != output_types: - raise TypeError( - "Mismatch between input types and output types for training loop " - "body: {} vs {}".format(input_types, output_types)) - - # Add the dequeue operations to output_operations to ensure they are run - # by the loop, even if the programmer's loop body does not use them. - output_operations += dequeue_ops - - # Add a dummy output, if needed. - if not output_tensors: - output_tensors = array_ops.constant(0) - - if output_operations: - # TODO(phawkins): in principle this is too restrictive since it serializes - # the training loop steps. In practice it does not matter since this loop - # will be compiled by XLA. - output_tensors = control_flow_ops.tuple(output_tensors, - control_inputs=output_operations) - - if tensor_tracer.TensorTracer.is_enabled(): - num_replicas = tpu_function.get_tpu_context().number_of_shards - if num_replicas is None: - num_replicas = 1 - tt = tensor_tracer.TensorTracer() - output_tensors = tt.trace_tpu(ops.get_default_graph(), - output_tensors, None, - num_replicas) - return output_tensors - - # If the body has arity 0, add a dummy loop-carried value to which we can add - # control dependencies from any side-effecting operations. - if input_arity == 0: - inputs = [array_ops.constant(0)] - return control_flow_ops.while_loop( - condition_wrapper, body_wrapper, inputs, name="", parallel_iterations=1) - - -def repeat(n, body, inputs=None, infeed_queue=None, name=None): - """Builds a training loop that executes a fixed number of iterations. - - The set of loop-carried tensors correspond to `inputs`. - `body` must be a function that takes and returns the values of the - loop-carried tensors. - - Args: - n: the number of loop iterations - body: a Python function that builds the loop body. - inputs: a list of initial values passed into the training loop or - None (equivalent to an empty list). - infeed_queue: if not None, the infeed queue from which to append a tuple - of arguments as inputs to condition. - name: (Deprecated) Does nothing. - Returns: - The final values of the loop-carried tensors. - Raises: - ValueError: if there is a type error. - """ - def _convert_to_list(xs): - if not isinstance(xs, (list, tuple)): - return [xs] - else: - return list(xs) - - def cond(i, *args): - del args - return i < n - - def body_wrapper(i, *args): - return [i + 1] + _convert_to_list(body(*args)) - - inputs = [0] if inputs is None else [0] + _convert_to_list(inputs) - outputs = while_loop( - cond, body_wrapper, inputs=inputs, infeed_queue=infeed_queue, name=name) - outputs = _convert_to_list(outputs) - if len(outputs) == 1: - # Returns the Op rather than an empty list. - return outputs[0].op - else: - return outputs[1:] +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.training_loop import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/tpu/python/tpu/util.py b/tensorflow/contrib/tpu/python/tpu/util.py index dfb8ce1d1821da05c853bb0d10b1db3a857ccb1b..8d9b70d46eb42c9a525eeafc51d07f0ad4241d52 100644 --- a/tensorflow/contrib/tpu/python/tpu/util.py +++ b/tensorflow/contrib/tpu/python/tpu/util.py @@ -1,51 +1,23 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# Copyright 2019 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# =================================================================== - -"""Utilities for the functionalities.""" +# ============================================================================== +"""Stub file to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import time -import six - -from tensorflow.python.platform import tf_logging as logging -from tensorflow.python.training import training - -def check_positive_integer(value, name): - """Checks whether `value` is a positive integer.""" - if not isinstance(value, six.integer_types): - raise TypeError('{} must be int, got {}'.format(name, type(value))) - - if value <= 0: - raise ValueError('{} must be positive, got {}'.format(name, value)) - - -# TODO(b/118302029) Remove this copy of MultiHostDatasetInitializerHook after we -# release a tensorflow_estimator with MultiHostDatasetInitializerHook in -# python/estimator/util.py. -class MultiHostDatasetInitializerHook(training.SessionRunHook): - """Creates a SessionRunHook that initializes all passed iterators.""" - - def __init__(self, dataset_initializers): - self._initializers = dataset_initializers - - def after_create_session(self, session, coord): - del coord - start = time.time() - session.run(self._initializers) - logging.info('Initialized dataset iterators in %d seconds', - time.time() - start) +# pylint: disable=wildcard-import,unused-import +from tensorflow.python.tpu.util import * +# pylint: enable=wildcard-import,unused-import diff --git a/tensorflow/contrib/training/python/training/hparam.py b/tensorflow/contrib/training/python/training/hparam.py index 27f0d9b2e38c433d4fb4573285ecb8c9946112e8..cb0a25f333b2bba9c4eee991180eab2a083eeb31 100644 --- a/tensorflow/contrib/training/python/training/hparam.py +++ b/tensorflow/contrib/training/python/training/hparam.py @@ -353,8 +353,10 @@ class HParams(object): def my_program(): # Create a HParams object specifying the names and values of the # model hyperparameters: - hparams = tf.HParams(learning_rate=0.1, num_hidden_units=100, - activations=['relu', 'tanh']) + hparams = tf.contrib.training.HParams( + learning_rate=0.1, + num_hidden_units=100, + activations=['relu', 'tanh']) # Override hyperparameters values by parsing the command line hparams.parse(args.hparams) @@ -387,7 +389,7 @@ class HParams(object): # Define 3 hyperparameters: 'learning_rate' is a float parameter, # 'num_hidden_units' an integer parameter, and 'activation' a string # parameter. - hparams = tf.HParams( + hparams = tf.contrib.training.HParams( learning_rate=0.1, num_hidden_units=100, activation='relu') hparams.activation ==> 'relu' diff --git a/tensorflow/contrib/verbs/grpc_verbs_service.h b/tensorflow/contrib/verbs/grpc_verbs_service.h index 444c863b942ef8bce8d54d59765563b12eb6087e..e616778665a9c95b30099b128ec5d1e181ba0618 100644 --- a/tensorflow/contrib/verbs/grpc_verbs_service.h +++ b/tensorflow/contrib/verbs/grpc_verbs_service.h @@ -25,12 +25,6 @@ limitations under the License. #include "tensorflow/core/distributed_runtime/rpc/grpc_call.h" #include "tensorflow/core/lib/core/refcount.h" -namespace grpc { -class ServerBuilder; -class ServerCompletionQueue; -class Alarm; -} // namespace grpc - namespace tensorflow { class GrpcVerbsService : public AsyncServiceInterface { diff --git a/tensorflow/contrib/verbs/grpc_verbs_service_impl.h b/tensorflow/contrib/verbs/grpc_verbs_service_impl.h index cfb9b7ddd7d88c150e47caff66f0865fcaec662c..2432c34ae2353d5d7bca03d80a043b5875ef8cce 100644 --- a/tensorflow/contrib/verbs/grpc_verbs_service_impl.h +++ b/tensorflow/contrib/verbs/grpc_verbs_service_impl.h @@ -27,14 +27,6 @@ limitations under the License. #include "tensorflow/contrib/verbs/verbs_service.pb.h" -namespace grpc { -class CompletionQueue; -class Channel; -class RpcService; -class ServerCompletionQueue; -class ServerContext; -} // namespace grpc - namespace tensorflow { namespace grpc { diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD index 906e8695cd36722a69810f5e20eb31d92528b554..8f5de68322075b6cf44268ae815f42e7a288e12f 100644 --- a/tensorflow/core/BUILD +++ b/tensorflow/core/BUILD @@ -128,6 +128,9 @@ load( "tf_additional_libdevice_srcs", "tf_additional_minimal_lib_srcs", "tf_additional_mpi_lib_defines", + "tf_additional_numa_deps", + "tf_additional_numa_lib_defines", + "tf_additional_numa_copts", "tf_additional_proto_hdrs", "tf_additional_proto_srcs", "tf_additional_test_deps", @@ -236,6 +239,8 @@ ADDITIONAL_CORE_PROTO_SRCS = [ "protobuf/meta_graph.proto", "protobuf/named_tensor.proto", "protobuf/saved_model.proto", + "protobuf/saved_object_graph.proto", + "protobuf/struct.proto", "protobuf/tensorflow_server.proto", "protobuf/transport_options.proto", "util/test_log.proto", @@ -386,15 +391,15 @@ cc_library( ":platform_port_hdrs", ":platform_port_internal_hdrs", ], - copts = tf_copts(), + copts = tf_copts() + tf_additional_numa_copts(), visibility = ["//tensorflow/core:__subpackages__"], deps = [ ":lib_platform", ":platform_base", - "//tensorflow/core/platform/default/build_config:port", "@com_google_absl//absl/base", + "//tensorflow/core/platform/default/build_config:port", "@snappy", - ], + ] + tf_additional_numa_deps(), ) filegroup( @@ -966,7 +971,10 @@ tf_cuda_library( "util/mkl_util.h", ]), visibility = ["//visibility:public"], - deps = [":framework_internal"], + deps = [ + ":framework_internal", + "@com_google_absl//absl/base", + ], ) cc_library( @@ -2273,11 +2281,14 @@ LIB_INTERNAL_PUBLIC_HEADERS = tf_additional_lib_hdrs() + [ ] # Replicated for lib_internal and lib_internal_impl. -LIB_INTERNAL_DEFINES = (tf_additional_lib_defines() + [ - "TF_USE_SNAPPY", - ] + tf_additional_verbs_lib_defines() + - tf_additional_mpi_lib_defines() + - tf_additional_gdr_lib_defines()) +LIB_INTERNAL_DEFINES = ( + tf_additional_lib_defines() + [ + "TF_USE_SNAPPY", + ] + tf_additional_verbs_lib_defines() + + tf_additional_mpi_lib_defines() + + tf_additional_gdr_lib_defines() + + tf_additional_numa_lib_defines() +) cc_library( name = "lib_internal", @@ -2350,19 +2361,20 @@ cc_library( copts = tf_copts(), defines = LIB_INTERNAL_DEFINES, deps = tf_additional_lib_deps() + [ - ":lib_hash_crc32c_accelerate_internal", - ":lib_proto_parsing", - ":abi", - ":core_stringpiece", - "@com_google_absl//absl/memory", - "@com_google_absl//absl/strings", - "//third_party/eigen3", - "//tensorflow/core/platform/default/build_config:platformlib", - "@snappy", - "@zlib_archive//:zlib", - "@double_conversion//:double-conversion", - "@protobuf_archive//:protobuf", - ] + tf_protos_all_impl() + tf_protos_grappler_impl(), + ":lib_hash_crc32c_accelerate_internal", + ":lib_proto_parsing", + ":abi", + ":core_stringpiece", + "@com_google_absl//absl/memory", + "@com_google_absl//absl/strings", + "//third_party/eigen3", + "//tensorflow/core/platform/default/build_config:platformlib", + "@snappy", + "@zlib_archive//:zlib", + "@double_conversion//:double-conversion", + "@protobuf_archive//:protobuf", + ] + tf_protos_all_impl() + tf_protos_grappler_impl() + + tf_additional_numa_deps(), ) # File compiled with extra flags to get cpu-specific acceleration. @@ -2786,6 +2798,7 @@ tf_cuda_library( ":stats_calculator_portable", ":version_lib", "@com_google_absl//absl/base", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/strings", "//tensorflow/core/platform/default/build_config:platformlib", diff --git a/tensorflow/core/api_def/base_api/api_def_AllToAll.pbtxt b/tensorflow/core/api_def/base_api/api_def_AllToAll.pbtxt index d6f28bd022bcd843aa3a7aeb8b1b257a3b3ddfd3..e585bae4a373c6d5afe217b74acd37caa0262023 100644 --- a/tensorflow/core/api_def/base_api/api_def_AllToAll.pbtxt +++ b/tensorflow/core/api_def/base_api/api_def_AllToAll.pbtxt @@ -1,5 +1,6 @@ op { graph_op_name: "AllToAll" + visibility: HIDDEN in_arg { name: "input" description: <