diff --git a/.bazelrc b/.bazelrc index ceba7bfdbac74d1e44aadc3010e5e84bd36ce3ee..17285afdb381018d0054e771475327b1f7ed9866 100644 --- a/.bazelrc +++ b/.bazelrc @@ -25,12 +25,14 @@ build --define framework_shared_object=true # If you would like to use a local MKL instead of downloading, please set the # environment variable "TF_MKL_ROOT" every time before build. build:mkl --define=build_with_mkl=true --define=enable_mkl=true +build:mkl --define=tensorflow_mkldnn_contraction_kernel=0 build:mkl -c opt # This config option is used to enable MKL-DNN open source library only, # without depending on MKL binary version. build:mkl_open_source_only --define=build_with_mkl_dnn_only=true build:mkl_open_source_only --define=build_with_mkl=true --define=enable_mkl=true +build:mkl_open_source_only --define=tensorflow_mkldnn_contraction_kernel=0 build:download_clang --crosstool_top=@local_config_download_clang//:toolchain build:download_clang --define=using_clang=true @@ -78,7 +80,7 @@ build --define=use_fast_cpp_protos=true build --define=allow_oversize_protos=true build --spawn_strategy=standalone -build --genrule_strategy=standalone +build --strategy=Genrule=standalone build -c opt # Other build flags. @@ -88,14 +90,17 @@ build --define=grpc_no_ares=true build:dynamic_kernels --define=dynamic_loaded_kernels=true build:dynamic_kernels --copt=-DAUTOLOAD_DYNAMIC_KERNELS +# Build TF with C++ 17 features. +build:c++17 --cxxopt=-std=c++1z +build:c++17 --cxxopt=-stdlib=libc++ +build:c++1z --cxxopt=-std=c++1z +build:c++1z --cxxopt=-stdlib=libc++ + # Default paths for TF_SYSTEM_LIBS build --define=PREFIX=/usr build --define=LIBDIR=$(PREFIX)/lib build --define=INCLUDEDIR=$(PREFIX)/include -# Disable MKL-DNN contraction kernels by default. -build --define=tensorflow_mkldnn_contraction_kernel=0 - # Default options should come above this line # Options from ./configure diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 4a296f265f7b9521c46d350cec26ff199f43eb6c..b978f89f9e1d79dd4f7481711a59c2b94e8bf01b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -150,41 +150,45 @@ may exist in your changes. There are two ways to run TensorFlow unit tests. -1. Using tools and libraries installed directly on your system. +1. Using tools and libraries installed directly on your system. - Refer to the - [CPU-only developer Dockerfile](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/docker/Dockerfile.devel) and - [GPU developer Dockerfile](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/docker/Dockerfile.devel-gpu) - for the required packages. Alternatively, use the said - [Docker images](https://hub.docker.com/r/tensorflow/tensorflow/tags/), e.g., - `tensorflow/tensorflow:nightly-devel` and `tensorflow/tensorflow:nightly-devel-gpu` - for development to avoid installing the packages directly on your system. + Refer to the + [CPU-only developer Dockerfile](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/docker/Dockerfile.devel) + and + [GPU developer Dockerfile](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/docker/Dockerfile.devel-gpu) + for the required packages. Alternatively, use the said + [Docker images](https://hub.docker.com/r/tensorflow/tensorflow/tags/), e.g., + `tensorflow/tensorflow:nightly-devel` and + `tensorflow/tensorflow:nightly-devel-gpu` for development to avoid + installing the packages directly on your system (in which case remember to + change directory from `/root` to `/tensorflow` once you get into the running + container so `bazel` can find the `tensorflow` workspace). - Once you have the packages installed, you can run a specific unit test in - bazel by doing as follows: + Once you have the packages installed, you can run a specific unit test in + bazel by doing as follows: - If the tests are to be run on GPU, add CUDA paths to LD_LIBRARY_PATH and add - the `cuda` option flag + If the tests are to be run on GPU, add CUDA paths to LD_LIBRARY_PATH and add + the `cuda` option flag - ```bash - export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH" + ```bash + export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH" - export flags="--config=opt --config=cuda -k" - ``` + export flags="--config=opt --config=cuda -k" + ``` - For example, to run all tests under tensorflow/python, do: + For example, to run all tests under tensorflow/python, do: - ```bash - bazel test ${flags} //tensorflow/python/... - ``` + ```bash + bazel test ${flags} //tensorflow/python/... + ``` -2. Using [Docker](https://www.docker.com) and TensorFlow's CI scripts. +2. Using [Docker](https://www.docker.com) and TensorFlow's CI scripts. - ```bash - # Install Docker first, then this will build and run cpu tests - tensorflow/tools/ci_build/ci_build.sh CPU bazel test //tensorflow/... - ``` - - See - [TensorFlow Builds](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/tools/ci_build) for details. + ```bash + # Install Docker first, then this will build and run cpu tests + tensorflow/tools/ci_build/ci_build.sh CPU bazel test //tensorflow/... + ``` + See + [TensorFlow Builds](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/tools/ci_build) + for details. diff --git a/README.md b/README.md index 519815d006cc33be10132909baf414a4bd843435..96a8ecf4f693d5634da63f4ecc6f4e9c35751f5b 100644 --- a/README.md +++ b/README.md @@ -24,7 +24,8 @@ organization for the purposes of conducting machine learning and deep neural networks research. The system is general enough to be applicable in a wide variety of other domains, as well. -TensorFlow provides stable Python API and C APIs as well as without API backwards compatibility guarantee like C++, Go, Java, JavaScript and Swift. +TensorFlow provides stable Python and C APIs as well as non-guaranteed backwards +compatible API's for C++, Go, Java, JavaScript and Swift. Keep up to date with release announcements and security updates by subscribing to @@ -57,21 +58,24 @@ Simply run `pip install tf-nightly` or `pip install tf-nightly-gpu` in a clean environment to install the nightly TensorFlow build. We support CPU and GPU packages on Linux, Mac, and Windows. - #### *Try your first TensorFlow program* + ```shell $ python ``` + ```python >>> import tensorflow as tf >>> tf.enable_eager_execution() ->>> tf.add(1, 2) +>>> tf.add(1, 2).numpy() 3 >>> hello = tf.constant('Hello, TensorFlow!') >>> hello.numpy() 'Hello, TensorFlow!' ``` -Learn more examples about how to do specific tasks in TensorFlow at the [tutorials page of tensorflow.org](https://www.tensorflow.org/tutorials/). + +Learn more examples about how to do specific tasks in TensorFlow at the +[tutorials page of tensorflow.org](https://www.tensorflow.org/tutorials/). ## Contribution guidelines diff --git a/WORKSPACE b/WORKSPACE index 2277e83a3f67b62cf4ee1311767ee06c0549c697..9f07b9fd47136d058cc4039ed6948db539485039 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -4,11 +4,11 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive", "http_file" http_archive( name = "io_bazel_rules_closure", - sha256 = "a38539c5b5c358548e75b44141b4ab637bba7c4dc02b46b1f62a96d6433f56ae", - strip_prefix = "rules_closure-dbb96841cc0a5fb2664c37822803b06dab20c7d1", + sha256 = "43c9b882fa921923bcba764453f4058d102bece35a37c9f6383c713004aacff1", + strip_prefix = "rules_closure-9889e2348259a5aad7e805547c1a0cf311cfcd91", urls = [ - "https://mirror.bazel.build/github.com/bazelbuild/rules_closure/archive/dbb96841cc0a5fb2664c37822803b06dab20c7d1.tar.gz", - "https://github.com/bazelbuild/rules_closure/archive/dbb96841cc0a5fb2664c37822803b06dab20c7d1.tar.gz", # 2018-04-13 + "https://mirror.bazel.build/github.com/bazelbuild/rules_closure/archive/9889e2348259a5aad7e805547c1a0cf311cfcd91.tar.gz", + "https://github.com/bazelbuild/rules_closure/archive/9889e2348259a5aad7e805547c1a0cf311cfcd91.tar.gz", # 2018-12-21 ], ) @@ -29,7 +29,7 @@ load( bazel_toolchains_repositories() load( - "@io_bazel_rules_docker//container:container.bzl", + "@io_bazel_rules_docker//repositories:repositories.bzl", container_repositories = "repositories", ) @@ -43,29 +43,17 @@ remote_config_workspace() # Apple and Swift rules. http_archive( name = "build_bazel_rules_apple", - sha256 = "4fe4ee824200b48821730f89ff260984332dc3551db587c24691235d1d96a8a7", - strip_prefix = "rules_apple-0.10.0", - urls = ["https://github.com/bazelbuild/rules_apple/archive/0.10.0.tar.gz"], -) -http_archive( - name = "build_bazel_rules_swift", - sha256 = "6544ff5615febec0342de1127144d2f3e43ea80fb7f9b1ade65e6a184e39e618", - strip_prefix = "rules_swift-0.5.0", - urls = ["https://github.com/bazelbuild/rules_swift/archive/0.5.0.tar.gz"], -) -http_archive( - name = "bazel_skylib", - sha256 = "eb5c57e4c12e68c0c20bc774bfbc60a568e800d025557bc4ea022c6479acc867", - strip_prefix = "bazel-skylib-0.6.0", - urls = ["https://github.com/bazelbuild/bazel-skylib/archive/0.6.0.tar.gz"], + sha256 = "73b4980a318d203d3307f850e27e66ec5cc8d223147a3475a6f11597eb6438a5", + strip_prefix = "rules_apple-0.13.0", + urls = ["https://github.com/bazelbuild/rules_apple/archive/0.13.0.tar.gz"], ) http_file( name = "xctestrunner", executable = 1, - urls = ["https://github.com/google/xctestrunner/releases/download/0.2.5/ios_test_runner.par"], + urls = ["https://github.com/google/xctestrunner/releases/download/0.2.6/ios_test_runner.par"], ) load("@build_bazel_rules_apple//apple:repositories.bzl", "apple_rules_dependencies") -apple_rules_dependencies(ignore_version_differences = True) +apple_rules_dependencies() load("@build_bazel_rules_swift//swift:repositories.bzl", "swift_rules_dependencies") swift_rules_dependencies() @@ -73,7 +61,7 @@ swift_rules_dependencies() # files, in case the parsing of those build files depends on the bazel # version we require here. load("//tensorflow:version_check.bzl", "check_bazel_version_at_least") -check_bazel_version_at_least("0.18.0") +check_bazel_version_at_least("0.19.0") load("//tensorflow:workspace.bzl", "tf_workspace") @@ -134,4 +122,3 @@ http_archive( "http://download.tensorflow.org/models/speech_commands_v0.01.zip", ], ) - diff --git a/configure.py b/configure.py index bb59063f7932dd1f74e12a39db029a86268a7c87..14fca1f73236eb01ec4bc24499544453fb0807f8 100644 --- a/configure.py +++ b/configure.py @@ -256,6 +256,7 @@ def reset_tf_configure_bazelrc(): """Reset file that contains customized config settings.""" open(_TF_BAZELRC, 'w').close() + def cleanup_makefile(): """Delete any leftover BUILD files from the Makefile build. @@ -785,8 +786,7 @@ def set_gcc_host_compiler_path(environ_cp): environ_cp, var_name='GCC_HOST_COMPILER_PATH', var_default=default_gcc_host_compiler_path, - ask_for_var= - 'Please specify which gcc should be used by nvcc as the host compiler.', + ask_for_var='Please specify which gcc should be used by nvcc as the host compiler.', check_success=os.path.exists, error_msg='Invalid gcc path. %s cannot be found.', ) @@ -1237,6 +1237,7 @@ def set_tf_nccl_install_path(environ_cp): environ_cp['TF_NCCL_VERSION'] = tf_nccl_version write_action_env_to_bazelrc('TF_NCCL_VERSION', tf_nccl_version) + def get_native_cuda_compute_capabilities(environ_cp): """Get native cuda compute capabilities. @@ -1482,7 +1483,7 @@ def set_other_mpi_vars(environ_cp): else: raise ValueError( 'Cannot find the MPI library file in %s/lib or %s/lib64 or %s/lib32' % - mpi_home, mpi_home, mpi_home) + (mpi_home, mpi_home, mpi_home)) def set_system_libs_flag(environ_cp): @@ -1513,10 +1514,6 @@ def set_windows_build_flags(environ_cp): # The host and target platforms are the same in Windows build. So we don't # have to distinct them. This avoids building the same targets twice. write_to_bazelrc('build --distinct_host_configuration=false') - # Enable short object file path to avoid long path issue on Windows. - # TODO(pcloudy): Remove this flag when upgrading Bazel to 0.16.0 - # Short object file path will be enabled by default. - write_to_bazelrc('build --experimental_shortened_obj_file_path=true') if get_var( environ_cp, 'TF_OVERRIDE_EIGEN_STRONG_INLINE', 'Eigen strong inline', @@ -1556,7 +1553,7 @@ def main(): # environment variables. environ_cp = dict(os.environ) - check_bazel_version('0.19.0', '0.21.0') + check_bazel_version('0.19.0', '0.22.0') reset_tf_configure_bazelrc() @@ -1687,14 +1684,15 @@ 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('dynamic_kernels', - '(Experimental) Build kernels into separate shared objects.') + config_info_line( + 'dynamic_kernels', + '(Experimental) Build kernels into separate shared objects.') print('Preconfigured Bazel build configs to DISABLE default on features:') config_info_line('noaws', 'Disable AWS S3 filesystem support.') config_info_line('nogcp', 'Disable GCP support.') config_info_line('nohdfs', 'Disable HDFS support.') - config_info_line('noignite', 'Disable Apacha Ignite support.') + config_info_line('noignite', 'Disable Apache Ignite support.') config_info_line('nokafka', 'Disable Apache Kafka support.') config_info_line('nonccl', 'Disable NVIDIA NCCL support.') diff --git a/tensorflow/BUILD b/tensorflow/BUILD index 0b4498ef1257e7a106ece88c7bd4e9440d85557a..0b63ee4056c57a58aa22560bea63dd3fac623602 100644 --- a/tensorflow/BUILD +++ b/tensorflow/BUILD @@ -40,12 +40,16 @@ load( # @unused TENSORFLOW_API_INIT_FILES_V2 = ( - TENSORFLOW_API_INIT_FILES + get_compat_files(TENSORFLOW_API_INIT_FILES_V1, 1) + TENSORFLOW_API_INIT_FILES + + get_compat_files(TENSORFLOW_API_INIT_FILES, 2) + + get_compat_files(TENSORFLOW_API_INIT_FILES_V1, 1) ) # @unused -TENSORFLOW_API_INIT_FILES_V1_WITH_COMPAT = ( - TENSORFLOW_API_INIT_FILES_V1 + get_compat_files(TENSORFLOW_API_INIT_FILES_V1, 1) +TENSORFLOW_API_INIT_FILES_V1 = ( + TENSORFLOW_API_INIT_FILES_V1 + + get_compat_files(TENSORFLOW_API_INIT_FILES, 2) + + get_compat_files(TENSORFLOW_API_INIT_FILES_V1, 1) ) # Config setting used when building for products @@ -90,6 +94,12 @@ config_setting( visibility = ["//visibility:public"], ) +config_setting( + name = "emscripten", + values = {"crosstool_top": "//external:android/emscripten"}, + visibility = ["//visibility:public"], +) + config_setting( name = "raspberry_pi_armeabi", values = { @@ -343,6 +353,13 @@ config_setting( }, ) +config_setting( + name = "using_rocm_hipcc", + define_values = { + "using_rocm_hipcc": "true", + }, +) + config_setting( name = "with_mpi_support", values = {"define": "with_mpi_support=true"}, @@ -381,17 +398,7 @@ config_setting( package_group( name = "internal", - packages = [ - "-//third_party/tensorflow/python/estimator", - "//learning/deepmind/...", - "//learning/meta_rank/...", - "//learning/pathways/...", # While dataset C++ api requires internals - "//tensorflow/...", - "//tensorflow_estimator/contrib/...", - "//tensorflow_fold/llgtm/...", - "//tensorflow_text/...", - "//third_party/py/tensor2tensor/...", - ], + packages = ["//tensorflow/..."], ) load( @@ -600,13 +607,20 @@ gen_api_init_files( name = "tf_python_api_gen_v1", srcs = [ "api_template_v1.__init__.py", + "compat_template.__init__.py", "compat_template_v1.__init__.py", ], api_version = 1, - compat_api_versions = [1], - compat_init_templates = ["compat_template_v1.__init__.py"], + compat_api_versions = [ + 1, + 2, + ], + compat_init_templates = [ + "compat_template_v1.__init__.py", + "compat_template.__init__.py", + ], output_dir = "_api/v1/", - output_files = TENSORFLOW_API_INIT_FILES_V1_WITH_COMPAT, + output_files = TENSORFLOW_API_INIT_FILES_V1, output_package = "tensorflow._api.v1", root_file_name = "v1.py", root_init_template = "api_template_v1.__init__.py", @@ -616,11 +630,18 @@ gen_api_init_files( name = "tf_python_api_gen_v2", srcs = [ "api_template.__init__.py", + "compat_template.__init__.py", "compat_template_v1.__init__.py", ], api_version = 2, - compat_api_versions = [1], - compat_init_templates = ["compat_template_v1.__init__.py"], + compat_api_versions = [ + 1, + 2, + ], + compat_init_templates = [ + "compat_template_v1.__init__.py", + "compat_template.__init__.py", + ], output_dir = "_api/v2/", output_files = TENSORFLOW_API_INIT_FILES_V2, output_package = "tensorflow._api.v2", diff --git a/tensorflow/api_template.__init__.py b/tensorflow/api_template.__init__.py index a93799bfe84b0f9c4743e1ad0effd6e69ad7f3f2..a6eb4755f32d2504ae1aab747f110d68d72a0d5f 100644 --- a/tensorflow/api_template.__init__.py +++ b/tensorflow/api_template.__init__.py @@ -111,5 +111,10 @@ try: except NameError: pass +# Add module aliases +if hasattr(_current_module, 'keras'): + losses = keras.losses + metrics = keras.metrics + optimizers = keras.optimizers # pylint: enable=undefined-variable diff --git a/tensorflow/api_template_v1.__init__.py b/tensorflow/api_template_v1.__init__.py index 514aba1b59631f882523396aab0f4d3d5e88a893..eeca8f0d566a6401cb64e4fe3f0ee3c5aeb4ece2 100644 --- a/tensorflow/api_template_v1.__init__.py +++ b/tensorflow/api_template_v1.__init__.py @@ -62,11 +62,15 @@ if '__all__' in vars(): vars()['__all__'].append('contrib') from tensorflow.python.platform import flags # pylint: disable=g-import-not-at-top +# The 'app' module will be imported as part of the placeholder section above. app.flags = flags # pylint: disable=undefined-variable +# Also use 'app' module (choice is arbitrary) to derive the API directory below. +_API_MODULE = app # pylint: disable=undefined-variable + # Make sure directory containing top level submodules is in # the __path__ so that "from tensorflow.foo import bar" works. -_tf_api_dir = _os.path.dirname(_os.path.dirname(app.__file__)) # pylint: disable=undefined-variable +_tf_api_dir = _os.path.dirname(_os.path.dirname(_API_MODULE.__file__)) # pylint: disable=undefined-variable if not hasattr(_current_module, '__path__'): __path__ = [_tf_api_dir] elif _tf_api_dir not in __path__: diff --git a/tensorflow/c/BUILD b/tensorflow/c/BUILD index 9d267e9e596efc1cb3353039f6702906b21bafad..ef7863dc0d5cbd57da30baa6e04278c2a0354b25 100644 --- a/tensorflow/c/BUILD +++ b/tensorflow/c/BUILD @@ -67,6 +67,23 @@ tf_cuda_library( tf_cuda_library( name = "c_api", + hdrs = ["c_api.h"], + copts = tf_copts(), + visibility = ["//visibility:public"], + deps = [ + ":c_api_no_xla", + ":c_api_internal", + ] + select({ + "//tensorflow:with_xla_support": [ + "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/jit", + ], + "//conditions:default": [], + }), +) + +tf_cuda_library( + name = "c_api_no_xla", srcs = [ "c_api.cc", "c_api_function.cc", @@ -75,15 +92,13 @@ tf_cuda_library( "c_api.h", ], copts = tf_copts(), - visibility = ["//visibility:public"], - deps = select({ + visibility = ["//tensorflow/c:__subpackages__"], + deps = [":c_api_internal"] + select({ "//tensorflow:android": [ - ":c_api_internal", "//tensorflow/core:android_tensorflow_lib_lite", ], "//conditions:default": [ - ":c_api_internal", - "//tensorflow/cc/saved_model:loader", + "//tensorflow/cc/saved_model:loader_lite", "//tensorflow/cc:gradients", "//tensorflow/cc:ops", "//tensorflow/cc:grad_ops", @@ -97,13 +112,8 @@ tf_cuda_library( "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core/distributed_runtime:server_lib", + "//tensorflow/core/kernels:logging_ops", ], - }) + select({ - "//tensorflow:with_xla_support": [ - "//tensorflow/compiler/tf2xla:xla_compiler", - "//tensorflow/compiler/jit", - ], - "//conditions:default": [], }), ) @@ -129,6 +139,7 @@ tf_cuda_library( "//tensorflow/core:lib_platform", "//tensorflow/core:protos_all_cc", "//tensorflow/core/common_runtime/eager:attr_builder", + "@com_google_absl//absl/strings", ], ) @@ -155,8 +166,8 @@ tf_cuda_library( hdrs = ["tf_status_helper.h"], visibility = ["//visibility:public"], deps = [ - ":c_api", ":c_api_internal", + ":c_api_no_xla", "//tensorflow/core:lib", ], ) @@ -212,13 +223,13 @@ tf_cuda_library( visibility = ["//visibility:public"], deps = select({ "//tensorflow:android": [ - ":c_api", + ":c_api_no_xla", ":c_api_internal", ":tf_status_helper", "//tensorflow/core:android_tensorflow_lib_lite", ], "//conditions:default": [ - ":c_api", + ":c_api_no_xla", ":c_api_internal", ":tf_status_helper", "//tensorflow/core:framework", @@ -288,13 +299,23 @@ tf_cuda_cc_test( "//tensorflow/cc/saved_model:signature_constants", "//tensorflow/cc/saved_model:tag_constants", "//tensorflow/compiler/jit", + "//tensorflow/core:array_ops_op_lib", + "//tensorflow/core:bitwise_ops_op_lib", + "//tensorflow/core:control_flow_ops_op_lib", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:direct_session", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", + "//tensorflow/core:functional_ops_op_lib", "//tensorflow/core:lib", + "//tensorflow/core:math_ops_op_lib", + "//tensorflow/core:nn_ops_op_lib", + "//tensorflow/core:no_op_op_lib", "//tensorflow/core:proto_text", "//tensorflow/core:protos_all_cc", + "//tensorflow/core:sendrecv_ops_op_lib", + "//tensorflow/core:spectral_ops_op_lib", + "//tensorflow/core:state_ops_op_lib", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core/kernels:array", @@ -318,6 +339,7 @@ tf_cc_test( deps = [ ":c_api", ":c_api_experimental", + ":c_api_internal", ":c_test_util", "//tensorflow/c/eager:c_api", "//tensorflow/c/eager:c_api_test_util", @@ -334,6 +356,7 @@ tf_cc_test( srcs = ["c_api_function_test.cc"], deps = [ ":c_api", + ":c_api_internal", ":c_test_util", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", diff --git a/tensorflow/c/c_api.cc b/tensorflow/c/c_api.cc index 9580215a317b1a6b1cdacbd430a1764af61be990..ef22b67fe95364a0513ebd7a59d116a2d78cc2e9 100644 --- a/tensorflow/c/c_api.cc +++ b/tensorflow/c/c_api.cc @@ -27,6 +27,7 @@ limitations under the License. #include "tensorflow/cc/ops/while_loop.h" #include "tensorflow/cc/saved_model/loader.h" #include "tensorflow/core/framework/op_gen_lib.h" +#include "tensorflow/core/kernels/logging_ops.h" #endif #include "tensorflow/c/c_api_internal.h" #include "tensorflow/core/common_runtime/device_mgr.h" @@ -257,6 +258,74 @@ int64_t TF_Dim(const TF_Tensor* t, int dim_index) { size_t TF_TensorByteSize(const TF_Tensor* t) { return t->buffer->size(); } void* TF_TensorData(const TF_Tensor* t) { return t->buffer->data(); } +int64_t TF_TensorElementCount(const TF_Tensor* t) { + int64_t result = 1; + int rank = TF_NumDims(t); + for (int dim = 0; dim < rank; ++dim) { + result *= TF_Dim(t, dim); + } + return result; +} + +// Returns the number of elements that would be present in a tensor with the +// given shape. +static int64_t ShapeNumElements(const int64_t* dims, int num_dims) { + int64_t result = 1; + for (int dim = 0; dim < num_dims; ++dim) { + result *= dims[dim]; + } + return result; +} + +static void UnrefIfNonNull(::tensorflow::TensorBuffer* buf) { + if (buf != nullptr) { + buf->Unref(); + } +} + +static void RefIfNonNull(::tensorflow::TensorBuffer* buf) { + if (buf != nullptr) { + buf->Ref(); + } +} + +void TF_TensorBitcastFrom(const TF_Tensor* from, TF_DataType type, + TF_Tensor* to, const int64_t* new_dims, + int num_new_dims, TF_Status* status) { + TF_SetStatus(status, TF_OK, ""); + size_t in_size = TF_DataTypeSize(TF_TensorType(from)); + if (in_size == 0) { + TF_SetStatus(status, TF_INVALID_ARGUMENT, + "input tensor has a zero-sized data type"); + return; + } + size_t out_size = TF_DataTypeSize(type); + if (out_size == 0) { + TF_SetStatus(status, TF_INVALID_ARGUMENT, + "output tensor has a zero-sized data type"); + return; + } + + if (ShapeNumElements(new_dims, num_new_dims) * out_size != + TF_TensorElementCount(from) * in_size) { + TF_SetStatus(status, TF_INVALID_ARGUMENT, + "input tensor is not compatible with output shape"); + return; + } + + tensorflow::TensorShapeProto p; + for (int i = 0; i < num_new_dims; ++i) { + p.add_dim()->set_size(new_dims[i]); + } + to->shape = tensorflow::TensorShape(p); + to->dtype = type; + if (to->buffer != from->buffer) { + UnrefIfNonNull(to->buffer); + to->buffer = from->buffer; + RefIfNonNull(to->buffer); + } +} + // -------------------------------------------------------------------------- size_t TF_StringEncode(const char* src, size_t src_len, char* dst, size_t dst_len, TF_Status* status) { @@ -1242,6 +1311,13 @@ void TF_SetAttrTypeList(TF_OperationDescription* desc, const char* attr_name, reinterpret_cast(values), num_values)); } +void TF_SetAttrPlaceholder(TF_OperationDescription* desc, const char* attr_name, + const char* placeholder) { + tensorflow::AttrValue attr_value; + attr_value.set_placeholder(placeholder); + desc->node_builder.Attr(attr_name, attr_value); +} + void TF_SetAttrFuncName(TF_OperationDescription* desc, const char* attr_name, const char* value, size_t length) { tensorflow::NameAttrList func_name; @@ -2881,6 +2957,16 @@ const char* TF_ServerTarget(TF_Server* server) { #endif } -void TF_DeleteServer(TF_Server* server) { delete server; } +void TF_DeleteServer(TF_Server* server) { +#ifndef __ANDROID__ + delete server; +#endif +} + +void TF_RegisterLogListener(void (*listener)(const char*)) { +#ifndef __ANDROID__ + tensorflow::logging::RegisterListener(listener); +#endif +} } // end extern "C" diff --git a/tensorflow/c/c_api.h b/tensorflow/c/c_api.h index c7abba85521fccec07983cd5ab4f94a8368d6181..88b8b49b016efec4eb21271d275e50c786e1e602 100644 --- a/tensorflow/c/c_api.h +++ b/tensorflow/c/c_api.h @@ -272,6 +272,39 @@ TF_CAPI_EXPORT extern size_t TF_TensorByteSize(const TF_Tensor*); // Return a pointer to the underlying data buffer. TF_CAPI_EXPORT extern void* TF_TensorData(const TF_Tensor*); +// Returns the number of elements in the tensor. +TF_CAPI_EXPORT extern int64_t TF_TensorElementCount(const TF_Tensor* tensor); + +// Copy the internal data representation of `from` to `to`. `new_dims` and +// `num_new_dims` specify the new shape of the `to` tensor, `type` specifies its +// data type. On success, *status is set to TF_OK and the two tensors share the +// same data buffer. +// +// This call requires that the `from` tensor and the given type and shape (dims +// and num_dims) are "compatible" (i.e. they occupy the same number of bytes). +// Specifically, given from_type_size = TF_DataTypeSize(TF_TensorType(from)): +// +// ShapeElementCount(dims, num_dims) * TF_DataTypeSize(type) +// +// must equal +// +// TF_TensorElementCount(from) * from_type_size +// +// where TF_ShapeElementCount would be the number of elements in a tensor with +// the given shape. +// +// In addition, this function requires: +// * TF_DataTypeSize(TF_TensorType(from)) != 0 +// * TF_DataTypeSize(type) != 0 +// +// If any of the requirements are not met, *status is set to +// TF_INVALID_ARGUMENT. +TF_CAPI_EXPORT extern void TF_TensorBitcastFrom(const TF_Tensor* from, + TF_DataType type, TF_Tensor* to, + const int64_t* new_dims, + int num_new_dims, + TF_Status* status); + // -------------------------------------------------------------------------- // Encode the string `src` (`src_len` bytes long) into `dst` in the format // required by TF_STRING tensors. Does not write to memory more than `dst_len` @@ -516,6 +549,10 @@ TF_CAPI_EXPORT extern void TF_SetAttrTypeList(TF_OperationDescription* desc, const char* attr_name, const TF_DataType* values, int num_values); +TF_CAPI_EXPORT extern void TF_SetAttrPlaceholder(TF_OperationDescription* desc, + const char* attr_name, + const char* placeholder); + // Set a 'func' attribute to the specified name. // `value` must point to a string of length `length` bytes. TF_CAPI_EXPORT extern void TF_SetAttrFuncName(TF_OperationDescription* desc, @@ -1710,6 +1747,14 @@ TF_CAPI_EXPORT extern const char* TF_ServerTarget(TF_Server* server); // it will be stopped and joined. TF_CAPI_EXPORT extern void TF_DeleteServer(TF_Server* server); +// Register a listener method that processes printed messages. +// +// If any listeners are registered, the print operator will call all listeners +// with the printed messages and immediately return without writing to the +// logs. +TF_CAPI_EXPORT extern void TF_RegisterLogListener( + void (*listener)(const char*)); + #ifdef __cplusplus } /* end extern "C" */ #endif diff --git a/tensorflow/c/c_api_experimental.cc b/tensorflow/c/c_api_experimental.cc index f04b285037dff403428ed74fe90eac60339fe36b..a8325ce494c4f57fcd7e64b2d233ee4e6666bc4e 100644 --- a/tensorflow/c/c_api_experimental.cc +++ b/tensorflow/c/c_api_experimental.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/c/c_api_experimental.h" +#include "absl/strings/substitute.h" #include "tensorflow/c/c_api.h" #include "tensorflow/c/c_api_internal.h" #include "tensorflow/c/eager/c_api.h" @@ -128,6 +129,14 @@ const char* TF_GraphDebugString(TF_Graph* graph, size_t* len) { return ret; } +char* TF_FunctionDebugString(TF_Function* func, size_t* len) { + const auto& debug_str = func->fdef.DebugString(); + *len = debug_str.size(); + char* ret = static_cast(malloc(*len + 1)); + memcpy(ret, debug_str.c_str(), *len + 1); + return ret; +} + // On success, returns a set of TF_Function instances from `text_proto` of // GraphDef type. These functions must be deleted by calling TF_DeleteFunction. // @@ -8737,6 +8746,12 @@ static void CheckOk(TF_Status* status) { void TFE_TensorHandlePrintDebugString(TFE_TensorHandle* handle) { auto* status = TF_NewStatus(); + if (!TFE_TensorHandleIsConcrete(handle)) { + VLOG(1) << "Symbolic tensor: " << handle; + TF_DeleteStatus(status); + return; + } + TF_Tensor* t = TFE_TensorHandleResolve(handle, status); CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); @@ -8748,6 +8763,11 @@ void TFE_TensorHandlePrintDebugString(TFE_TensorHandle* handle) { TF_DeleteStatus(status); } +void TFE_OpPrintDebugString(TFE_Op* op) { + VLOG(1) << "TFE_OpPrintDebugString() over " << op; + LOG(INFO) << op->operation.DebugString(); +} + struct TFE_ExecuteOpNotification { TFE_ExecuteOpNotification() : status(TF_NewStatus(), TF_DeleteStatus) {} tensorflow::Notification n; @@ -8941,3 +8961,166 @@ TF_CAPI_EXPORT extern void TFE_EnableCollectiveOps(TFE_Context* ctx, } status->status = EnableCollectiveOps(server_def, ctx); } + +std::string tensorflow::getTF_OutputDebugString(TF_Output node) { + return absl::Substitute("TF_Output($0, $1)", node.oper, node.index); +} + +using tensorflow::getTF_OutputDebugString; + +TFE_TensorHandle* TFE_NewTensorHandleFromTFOutput(TF_Output t, + TF_DataType dtype) { + auto ret = new TFE_TensorHandle(t, dtype); + VLOG(1) << "Storing TFOutput " << getTF_OutputDebugString(t) + << " into tensor handle " << ret << " with internal handle " + << ret->handle; + return ret; +} + +unsigned char TFE_TensorHandleIsConcrete(TFE_TensorHandle* handle) { + assert(handle->handle != nullptr); + return handle->handle->getSymbolicTensor() == nullptr; +} + +TF_Output TFE_GetTFOutputFromTensorHandle(TFE_TensorHandle* handle, + TF_Status* status) { + if (TFE_TensorHandleIsConcrete(handle)) { + status->status = + tensorflow::errors::Internal("Not a symbolic tensor: ", handle); + return TF_Output{nullptr, -1}; + } + + auto* sym_tensor = handle->handle->getSymbolicTensor(); + CHECK(sym_tensor != nullptr); + auto ret = TF_Output{sym_tensor->oper, sym_tensor->index}; + VLOG(1) << "Retrieving " << getTF_OutputDebugString(ret) + << " from tensor handle " << handle; + CHECK_GE(sym_tensor->index, 0); + return ret; +} + +TFE_TraceContext* TFE_NewTraceContext(TF_Graph* graph) { + return new TFE_TraceContext(graph); +} + +void TFE_DeleteTraceContext(TFE_TraceContext* trace_ctx) { delete trace_ctx; } + +// If `handle` is already symbolic, return it. Otherwise map it to a new +// symbolic tensor (a PlaceHolder op) and return that. +static TF_Output getOrCreateSymbolicTensor(TFE_TraceContext* trace_ctx, + tensorflow::TensorHandle* handle, + TF_Status* status) { + VLOG(1) << "Getting symbolic tensor for input tensor handle " << handle + << ": " << handle->DebugString(); + + auto* sym_tensor = handle->getSymbolicTensor(); + if (sym_tensor != nullptr) { + auto ret = TF_Output{sym_tensor->oper, sym_tensor->index}; + VLOG(1) << "This handle is a symbolic tensor " << sym_tensor << ": " + << getTF_OutputDebugString(ret); + return ret; + } + + auto find_it = trace_ctx->input_tensor_map.find(handle); + if (find_it != trace_ctx->input_tensor_map.end()) { + VLOG(1) << "There exists a map entry from this concrete tensor to: " + << getTF_OutputDebugString(find_it->second); + return find_it->second; + } + + auto node_name = tensorflow::strings::StrCat("additional_input_", + trace_ctx->node_counter++); + VLOG(1) << "Adding a place holder node named " << node_name; + auto* desc = + TF_NewOperation(trace_ctx->graph, "Placeholder", node_name.c_str()); + TF_SetAttrType(desc, "dtype", + static_cast(handle->dtype) /*TF_FLOAT*/); + auto* result = TF_FinishOperation(desc, status); + if (!status->status.ok()) { + return TF_Output{nullptr, -1}; + } + + auto ret = TF_Output{result, 0}; + VLOG(1) << "Creating a new map entry to map to: " + << getTF_OutputDebugString(ret); + trace_ctx->input_tensor_map[handle] = ret; + // `handle` could be destroyed before it's read from `input_tensor_map` (say + // during a subsequent TFE_FinalizeInputTensorsFromTraceContext() call), so we + // increment its ref count to extend its life span to that of `trace_ctx`. + handle->Ref(); + VLOG(1) << "Ref count for handle " << handle + << " is 1?: " << handle->RefCountIsOne(); + return ret; +} + +TF_Operation* TFE_AddEagerOpToGraph(TFE_Op* op, TFE_TraceContext* trace_ctx, + TFE_TensorHandle** retvals, + int* num_retvals, TF_Status* status) { + VLOG(1) << "Calling TFE_AddEagerOpToGraph() with op " << op << ": " + << op->operation.DebugString(); + + const auto& op_type = op->operation.Name(); + auto op_name = + tensorflow::strings::StrCat(op_type, "_", trace_ctx->node_counter++); + auto* desc = + TF_NewOperation(trace_ctx->graph, op_type.c_str(), op_name.c_str()); + for (auto* input : op->operation.Inputs()) { + auto symbolic_input = getOrCreateSymbolicTensor(trace_ctx, input, status); + if (!status->status.ok()) return nullptr; + TF_AddInput(desc, symbolic_input); + } + + VLOG(1) << "Adding attrs."; + tensorflow::AttrValueMap attrs; + op->operation.Attrs().FillAttrValueMap(&attrs); + for (const auto& attr : attrs) { + desc->node_builder.Attr(attr.first, attr.second); + } + + auto* graph_op = TF_FinishOperation(desc, status); + if (!status->status.ok()) return nullptr; + + VLOG(1) << "Op finalized; setting return tensors."; + *num_retvals = TF_OperationNumOutputs(graph_op); + VLOG(1) << "This op has " << *num_retvals << " outputs."; + for (int i = 0; i < *num_retvals; ++i) { + auto output = TF_Output{graph_op, i}; + auto dtype = TF_OperationOutputType(output); + retvals[i] = TFE_NewTensorHandleFromTFOutput(output, dtype); + } + return graph_op; +} + +int TFE_FinalizeInputTensorsFromTraceContext(TFE_TraceContext* trace_ctx) { + if (trace_ctx->input_tensors == nullptr) { + trace_ctx->input_tensors = + new std::vector>(); + trace_ctx->input_tensors->reserve(trace_ctx->input_tensor_map.size()); + + for (auto input : trace_ctx->input_tensor_map) { + trace_ctx->input_tensors->emplace_back(input.first, input.second); + } + } + return trace_ctx->input_tensor_map.size(); +} + +TF_Output TFE_GetInputGraphNodeFromTraceContext(TFE_TraceContext* trace_ctx, + unsigned int idx) { + CHECK(trace_ctx->input_tensors != nullptr); + CHECK(trace_ctx->input_tensors->size() > idx); + return trace_ctx->input_tensors->at(idx).second; +} + +TFE_TensorHandle* TFE_ConsumeInputConcreteTensorFromTraceContext( + TFE_TraceContext* trace_ctx, unsigned int idx) { + CHECK(trace_ctx->input_tensors != nullptr); + CHECK(trace_ctx->input_tensors->size() > idx); + auto* handle = trace_ctx->input_tensors->at(idx).first; + VLOG(1) << "Ref count for internal handle " << handle + << " is 1?: " << handle->RefCountIsOne(); + handle->Ref(); + auto* ret = new TFE_TensorHandle(handle); + VLOG(1) << "Returning a new tensor handle " << ret << ": " + << handle->DebugString(); + return ret; +} diff --git a/tensorflow/c/c_api_experimental.h b/tensorflow/c/c_api_experimental.h index e6d04d0c2b25a3f7b1ebf50c58268f003595a520..8d1a8b82fbaf9901b6d9aecf6d092ae298c8dba3 100644 --- a/tensorflow/c/c_api_experimental.h +++ b/tensorflow/c/c_api_experimental.h @@ -84,6 +84,15 @@ TF_CAPI_EXPORT extern TF_Buffer* TF_CreateRunOptions( TF_CAPI_EXPORT extern const char* TF_GraphDebugString(TF_Graph* graph, size_t* len); +// Returns the function content in a human-readable format, with length set in +// `len`. The format is subject to change in the future. +// The returned string is heap-allocated, and caller should call free() on it. +// +// Do not return const char*, because some foreign language binding +// (e.g. swift) cannot then call free() on the returned pointer. +TF_CAPI_EXPORT extern char* TF_FunctionDebugString(TF_Function* func, + size_t* len); + // Creates a stack of data set + iterator nodes, currently hard-coded to return // a sequence of 3 float values <42.0, 43.0, 44.0> over 3 calls. On success, // returns the IteratorGetNext node, which caller can run or feed into an node. @@ -181,6 +190,8 @@ TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_DequeueVariantTensor( TF_CAPI_EXPORT extern void TFE_TensorHandlePrintDebugString( TFE_TensorHandle* handle); +TF_CAPI_EXPORT extern void TFE_OpPrintDebugString(TFE_Op* op); + typedef struct TFE_ExecuteOpNotification TFE_ExecuteOpNotification; // Allows invoking a kernel asynchronously, and explicitly returns a @@ -255,6 +266,54 @@ TF_CAPI_EXPORT extern void TFE_EnableCollectiveOps(TFE_Context* ctx, const void* proto, size_t proto_len, TF_Status* status); + +// Create a symbolic tensor from the input graph node. +TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_NewTensorHandleFromTFOutput( + TF_Output t, TF_DataType data_type); + +// Returns 0 if the input tensor handle represents a symbolic tensor (i.e., a +// graph node). Otherwise returns non-0. +TF_CAPI_EXPORT extern unsigned char TFE_TensorHandleIsConcrete( + TFE_TensorHandle* handle); + +// If `handle` is a symbolic tensor, return the corresponding graph node +// represented by TF_Output. Otherwise, return an error status. +TF_CAPI_EXPORT extern TF_Output TFE_GetTFOutputFromTensorHandle( + TFE_TensorHandle* handle, TF_Status* status); + +typedef struct TFE_TraceContext TFE_TraceContext; + +// A trace context contains a trace graph, to which TFE_AddEagerOpToGraph() +// calls add graph nodes as a way to symbolically execute the eager ops. +// +// It also contains a hash map from concrete input tensors to symbolic +// tensors. That map will be used to create input tensors to the trace graph. +TF_CAPI_EXPORT extern TFE_TraceContext* TFE_NewTraceContext(TF_Graph* graph); + +TF_CAPI_EXPORT extern void TFE_DeleteTraceContext(TFE_TraceContext* trace_ctx); + +// Symbolically executes `op`, by adding a corresponding node to the graph +// associated with `trace_ctx`. This graph node outputs a set of symbolic +// tensors in `retvals` and `num_retvals`. Returns the corresponding graph +// operation on success, otherwise returns nullptr. +TF_CAPI_EXPORT extern TF_Operation* TFE_AddEagerOpToGraph( + TFE_Op* op, TFE_TraceContext* trace_ctx, TFE_TensorHandle** retvals, + int* num_retvals, TF_Status* status); + +// Finalizes the trace graph and its inputs, and returns the number of inputs. +// After this call, the next two APIs can be called to iterate over the input +// tensors. +TF_CAPI_EXPORT extern int TFE_FinalizeInputTensorsFromTraceContext( + TFE_TraceContext* trace_ctx); + +TF_CAPI_EXPORT extern TF_Output TFE_GetInputGraphNodeFromTraceContext( + TFE_TraceContext* trace_ctx, unsigned int idx); + +// Each input tensor should be consumed at most once. +TF_CAPI_EXPORT extern TFE_TensorHandle* +TFE_ConsumeInputConcreteTensorFromTraceContext(TFE_TraceContext* trace_ctx, + unsigned int idx); + #ifdef __cplusplus } /* end extern "C" */ #endif diff --git a/tensorflow/c/c_api_experimental_test.cc b/tensorflow/c/c_api_experimental_test.cc index daa7701b7fe7e8ce757b6504329cf6434ad39778..354ee5f49f373edbc10e7706aa8776f3cc2a17cd 100644 --- a/tensorflow/c/c_api_experimental_test.cc +++ b/tensorflow/c/c_api_experimental_test.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/c/c_api_experimental.h" +#include "tensorflow/c/c_api_internal.h" #include "tensorflow/c/c_test_util.h" #include "tensorflow/c/eager/c_api.h" #include "tensorflow/c/eager/c_api_test_util.h" @@ -296,5 +297,154 @@ TEST(CAPI_EXPERIMENTAL, TFE_ExecuteOpInNewThreadTest_Blocking) { TF_DeleteStatus(status); } +TEST(CAPI_EXPERIMENTAL, SymbolicTensor) { + TF_Status* status = TF_NewStatus(); + auto node = TF_Output{nullptr, 1}; + auto* sym_handle = TFE_NewTensorHandleFromTFOutput(node, TF_FLOAT); + TFE_TensorHandlePrintDebugString(sym_handle); + CHECK_EQ(TFE_TensorHandleDataType(sym_handle), TF_FLOAT); + ASSERT_FALSE(TFE_TensorHandleIsConcrete(sym_handle)); + auto same_node = TFE_GetTFOutputFromTensorHandle(sym_handle, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + ASSERT_EQ(same_node.oper, node.oper); + ASSERT_EQ(same_node.index, node.index); + TFE_DeleteTensorHandle(sym_handle); + + TFE_TensorHandle* m = TestMatrixTensorHandle(); + ASSERT_TRUE(TFE_TensorHandleIsConcrete(m)); + (void)TFE_GetTFOutputFromTensorHandle(m, status); + CHECK_EQ(TF_INTERNAL, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteTensorHandle(m); + + TF_DeleteStatus(status); +} + +class AddEagerOpToGraphTest : public ::testing::Test { + protected: + AddEagerOpToGraphTest() + : status_(TF_NewStatus()), + eager_ctx_(nullptr), + graph_(TF_NewGraph()), + trace_ctx_(TFE_NewTraceContext(graph_)) { + TFE_ContextOptions* opts = TFE_NewContextOptions(); + CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_); + eager_ctx_ = TFE_NewContext(opts, status_); + CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_); + TFE_DeleteContextOptions(opts); + } + + ~AddEagerOpToGraphTest() override { + TFE_DeleteTraceContext(trace_ctx_); + TF_DeleteGraph(graph_); + TFE_DeleteContext(eager_ctx_); + TF_DeleteStatus(status_); + } + + template + void AddEagerOpToGraphAndCheck(TFE_Op* op, Callable checker) { + TFE_TensorHandle* retvals[5]; + int num_retvals = 5; + // Symbolically execute this op, which adds a graph node to `trace_ctx_`. + TF_Operation* graph_op = + TFE_AddEagerOpToGraph(op, trace_ctx_, retvals, &num_retvals, status_); + CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_); + CHECK_NOTNULL(graph_op); + // Check the expectations. + checker(graph_op); + for (int i = 0; i < num_retvals; ++i) { + TFE_DeleteTensorHandle(retvals[i]); + } + } + + TF_Status* status_; + TFE_Context* eager_ctx_; + TF_Graph* graph_; + TFE_TraceContext* trace_ctx_; +}; + +TEST_F(AddEagerOpToGraphTest, DebugPrintAndSymbolicExecution) { + TFE_TensorHandle* m = TestMatrixTensorHandle(); + TFE_Op* op = MatMulOp(eager_ctx_, m, m); + + CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_); + TFE_OpPrintDebugString(op); + + TFE_TensorHandle* retvals[5]; + int num_retvals = 5; + // Symbolically execute this op, which adds a graph node to `trace_ctx`. + TFE_AddEagerOpToGraph(op, trace_ctx_, retvals, &num_retvals, status_); + CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_); + + int num_inputs = TFE_FinalizeInputTensorsFromTraceContext(trace_ctx_); + CHECK_EQ(num_inputs, 1); + auto input_sym_tensor = TFE_GetInputGraphNodeFromTraceContext(trace_ctx_, + /*idx*/ 0); + + LOG(INFO) << tensorflow::getTF_OutputDebugString(input_sym_tensor); + auto handle = TFE_ConsumeInputConcreteTensorFromTraceContext(trace_ctx_, + /*idx*/ 0); + TFE_TensorHandlePrintDebugString(handle); + TFE_DeleteTensorHandle(handle); + + CHECK_EQ(num_retvals, 1); + CHECK_EQ(TFE_TensorHandleDataType(retvals[0]), TF_FLOAT); + + TFE_DeleteTensorHandle(retvals[0]); + TFE_DeleteTensorHandle(m); + TFE_DeleteOp(op); +} + +TEST_F(AddEagerOpToGraphTest, ValueAttributesArePreserved) { + // Create MinOp + TFE_TensorHandle* axis = TestAxisTensorHandle(); + TFE_Op* op = MinOp(eager_ctx_, axis, axis); + CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_); + + // Check the attributes set by the call to MinOp above. + AddEagerOpToGraphAndCheck(op, [this, &axis](TF_Operation* graph_op) { + unsigned char value; + TF_OperationGetAttrBool(graph_op, "keep_dims", &value, status_); + CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_); + CHECK_EQ(value, 1); + TF_DataType dtype; + TF_OperationGetAttrType(graph_op, "Tidx", &dtype, status_); + CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_); + CHECK_EQ(dtype, TF_INT32); + TF_OperationGetAttrType(graph_op, "T", &dtype, status_); + CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_); + CHECK_EQ(dtype, TFE_TensorHandleDataType(axis)); + }); + TFE_DeleteTensorHandle(axis); + TFE_DeleteOp(op); +} + +TEST_F(AddEagerOpToGraphTest, ListAttributesArePreserved) { + // Create a "Squeeze" operator with list attributes. + TFE_TensorHandle* axis = TestAxisTensorHandle(); + TFE_Op* squeeze = TFE_NewOp(eager_ctx_, "Squeeze", status_); + CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_); + TFE_OpAddInput(squeeze, axis, status_); + TFE_OpSetAttrType(squeeze, "T", TF_INT32); + std::vector boundaries = {1, 2, 3, 4}; + TFE_OpSetAttrIntList(squeeze, "squeeze_dims", boundaries.data(), + boundaries.size()); + // Check attributes are preserved. + AddEagerOpToGraphAndCheck( + squeeze, [this, &boundaries](TF_Operation* squeeze_graph_op) { + TF_DataType dtype; + TF_OperationGetAttrType(squeeze_graph_op, "T", &dtype, status_); + CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_); + CHECK_EQ(dtype, TF_INT32); + std::unique_ptr list(new int64_t[boundaries.size()]); + TF_OperationGetAttrIntList(squeeze_graph_op, "squeeze_dims", list.get(), + boundaries.size(), status_); + CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_); + EXPECT_TRUE(std::equal(list.get(), list.get() + boundaries.size(), + boundaries.begin())); + }); + TFE_DeleteTensorHandle(axis); + TFE_DeleteOp(squeeze); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/c/c_api_function.cc b/tensorflow/c/c_api_function.cc index 28b9f8df9c873ee394eb6a241dd9ac06ba6c8796..45d6c33a1e7053451d1dbadff480cf300ea4abbb 100644 --- a/tensorflow/c/c_api_function.cc +++ b/tensorflow/c/c_api_function.cc @@ -162,6 +162,11 @@ Status FillFunctionBody( const std::vector& body_nodes, const std::unordered_map& tensor_renaming, FunctionDef* fdef) { + std::unordered_set func_attr_names; + for (const auto& func_attr : fdef->signature().attr()) { + func_attr_names.insert(func_attr.name()); + } + std::vector in_edges; std::vector control_edges; for (const Node* node : body_nodes) { @@ -243,6 +248,41 @@ Status FillFunctionBody( if (node->op_def().is_stateful()) { fdef->mutable_signature()->set_is_stateful(true); } + + // If this node has any attributes with placeholder value, add the + // attribute to FunctionDef signature. + for (const auto& iter : node->attrs()) { + if (iter.second.placeholder().empty()) { + continue; + } + + // If we already added the attribute, skip it. + string func_attr_name = iter.second.placeholder(); + if (func_attr_names.find(func_attr_name) != func_attr_names.end()) { + continue; + } + + // This node's attribute is a placeholder value, so it does not have type + // information. We check node's OpDef for attribute type. + string node_attr_name = iter.first; + const OpDef::AttrDef* node_attr_def = nullptr; + for (const auto& node_attr : node->op_def().attr()) { + if (node_attr.name() == node_attr_name) { + node_attr_def = &node_attr; + } + } + if (!node_attr_def) { + return errors::Unimplemented( + "Placeholder value is not supported for attributes not in OpDef. " + "Attribute: ", + node_attr_name, ", OpDef: ", node->op_def().DebugString()); + } + OpDef::AttrDef* attr_def = fdef->mutable_signature()->add_attr(); + attr_def->set_name(func_attr_name); + attr_def->set_type(node_attr_def->type()); + + func_attr_names.insert(func_attr_name); + } } return Status::OK(); } diff --git a/tensorflow/c/c_api_function_test.cc b/tensorflow/c/c_api_function_test.cc index 73fe73769bc1219ce865149d67d333c53371ccc5..946f8c4a2c3fb25f908d809e00bf579b40a8668b 100644 --- a/tensorflow/c/c_api_function_test.cc +++ b/tensorflow/c/c_api_function_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/c/c_api.h" +#include "tensorflow/c/c_api_internal.h" #include "tensorflow/c/c_test_util.h" #include "tensorflow/core/framework/function.pb.h" #include "tensorflow/core/framework/op_def.pb.h" @@ -1230,6 +1231,53 @@ void DefineFunction(const char* name, TF_Function** func, ASSERT_NE(*func, nullptr); } +REGISTER_OP("CustomOp") + .Output("output: float32") + .Attr("index: int") + .SetShapeFn(tensorflow::shape_inference::UnknownShape); + +void NodeWithPlaceholderAttrHelper(TF_Graph* graph, TF_Status* s, + const char* name, const char* placeholder, + TF_Operation** op) { + TF_OperationDescription* desc = TF_NewOperation(graph, "CustomOp", name); + TF_SetAttrPlaceholder(desc, "index", placeholder); + *op = TF_FinishOperation(desc, s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + ASSERT_NE(*op, nullptr); +} + +TEST_F(CApiFunctionTest, GraphToFunctionDefWithPlaceholderAttr) { + std::unique_ptr func_graph( + TF_NewGraph(), TF_DeleteGraph); + std::unique_ptr s(TF_NewStatus(), + TF_DeleteStatus); + + TF_Operation *node1, *node2, *node3; + NodeWithPlaceholderAttrHelper(func_graph.get(), s.get(), "node1", "v1", + &node1); + NodeWithPlaceholderAttrHelper(func_graph.get(), s.get(), "node2", "v1", + &node2); + NodeWithPlaceholderAttrHelper(func_graph.get(), s.get(), "node3", "v2", + &node3); + + TF_Output inputs[] = {}; + TF_Output outputs[] = {{node1, 0}, {node2, 0}, {node3, 0}}; + func_ = TF_GraphToFunction( + func_graph.get(), "func", /*append_hash_to_fn_name=*/false, -1, + /*opers=*/nullptr, 0, inputs, 3, outputs, + /*output_names=*/nullptr, + /*opts=*/nullptr, /*description=*/nullptr, s.get()); + ASSERT_EQ(TF_OK, TF_GetCode(s.get())) << TF_Message(s.get()); + ASSERT_NE(func_, nullptr); + + // Verify that FunctionDef has 2 attributes, "v1" and "v2". + ASSERT_EQ(func_->fdef.signature().attr().size(), 2); + EXPECT_EQ(func_->fdef.signature().attr(0).name(), "v1"); + EXPECT_EQ(func_->fdef.signature().attr(0).type(), "int"); + EXPECT_EQ(func_->fdef.signature().attr(1).name(), "v2"); + EXPECT_EQ(func_->fdef.signature().attr(1).type(), "int"); +} + TEST_F(CApiFunctionTest, SetGradientAndRun) { // Define the function and its grad DefineFunction(func_name_, &func_); diff --git a/tensorflow/c/c_api_internal.h b/tensorflow/c/c_api_internal.h index 5ba26d3c585350aa510f9970cbfc246a9a108543..73283d775639b297857b2a50007dc7c28b1f39a3 100644 --- a/tensorflow/c/c_api_internal.h +++ b/tensorflow/c/c_api_internal.h @@ -228,6 +228,8 @@ void RecordMutation(TF_Graph* graph, const TF_Operation& op, bool ExtendSessionGraphHelper(TF_Session* session, TF_Status* status) LOCKS_EXCLUDED(session->graph->mu, session->mu); +std::string getTF_OutputDebugString(TF_Output node); + } // end namespace tensorflow #endif // TENSORFLOW_C_C_API_INTERNAL_H_ diff --git a/tensorflow/c/c_api_test.cc b/tensorflow/c/c_api_test.cc index d5934a10395ae094f65d3bc8b6cd7b94dbd32410..2be03bf0de6277fc63c353ad6dc63bec096a6993 100644 --- a/tensorflow/c/c_api_test.cc +++ b/tensorflow/c/c_api_test.cc @@ -163,6 +163,7 @@ TEST(CAPI, AllocateTensor) { EXPECT_EQ(dims[0], TF_Dim(t, 0)); EXPECT_EQ(dims[1], TF_Dim(t, 1)); EXPECT_EQ(num_bytes, TF_TensorByteSize(t)); + EXPECT_EQ(6, TF_TensorElementCount(t)); TF_DeleteTensor(t); } @@ -1467,6 +1468,41 @@ TEST(CAPI, DeletingNullPointerIsSafe) { TF_DeleteStatus(status); } +TEST(CAPI, TestBitcastFrom_Reshape) { + int64_t dims[] = {2, 3}; + TF_Tensor* a = + TF_AllocateTensor(TF_UINT64, dims, 2, 6 * TF_DataTypeSize(TF_UINT64)); + TF_Tensor* b = + TF_AllocateTensor(TF_UINT64, nullptr, 0, TF_DataTypeSize(TF_UINT64)); + EXPECT_NE(a, nullptr); + EXPECT_NE(b, nullptr); + + EXPECT_EQ(6, TF_TensorElementCount(a)); + EXPECT_EQ(1, TF_TensorElementCount(b)); + EXPECT_EQ(6 * TF_DataTypeSize(TF_UINT64), TF_TensorByteSize(a)); + EXPECT_EQ(TF_DataTypeSize(TF_UINT64), TF_TensorByteSize(b)); + + int64_t new_dims[] = {3, 2}; + TF_Status* status = TF_NewStatus(); + TF_TensorBitcastFrom(a, TF_UINT64, b, new_dims, 2, status); + ASSERT_EQ(TF_OK, TF_GetCode(status)); + TF_DeleteStatus(status); + + EXPECT_EQ(6, TF_TensorElementCount(a)); + EXPECT_EQ(6, TF_TensorElementCount(b)); + EXPECT_EQ(6 * TF_DataTypeSize(TF_UINT64), TF_TensorByteSize(a)); + EXPECT_EQ(6 * TF_DataTypeSize(TF_UINT64), TF_TensorByteSize(b)); + + // Check that a write to one tensor shows up in the other. + *(static_cast(TF_TensorData(a))) = 4; + EXPECT_EQ(4, *(static_cast(TF_TensorData(b)))); + *(static_cast(TF_TensorData(b))) = 6; + EXPECT_EQ(6, *(static_cast(TF_TensorData(a)))); + + TF_DeleteTensor(a); + TF_DeleteTensor(b); +} + REGISTER_OP("TestOpWithNoGradient") .Input("x: T") .Output("y: T") diff --git a/tensorflow/c/c_test.c b/tensorflow/c/c_test.c index 100aefd58c6d29b3da90249da6a4b4c9cc801d1e..7468122cd567270c8454f886e478be34c2c15cbf 100644 --- a/tensorflow/c/c_test.c +++ b/tensorflow/c/c_test.c @@ -25,6 +25,16 @@ limitations under the License. #include "tensorflow/c/env.h" #include "tensorflow/c/kernels.h" +// A create function. This will never actually get called in this test, it's +// just nice to know that it compiles. +void* create(TF_OpKernelConstruction* ctx) { + TF_DataType type; + TF_Status* s = TF_NewStatus(); + TF_OpKernelConstruction_GetAttrType(ctx, "foobar", &type, s); + TF_DeleteStatus(s); + return NULL; +} + // A compute function. This will never actually get called in this test, it's // just nice to know that it compiles. void compute(void* kernel, TF_OpKernelContext* ctx) { @@ -75,7 +85,7 @@ int main(int argc, char** argv) { TF_StringStreamDone(s); TF_KernelBuilder* b = - TF_NewKernelBuilder("SomeOp", "SomeDevice", NULL, &compute, NULL); + TF_NewKernelBuilder("SomeOp", "SomeDevice", &create, &compute, NULL); TF_RegisterKernelBuilder("someKernel", b, status); TF_DeleteStatus(status); diff --git a/tensorflow/c/eager/BUILD b/tensorflow/c/eager/BUILD index 9f09ad1fc307ed96f327df289c0a45bf79325d7e..257be6379c09841d1427813a0aa25b10a205016d 100644 --- a/tensorflow/c/eager/BUILD +++ b/tensorflow/c/eager/BUILD @@ -3,11 +3,19 @@ licenses(["notice"]) # Apache 2.0 load( "//tensorflow:tensorflow.bzl", - "tf_cuda_cc_test", - "tf_cc_test", "tf_copts", - "tfe_xla_copts", + "tf_cuda_cc_test", "tf_cuda_library", + "tfe_xla_copts", +) +load( + "//tensorflow/core:platform/default/build_config.bzl", + "tf_additional_device_tracer_test_flags", + "tf_kernel_tests_linkstatic", +) +load( + "//tensorflow/core:platform/default/build_config_root.bzl", + "tf_cuda_tests_tags", ) tf_cuda_library( @@ -62,6 +70,7 @@ tf_cuda_library( "//tensorflow/core/distributed_runtime:remote_device", "//tensorflow/core/distributed_runtime:server_lib", "//tensorflow/core/distributed_runtime:worker_env", + "//tensorflow/core/profiler/lib:profiler_session", "//tensorflow/core:gpu_runtime", ], ) @@ -101,6 +110,7 @@ tf_cuda_library( "//tensorflow/core/distributed_runtime/rpc:grpc_worker_service", "//tensorflow/core/distributed_runtime/rpc:rpc_rendezvous_mgr", "//tensorflow/core/distributed_runtime/rpc/eager:grpc_eager_client", + "//tensorflow/core/profiler/lib:profiler_session", ], ) @@ -200,10 +210,36 @@ tf_cuda_library( "//tensorflow/core/distributed_runtime:remote_device", "//tensorflow/core/distributed_runtime:server_lib", "//tensorflow/core/distributed_runtime:worker_env", + "//tensorflow/core/profiler/rpc:profiler_server", "//tensorflow/core:gpu_runtime", ], ) +tf_cuda_cc_test( + name = "c_api_experimental_test", + size = "small", + srcs = [ + "c_api_experimental_test.cc", + ], + args = + ["--heap_check=local"] + tf_additional_device_tracer_test_flags(), + linkstatic = tf_kernel_tests_linkstatic(), + tags = tf_cuda_tests_tags() + ["nomac"], + deps = [ + ":c_api_experimental", + ":c_api_test_util", + "//tensorflow/c:c_test_util", + "//tensorflow/cc/profiler", + "//tensorflow/contrib/tpu/profiler:trace_events_proto_cc", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "//tensorflow/core/profiler:protos_all_cc", + "@com_google_absl//absl/strings", + ], +) + cc_library( name = "tape", hdrs = ["tape.h"], diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc index 027d752f420238da867cb9d8c116640e1730caaa..af13f487af91594fedd4d5f77592682a6f98c34f 100755 --- a/tensorflow/c/eager/c_api.cc +++ b/tensorflow/c/eager/c_api.cc @@ -356,6 +356,8 @@ TFE_TensorHandle* TFE_NewTensorHandle(TF_Tensor* t, TF_Status* status) { void TFE_DeleteTensorHandle(TFE_TensorHandle* h) { if (h == nullptr) return; + VLOG(1) << "Deleting tensor handle " << h << " with internal handle " + << h->handle; if (h->handle) { h->handle->Unref(); } @@ -443,15 +445,15 @@ TF_Tensor* TFE_TensorHandleResolve(TFE_TensorHandle* h, TF_Status* status) { return nullptr; } // TODO(agarwal): move this implementation inside TFE_TensorHandle. - tensorflow::Device* d = nullptr; - tensorflow::Device* op_device = nullptr; const tensorflow::Tensor* t = nullptr; - status->status = h->handle->TensorAndDevice(&t, &d, &op_device); - if (!status->status.ok()) return nullptr; tensorflow::TensorHandle* h_cpu = nullptr; - if (!IsCPU(d)) { - status->status = h->handle->CopyToDevice( - h->handle->Context(), h->handle->Context()->HostCPU(), &h_cpu); + tensorflow::Device* d = nullptr; + tensorflow::Device* op_device = nullptr; + + if (h->handle->IsRemote()) { + status->status = EagerCopyToDevice( + h->handle, h->handle->Context(), + h->handle->Context()->HostCPU()->name().c_str(), &h_cpu); if (!status->status.ok()) { return nullptr; } @@ -460,6 +462,22 @@ TF_Tensor* TFE_TensorHandleResolve(TFE_TensorHandle* h, TF_Status* status) { h_cpu->Unref(); return nullptr; } + } else { + status->status = h->handle->TensorAndDevice(&t, &d, &op_device); + if (!status->status.ok()) return nullptr; + + if (!IsCPU(d)) { + status->status = h->handle->CopyToDevice( + h->handle->Context(), h->handle->Context()->HostCPU(), &h_cpu); + if (!status->status.ok()) { + return nullptr; + } + status->status = h_cpu->TensorAndDevice(&t, &d, &op_device); + if (!status->status.ok()) { + h_cpu->Unref(); + return nullptr; + } + } } TF_Tensor* retval = tensorflow::TF_TensorFromTensor(*t, status); if (h_cpu != nullptr) { @@ -696,6 +714,7 @@ void TFE_OpSetAttrFunctionList(TFE_Op* op, const char* attr_name, void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals, TF_Status* status) { + VLOG(1) << "Calling TFE_Execute() on op " << op; tensorflow::gtl::InlinedVector handle_retvals( *num_retvals); status->status = @@ -738,6 +757,10 @@ void TFE_ContextAddFunction(TFE_Context* ctx, TF_Function* function, status->status = ctx->context.AddFunctionDef(function->fdef); } +unsigned char TFE_ContextHasFunction(TFE_Context* ctx, const char* name) { + return ctx->context.FindFunctionDef(name) != nullptr; +} + void TFE_ContextEnableRunMetadata(TFE_Context* ctx) { ctx->context.SetShouldStoreMetadata(true); } @@ -774,7 +797,7 @@ void TFE_ContextExportRunMetadata(TFE_Context* ctx, TF_Buffer* buf, if (!status->status.ok()) return; tensorflow::mutex_lock ml(*ctx->context.MetadataMu()); status->status = MessageToBuffer(*ctx->context.RunMetadataProto(), buf); - ctx->context.RunMetadataProto()->Clear(); + ctx->context.ClearRunMetadata(); } namespace { diff --git a/tensorflow/c/eager/c_api.h b/tensorflow/c/eager/c_api.h index 120748ab763a3358b6e38e64bb3b6fd2ea32f7c3..044dfb7415b027b707af05a197fdb41fe1f6d2e5 100755 --- a/tensorflow/c/eager/c_api.h +++ b/tensorflow/c/eager/c_api.h @@ -393,6 +393,10 @@ TF_CAPI_EXPORT extern void TFE_ContextAddFunction(TFE_Context* ctx, TF_Function* function, TF_Status* status); +// Checks whether a function is registered under `name`. +TF_CAPI_EXPORT unsigned char TFE_ContextHasFunction(TFE_Context* ctx, + const char* name); + // Enables tracing of RunMetadata on the ops executed from this context. TF_CAPI_EXPORT extern void TFE_ContextEnableRunMetadata(TFE_Context* ctx); diff --git a/tensorflow/c/eager/c_api_debug.cc b/tensorflow/c/eager/c_api_debug.cc index 52b0824552855860dfb138f3ac9a5d3afa7dc965..ffcd5ace0b98597363abe63201bf6c328a03212f 100644 --- a/tensorflow/c/eager/c_api_debug.cc +++ b/tensorflow/c/eager/c_api_debug.cc @@ -83,7 +83,7 @@ TF_CAPI_EXPORT extern TFE_TensorDebugInfo* TFE_TensorHandleTensorDebugInfo( } } - if (xla::ShapeUtil::IsTuple(padded_shape)) { + if (padded_shape.IsTuple()) { if (xla::ShapeUtil::TupleElementCount(padded_shape) != 2) { // Currently, the only case of XlaTensor containing a tuple shape is to // represent 64 bit ints, doubles, and complex numbers (we don't support @@ -99,7 +99,7 @@ TF_CAPI_EXPORT extern TFE_TensorDebugInfo* TFE_TensorHandleTensorDebugInfo( xla::Shape shape0 = xla::ShapeUtil::GetTupleElementShape(padded_shape, 0); const xla::Shape& shape1 = xla::ShapeUtil::GetTupleElementShape(padded_shape, 1); - if (xla::ShapeUtil::IsTuple(shape0) || xla::ShapeUtil::IsTuple(shape1)) { + if (shape0.IsTuple() || shape1.IsTuple()) { status->status = tensorflow::errors::InvalidArgument( "XlaTensors should not contain nested tuples. Shape: ", padded_shape.DebugString()); diff --git a/tensorflow/c/eager/c_api_experimental.cc b/tensorflow/c/eager/c_api_experimental.cc index 3461d81b93594021b4e977aa60ef5d1f92e1fc5c..06bbb4ac41256524b566657105cc5d5858234405 100644 --- a/tensorflow/c/eager/c_api_experimental.cc +++ b/tensorflow/c/eager/c_api_experimental.cc @@ -17,7 +17,54 @@ limitations under the License. #include "tensorflow/c/c_api.h" #include "tensorflow/c/eager/c_api_internal.h" +#include "tensorflow/core/profiler/rpc/profiler_server.h" + +using tensorflow::string; void TFE_OpConsumeInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status) { op->operation.ConsumeInput(h->handle); } + +TFE_Profiler* TFE_NewProfiler(TFE_ProfilerContext* ctx) { + return new TFE_Profiler(ctx); +} + +bool TFE_ProfilerIsOk(TFE_Profiler* profiler) { + return profiler->profiler->Status().ok(); +} + +void TFE_DeleteProfiler(TFE_Profiler* profiler) { delete profiler; } + +void TFE_ProfilerSerializeToString(TFE_Context* ctx, TFE_Profiler* profiler, + TF_Buffer* buf, TF_Status* status) { + TFE_ContextAsyncWait(ctx, status); + if (!status->status.ok()) return; + string content; + status->status = profiler->profiler->SerializeToString(&content); + void* data = tensorflow::port::Malloc(content.length()); + content.copy(static_cast(data), content.length(), 0); + buf->data = data; + buf->length = content.length(); + buf->data_deallocator = [](void* data, size_t length) { + tensorflow::port::Free(data); + }; +} + +TFE_ProfilerContext* TFE_NewProfilerContext() { + return new TFE_ProfilerContext; +} + +void TFE_ProfilerContextSetEagerContext(TFE_ProfilerContext* profiler_context, + TFE_Context* eager_context) { + profiler_context->profiler_context.eager_context = &eager_context->context; +} + +void TFE_DeleteProfilerContext(TFE_ProfilerContext* profiler_context) { + delete profiler_context; +} + +void TFE_StartProfilerServer(TFE_ProfilerContext* context, int port) { + // Release child thread intentionally. The child thread can be terminate by + // terminating the main thread. + tensorflow::StartProfilerServer(&context->profiler_context, port).release(); +} diff --git a/tensorflow/c/eager/c_api_experimental.h b/tensorflow/c/eager/c_api_experimental.h index 4ee6c066eef7f60fd2f50a197c4dc20ed9c9f927..51a5fa0d816a179bb52940f6d8aead867ec9a267 100644 --- a/tensorflow/c/eager/c_api_experimental.h +++ b/tensorflow/c/eager/c_api_experimental.h @@ -25,6 +25,48 @@ extern "C" { TF_CAPI_EXPORT extern void TFE_OpConsumeInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status); +typedef struct TFE_ProfilerContext TFE_ProfilerContext; + +// A profiler which will start profiling when creating the object and will stop +// when the object is destroyed. It will profile all operations run under the +// given TFE_Context. Multiple instance of it can be created, but at most one +// of them will profile for each TFE_Context. +// Thread-safety: TFE_Profiler is thread-safe. +typedef struct TFE_Profiler TFE_Profiler; + +TF_CAPI_EXPORT extern TFE_Profiler* TFE_NewProfiler(TFE_ProfilerContext* ctx); +TF_CAPI_EXPORT extern bool TFE_ProfilerIsOk(TFE_Profiler* profiler); +TF_CAPI_EXPORT extern void TFE_DeleteProfiler(TFE_Profiler* profiler); + +// The output string is a binary string of tensorflow.tpu.Trace. User can write +// the string to file for offline analysis by tensorboard. +TF_CAPI_EXPORT extern void TFE_ProfilerSerializeToString(TFE_Context* ctx, + TFE_Profiler* profiler, + TF_Buffer* buf, + TF_Status* status); + +// Return a new profiler context object. +TF_CAPI_EXPORT extern TFE_ProfilerContext* TFE_NewProfilerContext(void); + +// Set the eager context in TFE_ProfilerServerOptions +TF_CAPI_EXPORT extern void TFE_ProfilerContextSetEagerContext( + TFE_ProfilerContext* profiler_context, TFE_Context* eager_context); + +// Destroy a profiler context object. +TF_CAPI_EXPORT extern void TFE_DeleteProfilerContext( + TFE_ProfilerContext* profiler_context); + +// Start a profiler grpc server which listens to specified port. It will start +// the server on its own thread. It can be shutdown by terminating tensorflow. +// It can be used in both Eager mode and graph mode. Creating multiple profiler +// server is allowed. The service defined in +// tensorflow/contrib/tpu/profiler/tpu_profiler.proto. Please use +// tensorflow/contrib/tpu/profiler/capture_tpu_profile to capture tracable +// file following +// https://cloud.google.com/tpu/docs/cloud-tpu-tools#capture_trace. +TF_CAPI_EXPORT extern void TFE_StartProfilerServer(TFE_ProfilerContext* context, + int port); + #ifdef __cplusplus } /* end extern "C" */ #endif diff --git a/tensorflow/c/eager/c_api_experimental_test.cc b/tensorflow/c/eager/c_api_experimental_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..7b84c5426fa77e7e0d62b7e263abc20400b5a171 --- /dev/null +++ b/tensorflow/c/eager/c_api_experimental_test.cc @@ -0,0 +1,129 @@ +/* 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/c/eager/c_api_experimental.h" + +#include +#include "tensorflow/c/eager/c_api_test_util.h" +#include "tensorflow/cc/profiler/profiler.h" +#include "tensorflow/contrib/tpu/profiler/trace_events.pb.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/protobuf.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/platform/test_benchmark.h" + +using tensorflow::string; + +namespace tensorflow { +namespace { + +static bool HasSubstr(absl::string_view base, absl::string_view substr) { + bool ok = str_util::StrContains(base, substr); + EXPECT_TRUE(ok) << base << ", expected substring " << substr; + return ok; +} + +void ExecuteWithProfiling(bool async) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(async)); + TFE_Context* ctx = TFE_NewContext(opts, status); + TFE_ProfilerContext* profiler_context = TFE_NewProfilerContext(); + TFE_ProfilerContextSetEagerContext(profiler_context, ctx); + TFE_Profiler* profiler = TFE_NewProfiler(profiler_context); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + TFE_DeleteProfilerContext(profiler_context); + + TFE_TensorHandle* m = TestMatrixTensorHandle(); + TFE_Op* matmul = MatMulOp(ctx, m, m); + TFE_TensorHandle* retvals[1] = {nullptr}; + int num_retvals = 1; + + // Run op on GPU if it is present. + string gpu_device_name; + if (GetDeviceName(ctx, &gpu_device_name, "GPU")) { + TFE_OpSetDevice(matmul, gpu_device_name.c_str(), status); + ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status); + const char* device_name = TFE_OpGetDevice(matmul, status); + ASSERT_TRUE(strstr(device_name, "GPU:0") != nullptr); + } + + TFE_Execute(matmul, &retvals[0], &num_retvals, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteOp(matmul); + TFE_DeleteTensorHandle(m); + + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + ASSERT_EQ(1, num_retvals); + TF_Buffer* profiler_result = TF_NewBuffer(); + TFE_ProfilerSerializeToString(ctx, profiler, profiler_result, status); + TFE_DeleteProfiler(profiler); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + tensorflow::tpu::Trace profile_proto; + EXPECT_TRUE(profile_proto.ParseFromString( + {reinterpret_cast(profiler_result->data), + profiler_result->length})); + string profile_proto_str = profile_proto.DebugString(); + if (!gpu_device_name.empty()) { + EXPECT_TRUE(HasSubstr(profile_proto_str, "GPU:0")); + // device name with "stream:all" is collected by Device Tracer. + EXPECT_TRUE(HasSubstr(profile_proto_str, "stream:all")); + } + EXPECT_TRUE(HasSubstr(profile_proto_str, "CPU:0")); + TF_DeleteBuffer(profiler_result); + + TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status); + TFE_DeleteTensorHandle(retvals[0]); + TFE_DeleteContext(ctx); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + float product[4] = {0}; + EXPECT_EQ(sizeof(product), TF_TensorByteSize(t)); + memcpy(&product[0], TF_TensorData(t), TF_TensorByteSize(t)); + TF_DeleteTensor(t); + EXPECT_EQ(7, product[0]); + EXPECT_EQ(10, product[1]); + EXPECT_EQ(15, product[2]); + EXPECT_EQ(22, product[3]); + TF_DeleteStatus(status); +} +TEST(CAPI, ExecuteWithTracing) { ExecuteWithProfiling(false); } +TEST(CAPI, ExecuteWithTracingAsync) { ExecuteWithProfiling(true); } + +TEST(CAPI, MultipleProfilerSession) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetAsync(opts, static_cast(false)); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + TFE_ProfilerContext* profiler_context = TFE_NewProfilerContext(); + TFE_ProfilerContextSetEagerContext(profiler_context, ctx); + + TFE_Profiler* profiler1 = TFE_NewProfiler(profiler_context); + EXPECT_TRUE(TFE_ProfilerIsOk(profiler1)); + + TFE_Profiler* profiler2 = TFE_NewProfiler(profiler_context); + EXPECT_FALSE(TFE_ProfilerIsOk(profiler2)); + + TFE_DeleteProfiler(profiler1); + TFE_DeleteProfiler(profiler2); + TFE_DeleteProfilerContext(profiler_context); +} + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/c/eager/c_api_internal.h b/tensorflow/c/eager/c_api_internal.h index 67bc1bcd24605f8363d6a7c8d5d6a0836a42fc82..a563e4b8f50f2a90497736f4cb9ca234400bfa04 100644 --- a/tensorflow/c/eager/c_api_internal.h +++ b/tensorflow/c/eager/c_api_internal.h @@ -52,6 +52,7 @@ limitations under the License. #include "tensorflow/core/lib/gtl/stl_util.h" #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/thread_annotations.h" +#include "tensorflow/core/profiler/lib/profiler_session.h" #include "tensorflow/core/public/version.h" struct TFE_ContextOptions { @@ -82,6 +83,12 @@ struct TFE_TensorHandle { TFE_TensorHandle(tensorflow::TensorHandle* handle) : handle(handle) {} tensorflow::TensorHandle* handle; + + // Create a symbolic tensor. + TFE_TensorHandle(TF_Output t, TF_DataType dtype) + : handle(new tensorflow::TensorHandle( + tensorflow::OutputGraphNode{t.oper, t.index}, + static_cast(dtype))) {} }; struct TFE_TensorDebugInfo { @@ -100,6 +107,18 @@ struct TFE_Op { tensorflow::EagerOperation operation; }; +struct TFE_ProfilerContext { + tensorflow::ProfilerContext profiler_context; +}; + +struct TFE_Profiler { + TFE_Profiler(TFE_ProfilerContext* ctx) { + profiler = tensorflow::ProfilerSession::Create(&ctx->profiler_context); + } + + std::unique_ptr profiler; +}; + namespace tensorflow { // Set an AttrValue on the op. Doesn't handle the list types. void SetOpAttrValueScalar(TFE_Context* ctx, TFE_Op* op, @@ -107,4 +126,24 @@ void SetOpAttrValueScalar(TFE_Context* ctx, TFE_Op* op, const char* attr_name, TF_Status* status); } // namespace tensorflow +struct TFE_TraceContext { + TF_Graph* const graph; + + unsigned int node_counter = 0; + // Each tensor handle will have its ref count incremented when it's added as a + // map key, and decremented when this object is destroyed. + std::map input_tensor_map; + std::vector>* input_tensors = + nullptr; + + TFE_TraceContext(TF_Graph* graph) : graph(graph) {} + + ~TFE_TraceContext() { + delete input_tensors; + for (auto input : input_tensor_map) { + input.first->Unref(); + } + } +}; + #endif // TENSORFLOW_C_EAGER_C_API_INTERNAL_H_ diff --git a/tensorflow/c/eager/c_api_test.cc b/tensorflow/c/eager/c_api_test.cc index 6b39b79ee82f9c7baaf856e573a42b7da65691e5..3d1ca4fb4b561a03ea9d879b1876fb1fd08a3139 100644 --- a/tensorflow/c/eager/c_api_test.cc +++ b/tensorflow/c/eager/c_api_test.cc @@ -175,13 +175,8 @@ void TestRemoteExecute(bool async) { TFE_Execute(matmul, &retvals[0], &num_retvals, status); EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); - auto* retval_task0 = TFE_TensorHandleCopyToDevice( - retvals[0], ctx, "/job:localhost/replica:0/task:0/device:CPU:0", status); - ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); - - TF_Tensor* t = TFE_TensorHandleResolve(retval_task0, status); + TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status); ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); - TFE_DeleteTensorHandle(retval_task0); float product[4] = {0}; EXPECT_EQ(sizeof(product), TF_TensorByteSize(t)); memcpy(&product[0], TF_TensorData(t), TF_TensorByteSize(t)); diff --git a/tensorflow/c/kernels.cc b/tensorflow/c/kernels.cc index 2a4eaecb6cf2740a522b1e849d1306ebde6c4577..71181ae430ab64106e2a75937bd54fbf2efc61ac 100644 --- a/tensorflow/c/kernels.cc +++ b/tensorflow/c/kernels.cc @@ -48,9 +48,10 @@ TF_KernelBuilder* TF_NewKernelBuilder( } void TF_DeleteKernelBuilder(TF_KernelBuilder* builder) { - DCHECK_NE(builder, nullptr); - delete builder->cc_builder; - delete builder; + if (builder != nullptr) { + delete builder->cc_builder; + delete builder; + } } namespace tensorflow { @@ -158,3 +159,41 @@ void TF_SetOutput(TF_OpKernelContext* ctx, int i, const TF_Tensor* tensor, cc_ctx->set_output(i, cc_tensor); } } + +void TF_OpKernelConstruction_Failure(TF_OpKernelConstruction* ctx, + TF_Status* status) { + auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelConstruction*>(ctx); + ::tensorflow::Status s(::tensorflow::StatusFromTF_Status(status)); + cc_ctx->CtxFailure(s); +} + +void TF_OpKernelContext_Failure(TF_OpKernelContext* ctx, TF_Status* status) { + auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx); + ::tensorflow::Status s(::tensorflow::StatusFromTF_Status(status)); + cc_ctx->CtxFailure(s); +} + +#define DEFINE_TF_GETATTR(func, c_type, cc_type) \ + void TF_OpKernelConstruction_GetAttr##func(TF_OpKernelConstruction* ctx, \ + const char* attr_name, \ + c_type* val, TF_Status* status) { \ + TF_SetStatus(status, TF_OK, ""); \ + cc_type v; \ + auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelConstruction*>(ctx); \ + ::tensorflow::Status s = cc_ctx->GetAttr(attr_name, &v); \ + ::tensorflow::Set_TF_Status_from_Status(status, s); \ + if (s.ok()) { \ + *val = static_cast(v); \ + } \ + } + +DEFINE_TF_GETATTR(Type, TF_DataType, tensorflow::DataType) + +TF_DataType TF_ExpectedOutputDataType(TF_OpKernelContext* ctx, int i) { + auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx); + return static_cast(cc_ctx->expected_output_dtype(i)); +} + +int64_t TF_StepId(TF_OpKernelContext* ctx) { + return reinterpret_cast<::tensorflow::OpKernelContext*>(ctx)->step_id(); +} diff --git a/tensorflow/c/kernels.h b/tensorflow/c/kernels.h index cefc30bcdf89bdc14a4406299cc29f74153e77ac..c47bfa8aa3a721d422a0a1536b924f3e53793193 100644 --- a/tensorflow/c/kernels.h +++ b/tensorflow/c/kernels.h @@ -111,6 +111,32 @@ TF_CAPI_EXPORT extern void TF_SetOutput(TF_OpKernelContext* ctx, int i, const TF_Tensor* tensor, TF_Status* status); +// Notifies the given OpKernelConstruction that kernel construction has failed. +TF_CAPI_EXPORT extern void TF_OpKernelConstruction_Failure( + TF_OpKernelConstruction* ctx, TF_Status* status); + +// Notifies the given OpKernelContext that the kernel's compute function has +// failed. +TF_CAPI_EXPORT extern void TF_OpKernelContext_Failure(TF_OpKernelContext* ctx, + TF_Status* status); + +// Returns the expected output data type of the ith output. If i < 0 or +// i >= TF_NumOutputs(ctx), the program aborts. +TF_CAPI_EXPORT extern TF_DataType TF_ExpectedOutputDataType( + TF_OpKernelContext* ctx, int i); + +// Returns the step ID of the given context. +TF_CAPI_EXPORT extern int64_t TF_StepId(TF_OpKernelContext* ctx); + +// Interprets the named kernel construction attribute as a TF_DataType and +// places it into *val. *status is set to TF_OK. +// +// If the attribute could not be found or could not be interpreted as +// TF_DataType, *status is populated with an error. +TF_CAPI_EXPORT extern void TF_OpKernelConstruction_GetAttrType( + TF_OpKernelConstruction* ctx, const char* attr_name, TF_DataType* val, + TF_Status* status); + #ifdef __cplusplus } /* end extern "C" */ #endif diff --git a/tensorflow/c/kernels_test.cc b/tensorflow/c/kernels_test.cc index e659ee3c3d258a626ccf03a782ec031b5a703a48..608887722f7bca44c884a3426d5e378e9387a530 100644 --- a/tensorflow/c/kernels_test.cc +++ b/tensorflow/c/kernels_test.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/c/kernels.h" #include "tensorflow/c/c_api.h" +#include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/kernel_def.pb.h" #include "tensorflow/core/framework/node_def.pb_text.h" #include "tensorflow/core/framework/op.h" @@ -35,12 +36,24 @@ static void* MyCreateFunc(TF_OpKernelConstruction* ctx) { struct MyCustomKernel* s = new struct MyCustomKernel; s->created = true; s->compute_called = false; + + // Exercise attribute reads. + TF_DataType type; + TF_Status* status = TF_NewStatus(); + TF_OpKernelConstruction_GetAttrType(ctx, "SomeDataTypeAttr", &type, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)); + EXPECT_EQ(TF_FLOAT, type); + TF_DeleteStatus(status); + return s; } static void MyComputeFunc(void* kernel, TF_OpKernelContext* ctx) { struct MyCustomKernel* s = static_cast(kernel); s->compute_called = true; + if (ctx != nullptr) { + EXPECT_EQ(43, TF_StepId(ctx)); + } } static void MyDeleteFunc(void* kernel) { @@ -61,6 +74,11 @@ static std::unique_ptr GetFakeKernel(const char* device_name, def.set_device(device_name); def.add_input("input1"); def.add_input("input2"); + + AttrValue v; + v.set_type(DataType::DT_FLOAT); + (*def.mutable_attr())["SomeDataTypeAttr"] = v; + return CreateOpKernel(DeviceType(device_name), nullptr, nullptr, def, 1, status); } @@ -75,7 +93,8 @@ TEST(TestKernel, TestRegisterKernelBuilder) { REGISTER_OP(op_name) .Input("input1: double") .Input("input2: uint8") - .Output("output1: uint8"); + .Output("output1: uint8") + .Attr("SomeDataTypeAttr: type"); TF_KernelBuilder* builder = TF_NewKernelBuilder( op_name, device_name, &MyCreateFunc, &MyComputeFunc, &MyDeleteFunc); @@ -126,7 +145,8 @@ TEST(TestKernel, TestInputAndOutputCount) { REGISTER_OP(op_name) .Input("input1: double") .Input("input2: uint8") - .Output("output1: uint8"); + .Output("output1: uint8") + .Attr("SomeDataTypeAttr: type"); static int num_inputs = 0; static int num_outputs = 0; @@ -155,6 +175,8 @@ TEST(TestKernel, TestInputAndOutputCount) { TF_SetOutput(ctx, 24, input, s); EXPECT_EQ(TF_OUT_OF_RANGE, TF_GetCode(s)); + EXPECT_EQ(TF_UINT8, TF_ExpectedOutputDataType(ctx, 0)); + TF_DeleteStatus(s); if (input != nullptr) { TF_DeleteTensor(input); @@ -175,6 +197,7 @@ TEST(TestKernel, TestInputAndOutputCount) { OpKernelContext::Params p; DummyDevice dummy_device(nullptr, false); p.device = &dummy_device; + p.step_id = 43; Tensor t(tensorflow::uint8(123)); @@ -200,4 +223,8 @@ TEST(TestKernel, TestInputAndOutputCount) { } } +TEST(TestKernel, DeleteKernelBuilderIsOkOnNull) { + TF_DeleteKernelBuilder(nullptr); +} + } // namespace tensorflow diff --git a/tensorflow/cc/framework/cc_op_gen.cc b/tensorflow/cc/framework/cc_op_gen.cc index 39593370d1c243e84dc5b6091724d1d404c102b0..43a33cbea6e1e4a50f61cc7d6d8d70cac6a603d2 100644 --- a/tensorflow/cc/framework/cc_op_gen.cc +++ b/tensorflow/cc/framework/cc_op_gen.cc @@ -321,6 +321,7 @@ std::pair AttrTypeName(StringPiece attr_type) { {"tensor", {"TensorProto", true}}, {"list(tensor)", {"gtl::ArraySlice", true}}, {"func", {"NameAttrList", true}}, + {"list(func)", {"gtl::ArraySlice", true}}, }; auto entry = attr_type_map->find(attr_type); diff --git a/tensorflow/cc/profiler/BUILD b/tensorflow/cc/profiler/BUILD index cf65fe1ab99b49207a64e86310178141b30d07d7..e9838d9aba6554b40082187057851e9c896f8352 100644 --- a/tensorflow/cc/profiler/BUILD +++ b/tensorflow/cc/profiler/BUILD @@ -10,7 +10,7 @@ tf_cuda_cc_test( name = "profiler_test", srcs = ["profiler_test.cc"], tags = [ - "noguitar", # b/77649654 + "nogpu", # b/77649654 ], deps = [ ":profiler", diff --git a/tensorflow/cc/saved_model/loader.cc b/tensorflow/cc/saved_model/loader.cc index 85d3dd01fa51b3c3ba6fcbf5faac03f1ff5630e2..66260fcf4a9b24f78d45010c6e86d4ee398b6d3d 100644 --- a/tensorflow/cc/saved_model/loader.cc +++ b/tensorflow/cc/saved_model/loader.cc @@ -21,11 +21,11 @@ limitations under the License. #include "tensorflow/cc/saved_model/reader.h" #include "tensorflow/core/lib/io/path.h" #include "tensorflow/core/lib/monitoring/counter.h" +#include "tensorflow/core/lib/monitoring/sampler.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/protobuf_internal.h" -#include "tensorflow/core/protobuf/saved_model.pb.h" #include "tensorflow/core/protobuf/saver.pb.h" #include "tensorflow/core/public/session.h" #include "tensorflow/core/public/session_options.h" @@ -42,9 +42,28 @@ auto* load_latency = monitoring::Counter<1>::New( "/tensorflow/cc/saved_model/load_latency", "Latency in microseconds for SavedModels that were successfully loaded.", "model_path"); +auto* load_latency_by_stage = monitoring::Sampler<2>::New( + { + "/tensorflow/cc/saved_model/load_latency_by_stage", // metric name + "Distribution of wall time spent (in microseconds) in each stage " + "(restore graph from disk, run init graph op, etc) when loading the " + "model", + "model_path", + "stage", + }, + // Scale of 10, power of 1.8 with bucket count 33 (~20 minutes). + monitoring::Buckets::Exponential(10, 1.8, 33)); + constexpr char kLoadAttemptFail[] = "fail"; constexpr char kLoadAttemptSuccess[] = "success"; +uint64 GetLatencyMicroseconds(const uint64 start_microseconds) { + const uint64 end_microseconds = Env::Default()->NowMicros(); + // Avoid clock skew. + if (end_microseconds < start_microseconds) return 0; + return end_microseconds - start_microseconds; +} + Status LoadMetaGraphIntoSession(const MetaGraphDef& meta_graph_def, const SessionOptions& session_options, std::unique_ptr* session) { @@ -242,6 +261,7 @@ Status LoadSavedModelInternal(const SessionOptions& session_options, const string& export_dir, const std::unordered_set& tags, SavedModelBundle* const bundle) { + const uint64 read_start_microseconds = Env::Default()->NowMicros(); TF_RETURN_IF_ERROR(ReadMetaGraphDefFromSavedModel(export_dir, tags, &bundle->meta_graph_def)); @@ -256,12 +276,23 @@ Status LoadSavedModelInternal(const SessionOptions& session_options, bundle->meta_graph_def.saver_def().restore_op_name(), bundle->meta_graph_def.saver_def().filename_tensor_name(), asset_file_defs, bundle->session.get())); + // Record walltime spent in restoring graph from disk, but postpone metric + // increments until graph init finishes. + const uint64 restore_graph_walltime = + GetLatencyMicroseconds(read_start_microseconds); + + const uint64 graph_init_start_microseconds = Env::Default()->NowMicros(); string init_op_name; TF_RETURN_IF_ERROR( GetInitOp(export_dir, bundle->meta_graph_def, &init_op_name)); TF_RETURN_IF_ERROR(RunInitOp(run_options, export_dir, bundle->meta_graph_def, asset_file_defs, bundle->session.get(), init_op_name)); + load_latency_by_stage->GetCell(export_dir, "restore_graph") + ->Add(restore_graph_walltime); + // Record wall time spent in init op. + load_latency_by_stage->GetCell(export_dir, "init_graph") + ->Add(GetLatencyMicroseconds(graph_init_start_microseconds)); return Status::OK(); } @@ -275,16 +306,10 @@ Status LoadSavedModel(const SessionOptions& session_options, const uint64 start_microseconds = Env::Default()->NowMicros(); const Status status = LoadSavedModelInternal(session_options, run_options, export_dir, tags, bundle); - const uint64 load_latency_microsecs = [&]() -> uint64 { - const uint64 end_microseconds = Env::Default()->NowMicros(); - // Avoid clock skew. - if (end_microseconds < start_microseconds) return 0; - return end_microseconds - start_microseconds; - }(); auto log_and_count = [&](const string& status_str) { LOG(INFO) << "SavedModel load for tags { " << str_util::Join(tags, " ") << " }; Status: " << status_str << ". Took " - << load_latency_microsecs << " microseconds."; + << GetLatencyMicroseconds(start_microseconds) << " microseconds."; load_attempt_count->GetCell(export_dir, status_str)->IncrementBy(1); }; if (status.ok()) { @@ -292,7 +317,8 @@ Status LoadSavedModel(const SessionOptions& session_options, } else { log_and_count(kLoadAttemptFail); } - load_latency->GetCell(export_dir)->IncrementBy(load_latency_microsecs); + load_latency->GetCell(export_dir) + ->IncrementBy(GetLatencyMicroseconds(start_microseconds)); return status; } diff --git a/tensorflow/cc/tools/freeze_saved_model.cc b/tensorflow/cc/tools/freeze_saved_model.cc index 23e9dc40d23899b9cef168c9128b6d8ed1be3ee9..eeb910178902ca883ed211379ba3f188c139f92e 100644 --- a/tensorflow/cc/tools/freeze_saved_model.cc +++ b/tensorflow/cc/tools/freeze_saved_model.cc @@ -124,7 +124,9 @@ Status GetVariableNameToTensorMap( return Status::OK(); } std::vector variable_names; + variable_names.reserve(variable_names_set.size()); std::vector tensor_names; + tensor_names.reserve(variable_names_set.size()); for (const string& node_name : variable_names_set) { variable_names.push_back(node_name); NodeDef* node_def = name_to_node_map.at(node_name); diff --git a/tensorflow/compat_template.__init__.py b/tensorflow/compat_template.__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..05fd9cd981f70b9f54b65a59a2e92c5405a80f08 --- /dev/null +++ b/tensorflow/compat_template.__init__.py @@ -0,0 +1,51 @@ +# Copyright 2015 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. +# ============================================================================== +"""Bring in all of the public TensorFlow interface into this module.""" + +from __future__ import absolute_import as _absolute_import +from __future__ import division as _division +from __future__ import print_function as _print_function + +import os as _os +import sys as _sys + +# pylint: disable=g-bad-import-order +from tensorflow.python import pywrap_tensorflow # pylint: disable=unused-import + +# API IMPORTS PLACEHOLDER + +from tensorflow.python.tools import component_api_helper as _component_api_helper +_component_api_helper.package_hook( + parent_package_str=__name__, + child_package_str=( + 'tensorflow_estimator.python.estimator.api._v2.estimator')) +_component_api_helper.package_hook( + parent_package_str=__name__, + child_package_str=('tensorflow.python.keras.api._v2.keras')) + +# We would like the following to work for fully enabling 2.0 in a 1.0 install: +# +# import tensorflow.compat.v2 as tf +# tf.enable_v2_behavior() +# +# This make this one symbol available directly. +from tensorflow.python.compat.v2_compat import enable_v2_behavior # pylint: disable=g-import-not-at-top + +# Add module aliases +_current_module = _sys.modules[__name__] +if hasattr(_current_module, 'keras'): + losses = keras.losses + metrics = keras.metrics + optimizers = keras.optimizers diff --git a/tensorflow/compat_template_v1.__init__.py b/tensorflow/compat_template_v1.__init__.py index b966c22b2319aef3b87ef54a283911718d37cf84..9549a71c41a0ba2aac58abd8cfb182aa4eaf3b4f 100644 --- a/tensorflow/compat_template_v1.__init__.py +++ b/tensorflow/compat_template_v1.__init__.py @@ -28,7 +28,8 @@ from tensorflow.python import pywrap_tensorflow # pylint: disable=unused-import from tensorflow.python.tools import component_api_helper as _component_api_helper _component_api_helper.package_hook( parent_package_str=__name__, - child_package_str=('tensorflow_estimator.python.estimator.api.estimator')) + child_package_str=( + 'tensorflow_estimator.python.estimator.api._v1.estimator')) _component_api_helper.package_hook( parent_package_str=__name__, child_package_str=('tensorflow.python.keras.api._v1.keras')) diff --git a/tensorflow/compiler/aot/BUILD b/tensorflow/compiler/aot/BUILD index 16151e77737429f4fbf690fc34b12a70bacebdc4..af016bf80e7a10d8729a1eb385466af48b5810cd 100644 --- a/tensorflow/compiler/aot/BUILD +++ b/tensorflow/compiler/aot/BUILD @@ -30,6 +30,7 @@ cc_library( "flags.h", ], deps = [ + ":aot_only_var_handle_op", ":embedded_protocol_buffers", "//tensorflow/compiler/tf2xla", "//tensorflow/compiler/tf2xla:cpu_function_runtime", @@ -71,6 +72,7 @@ tf_cc_test( ":tfcompile_lib", "//tensorflow/compiler/xla:shape_util", "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "@com_google_absl//absl/strings", @@ -205,6 +207,15 @@ cc_library( ], ) +cc_library( + name = "aot_only_var_handle_op", + srcs = ["aot_only_var_handle_op.cc"], + deps = [ + "//tensorflow/compiler/tf2xla:xla_compiler", + ], + alwayslink = 1, +) + tf_cc_test( name = "benchmark_test", srcs = ["benchmark_test.cc"], diff --git a/tensorflow/compiler/aot/aot_only_var_handle_op.cc b/tensorflow/compiler/aot/aot_only_var_handle_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..0ce36a979f424610a5aa952afa8db2245ed971a9 --- /dev/null +++ b/tensorflow/compiler/aot/aot_only_var_handle_op.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/tf2xla/xla_context.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" + +namespace tensorflow { +namespace { + +// Implementation of varhandle that binds a VarHandleOp to an XlaResource of the +// same name. It is not safe to use this op in a JIT context. +class XlaAotOnlyVarHandleOp : public XlaOpKernel { + public: + explicit XlaAotOnlyVarHandleOp(OpKernelConstruction* c); + void Compile(XlaOpKernelContext* context) override; + + private: + string name_; +}; + +XlaAotOnlyVarHandleOp::XlaAotOnlyVarHandleOp(OpKernelConstruction* c) + : XlaOpKernel(c) { + OP_REQUIRES_OK(c, c->GetAttr("shared_name", &name_)); +} + +void XlaAotOnlyVarHandleOp::Compile(XlaOpKernelContext* context) { + // Look for a resource of the same name. TF also keys that on the container + // and type attributes, but that doesn't seem necessary. + for (const auto& resource : context->xla_context()->resources()) { + if (resource->kind() == XlaResource::kVariable && + resource->name() == name_) { + context->SetResourceOutput(0, resource.get()); + return; + } + } + context->SetStatus( + errors::InvalidArgument("Variable: ", name_, " not configured")); +} +} // namespace + +REGISTER_XLA_OP(Name("VarHandleOp").CompilationOnly(), XlaAotOnlyVarHandleOp); + +} // namespace tensorflow diff --git a/tensorflow/compiler/aot/codegen.cc b/tensorflow/compiler/aot/codegen.cc index 347a365dcb9b16383a8d5fc2f0d2430aa0d77673..da0598736a7d6b7f55458d76ca30fa6ad46a74f9 100644 --- a/tensorflow/compiler/aot/codegen.cc +++ b/tensorflow/compiler/aot/codegen.cc @@ -168,12 +168,12 @@ Status GenArgMethods(const tf2xla::Config& config, const xla::ProgramShapeProto& ps, const CompileResult& compile_result, string* methods) { size_t num_args = ps.parameters_size(); - if (config.feed_size() != num_args) { - return errors::InvalidArgument("mismatch between feed_size(", - config.feed_size(), ") and num_args(", - num_args, ")"); + if (config.feed_size() + config.variable_size() != num_args) { + return errors::InvalidArgument( + "mismatch between feed_size(", config.feed_size(), ")+variable_size(", + config.variable_size(), ") and num_args(", num_args, ")"); } - for (int i = 0; i < num_args; ++i) { + for (int i = 0; i < config.feed_size(); ++i) { std::vector> rewrites; TF_RETURN_IF_ERROR( AddRewritesForShape(i, xla::Shape(ps.parameters(i)), &rewrites)); @@ -212,12 +212,14 @@ Status GenResultMethods(const tf2xla::Config& config, // tuple result, and we rely on this to simplify code generation. return errors::Internal("codegen requires the XLA result to be a tuple"); } - if (config.fetch_size() != ps.result().tuple_shapes_size()) { + size_t num_results = ps.result().tuple_shapes_size(); + if (config.fetch_size() + config.variable_size() != num_results) { return errors::InvalidArgument("mismatch between fetch_size(", - config.feed_size(), ") and tuple_size(", + config.fetch_size(), ")+variable_size(", + config.variable_size(), ") and tuple_size(", ps.result().tuple_shapes_size(), ")"); } - for (int i = 0; i < ps.result().tuple_shapes_size(); ++i) { + for (int i = 0; i < config.fetch_size(); ++i) { std::vector> rewrites; TF_RETURN_IF_ERROR(AddRewritesForShape( i, xla::Shape(ps.result().tuple_shapes(i)), &rewrites)); @@ -245,6 +247,51 @@ Status GenResultMethods(const tf2xla::Config& config, return Status::OK(); } +// Generate methods for variables. +Status GenVariableMethods(const tf2xla::Config& config, + const xla::ProgramShapeProto& ps, string* methods) { + size_t num_args = ps.parameters_size(); + for (int i = config.feed_size(); i < num_args; ++i) { + std::vector> rewrites; + TF_RETURN_IF_ERROR( + AddRewritesForShape(i, xla::Shape(ps.parameters(i)), &rewrites)); + const string code = R"( + void set_var_{{NAME}}_input_data({{TYPE}}* data) { + set_arg_data({{I}}, data); + } +)"; + const tf2xla::Variable& var = config.variable(i - config.feed_size()); + *methods += RewriteWithName( + var.name().empty() ? var.node_name() : var.name(), code, rewrites); + } + size_t num_results = ps.result().tuple_shapes_size(); + for (int i = config.fetch_size(); i < num_results; ++i) { + std::vector> rewrites; + TF_RETURN_IF_ERROR(AddRewritesForShape( + i, xla::Shape(ps.result().tuple_shapes(i)), &rewrites)); + string code = R"( + {{TYPE}}* var_{{NAME}}_result_data() { + return static_cast<{{TYPE}}*>(result_data({{I}})); + } + {{TYPE}}& var_{{NAME}}_result({{DIM_VARS}}) { + return (*static_cast<{{TYPE}}(*){{DIM_SIZES}}>( + result_data({{I}}))){{INDICES}}; + } + const {{TYPE}}* var_{{NAME}}_result_data() const { + return static_cast(result_data({{I}})); + } + const {{TYPE}}& var_{{NAME}}_result({{DIM_VARS}}) const { + return (*static_cast( + result_data({{I}}))){{INDICES}}; + } +)"; + const tf2xla::Variable& var = config.variable(i - config.fetch_size()); + *methods += RewriteWithName( + var.name().empty() ? var.node_name() : var.name(), code, rewrites); + } + return Status::OK(); +} + // Generates code implementing {Arg,Result}Names(), where T is one of // tf2xla::{Feed,Fetch}. Each feed or fetch name results in a C-style string // literal in the array, with nullptr terminating the array. @@ -291,6 +338,14 @@ Status ValidateFeedFetchCppNames(const tf2xla::Config& config) { TF_RETURN_IF_ERROR(ValidateCppIdent(fetch.name(), "fetch name")); } } + for (const tf2xla::Variable& variable : config.variable()) { + if (!variable.name().empty()) { + TF_RETURN_IF_ERROR(ValidateCppIdent(variable.name(), "variable name")); + } else { + TF_RETURN_IF_ERROR( + ValidateCppIdent(variable.node_name(), "variable name")); + } + } return Status::OK(); } @@ -339,9 +394,10 @@ Status GenerateHeader(const CodegenOpts& opts, const tf2xla::Config& config, std::vector buffer_infos_for_temps = ExtractTempBufferInfos(buffer_infos); const xla::ProgramShapeProto& ps = compile_result.program_shape; - string methods_arg, methods_result; + string methods_arg, methods_result, methods_variable; TF_RETURN_IF_ERROR(GenArgMethods(config, ps, compile_result, &methods_arg)); TF_RETURN_IF_ERROR(GenResultMethods(config, ps, &methods_result)); + TF_RETURN_IF_ERROR(GenVariableMethods(config, ps, &methods_variable)); const size_t arg_bytes_aligned = cpu_function_runtime::AlignedBufferBytes( buffer_infos_for_args.data(), buffer_infos_for_args.size(), /*allocate_entry_params=*/true); @@ -384,8 +440,9 @@ Status GenerateHeader(const CodegenOpts& opts, const tf2xla::Config& config, // calling HloProfilePrinter::profile_counters_size. const string assign_profile_counters_size = opts.gen_hlo_profile_printer_data - ? "data->set_profile_counters_size(" - "data->hlo_profile_printer_data()->profile_counters_size());" + ? "set_static_data_profile_counters_size(data, " + "get_static_data_hlo_profile_printer_data(data)->" + "profile_counters_size());" : ""; // Use a poor-man's text templating mechanism; first populate the full header @@ -449,7 +506,7 @@ extern "C" void {{ENTRY}}( // arg bytes aligned: {{ARG_BYTES_ALIGNED}} // temp bytes total: {{TEMP_BYTES_TOTAL}} // temp bytes aligned: {{TEMP_BYTES_ALIGNED}} -class {{CLASS}} : public tensorflow::XlaCompiledCpuFunction { +class {{CLASS}} final : public tensorflow::XlaCompiledCpuFunction { public: // Number of input arguments for the compiled computation. static constexpr size_t kNumArgs = {{ARG_NUM}}; @@ -464,16 +521,17 @@ class {{CLASS}} : public tensorflow::XlaCompiledCpuFunction { static XlaCompiledCpuFunction::StaticData* kStaticData = [](){ XlaCompiledCpuFunction::StaticData* data = new XlaCompiledCpuFunction::StaticData; - data->set_raw_function({{ENTRY}}); - data->set_buffer_infos(BufferInfos()); - data->set_num_buffers(kNumBuffers); - data->set_arg_index_table(ArgIndexToBufferIndex()); - data->set_num_args(kNumArgs); - data->set_result_index(kResultIndex); - data->set_arg_names(StaticArgNames()); - data->set_result_names(StaticResultNames()); - data->set_program_shape(StaticProgramShape()); - data->set_hlo_profile_printer_data(StaticHloProfilePrinterData()); + set_static_data_raw_function(data, {{ENTRY}}); + set_static_data_buffer_infos(data, BufferInfos()); + set_static_data_num_buffers(data, kNumBuffers); + set_static_data_arg_index_table(data, ArgIndexToBufferIndex()); + set_static_data_num_args(data, kNumArgs); + set_static_data_result_index(data, kResultIndex); + set_static_data_arg_names(data, StaticArgNames()); + set_static_data_result_names(data, StaticResultNames()); + set_static_data_program_shape(data, StaticProgramShape()); + set_static_data_hlo_profile_printer_data( + data, StaticHloProfilePrinterData()); {{ASSIGN_PROFILE_COUNTERS_SIZE}} return data; }(); @@ -521,6 +579,21 @@ class {{CLASS}} : public tensorflow::XlaCompiledCpuFunction { // buffers are managed internally, and may change after each call to Run. {{METHODS_RESULT}} + // Methods for managing variable buffers. Buffers are in row-major order. The + // input and output buffers may or may not be identical. + // + // void set_var_X_data(T* data) + // Sets the buffer for variable X. + // + // T* var_X_data() + // Returns the buffer of type T for variable X. + // + // T& var_X(...dim indices...) + // Returns a reference to the value of type T for variable X, + // with dim indices specifying which value. No bounds checking is performed + // on dim indices. +{{METHODS_VARIABLE}} + private: // Number of buffers for the compiled computation. static constexpr size_t kNumBuffers = {{NUM_BUFFERS}}; @@ -587,6 +660,7 @@ class {{CLASS}} : public tensorflow::XlaCompiledCpuFunction { include_hlo_profile_printer_data_proto}, {"{{METHODS_ARG}}\n", methods_arg}, {"{{METHODS_RESULT}}\n", methods_result}, + {"{{METHODS_VARIABLE}}\n", methods_variable}, {"{{NS_END}}\n", ns_end}, {"{{NS_START}}\n", ns_start}, {"{{PROGRAM_SHAPE}}", xla::ShapeUtil::HumanString(xla::ProgramShape(ps))}, diff --git a/tensorflow/compiler/aot/codegen_test.cc b/tensorflow/compiler/aot/codegen_test.cc index c1788ca32a1d099284eeb870f9513891051fd29e..5580e55b691bd10698b63d86bc0194b25da743b9 100644 --- a/tensorflow/compiler/aot/codegen_test.cc +++ b/tensorflow/compiler/aot/codegen_test.cc @@ -22,6 +22,8 @@ limitations under the License. #include "absl/strings/string_view.h" #include "llvm/Support/TargetSelect.h" #include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/core/framework/tensor_shape.pb.h" +#include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/lib/io/path.h" @@ -172,6 +174,15 @@ TEST(CodegenTest, Golden) { tf2xla::Fetch* fetch = config.add_fetch(); fetch->mutable_id()->set_node_name("fetch0"); fetch->set_name("myfetch"); + tf2xla::Variable* variable = config.add_variable(); + variable->set_node_name("myvar"); + variable->mutable_shape()->add_dim()->set_size(1); + variable->set_type(DT_FLOAT); + tf2xla::Variable* variable2 = config.add_variable(); + variable2->set_node_name("my/var"); + variable2->set_name("myvar2"); + variable2->mutable_shape()->add_dim()->set_size(5); + variable2->set_type(DT_INT32); CompileResult compile_result; compile_result.aot.reset(new xla::cpu::CpuAotCompilationResult( {}, @@ -186,9 +197,14 @@ TEST(CodegenTest, Golden) { { xla::ShapeUtil::MakeShape(xla::F32, {1, 2}), xla::ShapeUtil::MakeShape(xla::S64, {3, 4}), + xla::ShapeUtil::MakeShape(xla::F32, {1}), + xla::ShapeUtil::MakeShape(xla::S32, {5}), }, - xla::ShapeUtil::MakeTupleShape( - {xla::ShapeUtil::MakeShape(xla::U32, {5, 6})})) + xla::ShapeUtil::MakeTupleShape({ + xla::ShapeUtil::MakeShape(xla::U32, {5, 6}), + xla::ShapeUtil::MakeShape(xla::F32, {1}), + xla::ShapeUtil::MakeShape(xla::S32, {5}), + })) .ToProto(); compile_result.entry_point = "entry_point"; compile_result.pointer_size = 8; diff --git a/tensorflow/compiler/aot/codegen_test_h.golden b/tensorflow/compiler/aot/codegen_test_h.golden index 7fbf83af330285574023a5eb9e84f1e3ee82fbb9..b5f33d690d492489e9090786cd341e035ae7ca15 100644 --- a/tensorflow/compiler/aot/codegen_test_h.golden +++ b/tensorflow/compiler/aot/codegen_test_h.golden @@ -52,14 +52,14 @@ namespace bar { // is guaranteed that no thread may call a non-const method. // // The logical function signature is: -// ((unknown): f32[1,2], (unknown): s64[3,4]) -> (u32[5,6]) +// ((unknown): f32[1,2], (unknown): s64[3,4], (unknown): f32[1], (unknown): s32[5]) -> (u32[5,6], f32[1], s32[5]) // // Memory stats: // arg bytes total: 104 // arg bytes aligned: 192 // temp bytes total: 126 // temp bytes aligned: 320 -class MyClass : public tensorflow::XlaCompiledCpuFunction { +class MyClass final : public tensorflow::XlaCompiledCpuFunction { public: // Number of input arguments for the compiled computation. static constexpr size_t kNumArgs = 2; @@ -74,16 +74,17 @@ class MyClass : public tensorflow::XlaCompiledCpuFunction { static XlaCompiledCpuFunction::StaticData* kStaticData = [](){ XlaCompiledCpuFunction::StaticData* data = new XlaCompiledCpuFunction::StaticData; - data->set_raw_function(entry_point); - data->set_buffer_infos(BufferInfos()); - data->set_num_buffers(kNumBuffers); - data->set_arg_index_table(ArgIndexToBufferIndex()); - data->set_num_args(kNumArgs); - data->set_result_index(kResultIndex); - data->set_arg_names(StaticArgNames()); - data->set_result_names(StaticResultNames()); - data->set_program_shape(StaticProgramShape()); - data->set_hlo_profile_printer_data(StaticHloProfilePrinterData()); + set_static_data_raw_function(data, entry_point); + set_static_data_buffer_infos(data, BufferInfos()); + set_static_data_num_buffers(data, kNumBuffers); + set_static_data_arg_index_table(data, ArgIndexToBufferIndex()); + set_static_data_num_args(data, kNumArgs); + set_static_data_result_index(data, kResultIndex); + set_static_data_arg_names(data, StaticArgNames()); + set_static_data_result_names(data, StaticResultNames()); + set_static_data_program_shape(data, StaticProgramShape()); + set_static_data_hlo_profile_printer_data( + data, StaticHloProfilePrinterData()); return data; }(); @@ -213,6 +214,58 @@ class MyClass : public tensorflow::XlaCompiledCpuFunction { result_data(0)))[dim0][dim1]; } + // Methods for managing variable buffers. Buffers are in row-major order. The + // input and output buffers may or may not be identical. + // + // void set_var_X_data(T* data) + // Sets the buffer for variable X. + // + // T* var_X_data() + // Returns the buffer of type T for variable X. + // + // T& var_X(...dim indices...) + // Returns a reference to the value of type T for variable X, + // with dim indices specifying which value. No bounds checking is performed + // on dim indices. + + void set_var_myvar_input_data(float* data) { + set_arg_data(2, data); + } + + void set_var_myvar2_input_data(tensorflow::int32* data) { + set_arg_data(3, data); + } + + float* var_myvar_result_data() { + return static_cast(result_data(1)); + } + float& var_myvar_result() { + return (*static_cast( + result_data(1)))[0]; + } + const float* var_myvar_result_data() const { + return static_cast(result_data(1)); + } + const float& var_myvar_result() const { + return (*static_cast( + result_data(1)))[0]; + } + + tensorflow::int32* var_myvar2_result_data() { + return static_cast(result_data(2)); + } + tensorflow::int32& var_myvar2_result(size_t dim0) { + return (*static_cast( + result_data(2)))[dim0]; + } + const tensorflow::int32* var_myvar2_result_data() const { + return static_cast(result_data(2)); + } + const tensorflow::int32& var_myvar2_result(size_t dim0) const { + return (*static_cast( + result_data(2)))[dim0]; + } + private: // Number of buffers for the compiled computation. static constexpr size_t kNumBuffers = 6; @@ -256,7 +309,7 @@ class MyClass : public tensorflow::XlaCompiledCpuFunction { static const xla::ProgramShapeProto* StaticProgramShape() { static const xla::ProgramShapeProto* kShape = []() { xla::ProgramShapeProto* proto = new xla::ProgramShapeProto; - proto->ParseFromArray(&__tfcompile_foo_bar_MyClass_ProgramShapeProto_protobuf_array_contents[0], 64); + proto->ParseFromArray(&__tfcompile_foo_bar_MyClass_ProgramShapeProto_protobuf_array_contents[0], 132); return proto; }(); return kShape; diff --git a/tensorflow/compiler/aot/codegen_test_o.golden b/tensorflow/compiler/aot/codegen_test_o.golden index 7f7b96428572705f30144e6c95cd4cf9c44ce2a3..2884597abcf29583e6192296b0e4ce6825d7c01a 100644 Binary files a/tensorflow/compiler/aot/codegen_test_o.golden and b/tensorflow/compiler/aot/codegen_test_o.golden differ diff --git a/tensorflow/compiler/aot/compile.cc b/tensorflow/compiler/aot/compile.cc index 9fc223bdc7c0e207ce2005cb86250aa77e709df8..0e46a9f5e9d68fa2174f7bd9b9fa7c3a82dfb715 100644 --- a/tensorflow/compiler/aot/compile.cc +++ b/tensorflow/compiler/aot/compile.cc @@ -108,10 +108,13 @@ Status CompileGraph(const GraphDef& graph_def, const tf2xla::Config& config, computation.Snapshot()); // Serialize the HloSnapshot deterministically so that all the outputs of a // tf_library genrule are deterministic. - string proto; - TF_RET_CHECK(SerializeToStringDeterministic(*module, &proto)); + const size_t size = module->ByteSizeLong(); + auto serialized = absl::make_unique(size); + TF_RET_CHECK( + SerializeToBufferDeterministic(*module, serialized.get(), size)); TF_RETURN_IF_ERROR( - WriteStringToFile(Env::Default(), flags.out_session_module, proto)); + WriteStringToFile(Env::Default(), flags.out_session_module, + absl::string_view(serialized.get(), size))); } xla::cpu::CpuAotCompilationOptions aot_opts( flags.target_triple, flags.target_cpu, flags.target_features, diff --git a/tensorflow/compiler/aot/tests/BUILD b/tensorflow/compiler/aot/tests/BUILD index 10fa33ab5e84dcbc1629bee6214e8969046f19c2..444264ba6e1f59c33551796025ba845c62c02d43 100644 --- a/tensorflow/compiler/aot/tests/BUILD +++ b/tensorflow/compiler/aot/tests/BUILD @@ -69,6 +69,7 @@ genrule( "test_graph_tfmatmulandadd.pb", "test_graph_tfsplits.pb", "test_graph_tftop_k.pb", + "test_graph_tfvariable.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 @@ -222,6 +223,17 @@ tf_library( ], ) +tf_library( + name = "test_graph_tfvariable", + testonly = 1, + config = "test_graph_tfvariable.config.pbtxt", + cpp_class = "VariableComp", + graph = "test_graph_tfvariable.pb", + tags = [ + "manual", + ], +) + tf_cc_test( name = "tfcompile_test", srcs = ["tfcompile_test.cc"], @@ -241,6 +253,7 @@ tf_cc_test( ":test_graph_tfmatmulandadd_with_profiling", ":test_graph_tfsplits", ":test_graph_tftop_k", + ":test_graph_tfvariable", "//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 64b861a73091642b03573543a5c55618bf33915d..42f8812def0503824416d92daa2db71a64c3db88 100644 --- a/tensorflow/compiler/aot/tests/make_test_graphs.py +++ b/tensorflow/compiler/aot/tests/make_test_graphs.py @@ -50,7 +50,7 @@ def tfadd_with_ckpt(out_dir): y = variables.VariableV1(constant_op.constant([0]), name='y_saved') math_ops.add(x, y, name='x_y_sum') - init_op = variables.initialize_all_variables() + init_op = variables.global_variables_initializer() saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V1) with session.Session() as sess: sess.run(init_op) @@ -65,7 +65,7 @@ def tfadd_with_ckpt_saver(out_dir): y = variables.VariableV1(constant_op.constant([0]), name='y_saved') math_ops.add(x, y, name='x_y_sum') - init_op = variables.initialize_all_variables() + init_op = variables.global_variables_initializer() saver = saver_lib.Saver(name='abcprefix', write_version=saver_pb2.SaverDef.V1) with session.Session() as sess: sess.run(init_op) @@ -149,6 +149,14 @@ def tftop_k(_): array_ops.identity(output[1], name='indices') +def tfvariable(_): + x = variables.Variable(1000.0, name='x') + old_x = x.value() + with ops.control_dependencies([old_x]): + new_x = x.assign_add(42.0) + array_ops.stack([old_x, new_x], name='result') + + def write_graph(build_graph, out_dir): """Build a graph using build_graph and write it out.""" g = ops.Graph() @@ -171,6 +179,7 @@ def main(_): write_graph(tfmatmulandadd, FLAGS.out_dir) write_graph(tfsplits, FLAGS.out_dir) write_graph(tftop_k, FLAGS.out_dir) + write_graph(tfvariable, FLAGS.out_dir) if __name__ == '__main__': diff --git a/tensorflow/compiler/aot/tests/test_graph_tfvariable.config.pbtxt b/tensorflow/compiler/aot/tests/test_graph_tfvariable.config.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..9b4c4215a330b014f595edde001aba73ad7d8263 --- /dev/null +++ b/tensorflow/compiler/aot/tests/test_graph_tfvariable.config.pbtxt @@ -0,0 +1,12 @@ +# Text form of tensorflow.tf2xla.Config proto. +fetch { + id { node_name: "result" } +} + +variable { + node_name: "x" + shape { + dim { size: 1 } + } + type: DT_FLOAT +} diff --git a/tensorflow/compiler/aot/tests/tfcompile_test.cc b/tensorflow/compiler/aot/tests/tfcompile_test.cc index 4dd79e5882d7da61be029735ef2b165908c599f9..5f9316f3933713e12fc5960b9adfecc6e9bd99b5 100644 --- a/tensorflow/compiler/aot/tests/tfcompile_test.cc +++ b/tensorflow/compiler/aot/tests/tfcompile_test.cc @@ -30,6 +30,7 @@ limitations under the License. #include "tensorflow/compiler/aot/tests/test_graph_tfmatmulandadd_with_profiling.h" #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/xla/service/hlo_profile_printer.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" @@ -473,6 +474,28 @@ TEST(TFCompileTest, TopK) { EXPECT_EQ(expected_indices[1], fn.result1(1)); } +TEST(TFCompileTest, Variable) { + Eigen::ThreadPool tp(1); + Eigen::ThreadPoolDevice device(&tp, tp.NumThreads()); + + VariableComp fn; + float x = 23; + fn.set_var_x_input_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(x, 23); + x = fn.var_x_result(); + fn.Run(); + EXPECT_EQ(fn.result0(0, 0), 65); + EXPECT_EQ(fn.result0(1, 0), 107); + EXPECT_EQ(fn.var_x_result(), 107); +} + TEST(TFCompileTest, AssertEqAndReturnDiff) { // Assert is converted into a no-op in XLA, so there is no failure even if the // two args are different. diff --git a/tensorflow/compiler/aot/tfcompile.bzl b/tensorflow/compiler/aot/tfcompile.bzl index 4051664c24cacad4a2d151ad3ac9009015900609..2abe3e29b78dbbe719637b13418704acc213d050 100644 --- a/tensorflow/compiler/aot/tfcompile.bzl +++ b/tensorflow/compiler/aot/tfcompile.bzl @@ -207,7 +207,7 @@ def tf_library( # # Note that setting the local=1 attribute on a *test target* causes the # test infrastructure to skip that test. However this is a genrule, not - # a test target, and runs with --genrule_strategy=forced_forge, meaning + # a test target, and runs with --strategy=Genrule=forced_forge, meaning # the local=1 attribute is ignored, and the genrule is still run. # # https://www.bazel.io/versions/master/docs/be/general.html#genrule diff --git a/tensorflow/compiler/aot/tfcompile_main.cc b/tensorflow/compiler/aot/tfcompile_main.cc index d548de8c44285f6d21dd778db464a31e1b19645b..0b6ab7e723d6e3a55da2f1c30b75f44cbdaa75bb 100644 --- a/tensorflow/compiler/aot/tfcompile_main.cc +++ b/tensorflow/compiler/aot/tfcompile_main.cc @@ -136,6 +136,10 @@ int main(int argc, char** argv) { tensorflow::string usage = tensorflow::tfcompile::kUsageHeader; usage += tensorflow::Flags::Usage(argv[0], flag_list); + if (argc > 1 && absl::string_view(argv[1]) == "--help") { + std::cerr << usage << "\n"; + return 0; + } bool parsed_flags_ok = tensorflow::Flags::Parse(&argc, argv, flag_list); QCHECK(parsed_flags_ok) << "\n" << usage; diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index b9a87ba296abfc6b9d9aaeff3b3e26678e4e1b94..4e3229bc709d25195392afc84382a61703782255 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -175,12 +175,22 @@ cc_library( "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/service:stream_pool", + "//tensorflow/core:array_ops_op_lib", + "//tensorflow/core:control_flow_ops_op_lib", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", + "//tensorflow/core:dataset_ops_op_lib", "//tensorflow/core:framework", + "//tensorflow/core:functional_ops_op_lib", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", + "//tensorflow/core:math_ops_op_lib", + "//tensorflow/core:nn_ops_op_lib", + "//tensorflow/core:no_op_op_lib", "//tensorflow/core:protos_all_cc", + "//tensorflow/core:resource_variable_ops_op_lib", + "//tensorflow/core:sendrecv_ops_op_lib", + "//tensorflow/core:state_ops_op_lib", "//tensorflow/core:stream_executor_no_cuda", "//tensorflow/core/kernels:cast_op", "//tensorflow/core/kernels:constant_op", @@ -272,7 +282,6 @@ cc_library( "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:dump_graph", "//tensorflow/compiler/tf2xla:xla_compiler", - "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:local_client", @@ -455,7 +464,6 @@ cc_library( "//tensorflow/compiler/tf2xla:tf2xla_util", "//tensorflow/core:framework", "//tensorflow/core:graph", - "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", @@ -634,10 +642,10 @@ tf_cc_test( "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", + "//tensorflow/core:session_options", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", - "//tensorflow/core/grappler/optimizers/data:graph_utils", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", diff --git a/tensorflow/compiler/jit/build_xla_ops_pass.cc b/tensorflow/compiler/jit/build_xla_ops_pass.cc index 9f4042630edaec1b9519b6434d859a48372e8b15..285b1efa53d91922c9fa161cfd2de34e1434d0c4 100644 --- a/tensorflow/compiler/jit/build_xla_ops_pass.cc +++ b/tensorflow/compiler/jit/build_xla_ops_pass.cc @@ -115,6 +115,13 @@ void MergeOutgoingControlEdges(const Scope& s, Node* old_node, Node* new_node) { return; } + if (ctrl_edges.size() == 1 && ctrl_edges.front()->dst()->IsSink()) { + // Avoid creating a Merge node if we can just add an edge to _SINK + // instead. + s.graph()->AddControlEdge(new_node, s.graph()->sink_node()); + return; + } + // We can't merge control edges directly so we instead first "convert" them to // normal values that can be merged, merge the values and then "convert" the // merged value back into control. diff --git a/tensorflow/compiler/jit/build_xla_ops_pass_test.cc b/tensorflow/compiler/jit/build_xla_ops_pass_test.cc index 48a23a4c1711ac88a329723c46559112d5a39dbd..c14c7465c55b7d350d6b3a6853cef6692140ce78 100644 --- a/tensorflow/compiler/jit/build_xla_ops_pass_test.cc +++ b/tensorflow/compiler/jit/build_xla_ops_pass_test.cc @@ -24,7 +24,6 @@ limitations under the License. #include "tensorflow/compiler/jit/node_matchers.h" #include "tensorflow/core/common_runtime/device_factory.h" #include "tensorflow/core/graph/algorithm.h" -#include "tensorflow/core/grappler/optimizers/data/graph_utils.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/test.h" @@ -69,6 +68,8 @@ Status BuildXlaOps(const Scope& s, std::unique_ptr* result) { } } + FixupSourceAndSinkEdges(graph.get()); + GraphOptimizationPassOptions opt_options; opt_options.graph = &graph; BuildXlaOpsPass pass(/*enable_lazy_compilation=*/true); @@ -224,5 +225,23 @@ TEST_F(BuildXlaOpsTest, OnXlaDevice) { ASSERT_NE(write_op_new, nullptr); EXPECT_THAT(write_op_new, assign_var); } + +TEST_F(BuildXlaOpsTest, NoExtraMergeForEdgeToSink) { + Scope root = Scope::NewRootScope().ExitOnError(); + + FunctionDefLibrary flib_def = + CreateFunctionDefLibWithConstFunction("cluster_0"); + TF_ASSERT_OK(root.graph()->AddFunctionLibrary(flib_def)); + Node* call; + TF_ASSERT_OK(MakeXlaCompiledKernel(root.graph(), "cluster_0", "C", &call)); + + std::unique_ptr graph; + TF_ASSERT_OK(BuildXlaOps(root, &graph)); + + Node* sink_node = graph->sink_node(); + EXPECT_THAT(sink_node, NodeWith(CtrlDeps(NodeWith(Op("_XlaRun")), + NodeWith(Op("cluster_0")), + NodeWith(Op("NoOp"))))); +} } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/jit/deadness_analysis.cc b/tensorflow/compiler/jit/deadness_analysis.cc index 0562838f628c66b1eb03af9d2a5139c01dca31c5..4397eea9af266cbd0392f08323e59077c9395150 100644 --- a/tensorflow/compiler/jit/deadness_analysis.cc +++ b/tensorflow/compiler/jit/deadness_analysis.cc @@ -20,7 +20,10 @@ limitations under the License. #include "absl/strings/str_join.h" #include "tensorflow/compiler/jit/deadness_analysis_internal.h" #include "tensorflow/compiler/jit/xla_cluster_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/core/framework/tensor.pb.h" #include "tensorflow/core/graph/algorithm.h" +#include "tensorflow/core/graph/control_flow.h" #include "tensorflow/core/graph/tensor_id.h" #include "tensorflow/core/lib/hash/hash.h" @@ -110,7 +113,11 @@ class Predicate { enum class Kind { kAnd, kOr, kNot, kAndRecurrence, kSymbol }; virtual string ToString() const = 0; - int64 hash() const { return hash_; } + + // An ID assigned to the Predicate at construction time. Conceptually like a + // pointer, except that it is stable across runs. + int64 id() const { return id_; } + virtual absl::Span GetOperands() const = 0; virtual Kind kind() const = 0; @@ -123,29 +130,19 @@ class Predicate { static void Visit(Predicate* p, const FunctionTy& func); protected: - explicit Predicate(int64 hash) : hash_(hash) {} + explicit Predicate(int64 id) : id_(id) {} private: - const int64 hash_; + const int64 id_; TF_DISALLOW_COPY_AND_ASSIGN(Predicate); }; -int64 HashPredicateSequence(Predicate::Kind kind, - absl::Span preds) { - int64 hash = ::tensorflow::hash()(kind); - for (Predicate* pred : preds) { - hash = Hash64Combine(hash, pred->hash()); - } - return hash; -} - // Represents a logical conjunction of a set of predicates. class AndPredicate : public Predicate { public: - explicit AndPredicate(std::vector operands) - : Predicate(HashPredicateSequence(Kind::kAnd, operands)), - operands_(std::move(operands)) {} + explicit AndPredicate(int64 id, std::vector operands) + : Predicate(id), operands_(std::move(operands)) {} string ToString() const override { if (operands().empty()) { @@ -174,9 +171,8 @@ class AndPredicate : public Predicate { // Represents a logical disjunction of a set of predicates. class OrPredicate : public Predicate { public: - explicit OrPredicate(std::vector operands) - : Predicate(HashPredicateSequence(Kind::kOr, operands)), - operands_(std::move(operands)) {} + explicit OrPredicate(int64 id, std::vector operands) + : Predicate(id), operands_(std::move(operands)) {} string ToString() const override { if (operands().empty()) { @@ -204,9 +200,8 @@ class OrPredicate : public Predicate { // Represents a logical negation of a set of predicates. class NotPredicate : public Predicate { public: - explicit NotPredicate(Predicate* operand) - : Predicate(HashPredicateSequence(Kind::kNot, {operand})), - operands_({operand}) {} + explicit NotPredicate(int64 id, Predicate* operand) + : Predicate(id), operands_({operand}) {} string ToString() const override { return absl::StrCat("~", operand()->ToString()); @@ -222,29 +217,38 @@ class NotPredicate : public Predicate { std::array operands_; }; -// Represents an infinite list of predicates. +// Represents the liveness of an induction variable. For users inside the loop +// this represents the "current" liveness of the induction variable. For users +// outside the loop it represents the "last" liveness of the induction variable. +// +// More concretely, an and recurrence {S,&,X} represents the liveness of V +// in the following graph: // -// An AndRecurrence with start = S and step = X is printed as {S,&,X} and stands -// for the list of predicates: +// V = Merge(S', V_NextIt) +// V = Op(V, X') +// V_NextIt = NextIteration(V) // -// S, S & GenSym(X,1), S & GenSym(X,1) & GenSym(X,2), ... +// where Predicate(S') = S and Predicate(X') = X. // -// where GenSym(, ) renames every SymbolPredicate in -// by appending to it, in effect creating a "fresh" symbol. -// This means {P,&,Q} is not equal to "P on the first iteration; P&Q on -// subsequent iterations". +// `X` may contain symbolic predicates and the operations corresponding to these +// symbolic predicates are either in frame `loop` or outside it. The symbols +// that are inside frame `loop` are loop variant (i.e. can have different +// liveness in each loop iteration) and the symbols that are outside frame +// `loop` are loop invariant (i.e. have the same liveness across all +// iterations). class AndRecurrencePredicate : public Predicate { public: - explicit AndRecurrencePredicate(Predicate* start, Predicate* step) - : Predicate(HashPredicateSequence(Kind::kAndRecurrence, {start, step})), - operands_({start, step}) {} + explicit AndRecurrencePredicate(int64 id, Predicate* start, Predicate* step, + std::vector frame) + : Predicate(id), operands_({start, step}), frame_(std::move(frame)) {} Predicate* start() const { return operands_[0]; } Predicate* step() const { return operands_[1]; } + absl::Span frame() const { return frame_; } string ToString() const override { return absl::StrCat("{", start()->ToString(), ",&,", step()->ToString(), - "}"); + "}<", absl::StrJoin(frame(), ";"), ">"); } Kind kind() const override { return Kind::kAndRecurrence; } @@ -255,6 +259,7 @@ class AndRecurrencePredicate : public Predicate { private: std::array operands_; + std::vector frame_; }; // Represents an uninterpreted symbol in a logical predicate. @@ -264,8 +269,8 @@ class AndRecurrencePredicate : public Predicate { // symbols. class SymbolPredicate : public Predicate { public: - explicit SymbolPredicate(TensorId tensor_id, bool must_be_true) - : Predicate(Hash(tensor_id, must_be_true)), + explicit SymbolPredicate(int64 id, TensorId tensor_id, bool must_be_true) + : Predicate(id), tensor_id_(std::move(tensor_id)), must_be_true_(must_be_true) {} @@ -281,20 +286,13 @@ class SymbolPredicate : public Predicate { // "tensor_id() is live and evaluates to true". // // If `must_be_true()` is false then this SymbolPredicate represents the - // proposition "tensor_id() is live (and may evalutate to any value)" + // proposition "tensor_id() is live (and may evaluate to any value)" TensorId tensor_id() const { return tensor_id_; } bool must_be_true() const { return must_be_true_; } private: TensorId tensor_id_; bool must_be_true_; - - static int64 Hash(const TensorId tensor_id, bool must_be_true) { - return Hash64Combine( - ::tensorflow::hash()(must_be_true), - Hash64Combine(::tensorflow::hash()(Kind::kSymbol), - TensorId::Hasher{}(tensor_id))); - } }; template @@ -333,34 +331,58 @@ class PredicateFactory { } Predicate* MakeNotPredicate(Predicate* pred) { - SignatureForNot signature = pred; - auto it = interned_not_instances_.find(signature); - if (it == interned_not_instances_.end()) { - std::unique_ptr new_pred = Make(pred); - Predicate* new_pred_ptr = new_pred.get(); - interned_not_instances_.emplace(signature, std::move(new_pred)); - return new_pred_ptr; - } else { - return it->second.get(); + auto it = make_not_predicate_cache_.find(pred); + if (it != make_not_predicate_cache_.end()) { + return it->second; } + + Predicate* result = MakeNotPredicateImpl(pred); + + bool insert_successful = + make_not_predicate_cache_.insert({pred, result}).second; + (void)insert_successful; + DCHECK(insert_successful); + + return result; } - Predicate* MakeAndRecurrencePredicate(Predicate* start, Predicate* step) { - auto it = interned_and_rec_instances_.find({start, step}); + Predicate* MakeAndRecurrencePredicate(Predicate* start, Predicate* step, + std::vector frame) { + SignatureForAndRec signature(start, step, std::move(frame)); + auto it = interned_and_rec_instances_.find(signature); if (it != interned_and_rec_instances_.end()) { return it->second.get(); } - std::unique_ptr new_pred = - Make(start, step); + std::unique_ptr new_pred = Make( + std::get<0>(signature), std::get<1>(signature), std::get<2>(signature)); Predicate* new_pred_ptr = new_pred.get(); - CHECK(interned_and_rec_instances_ - .emplace(SignatureForAndRec(start, step), std::move(new_pred)) - .second); + bool inserted = + interned_and_rec_instances_.emplace(signature, std::move(new_pred)) + .second; + (void)inserted; + DCHECK(inserted); return new_pred_ptr; } - Predicate* MakeSymbolPredicate(TensorId tensor_id, bool must_be_true) { + Status MakeSymbolPredicate(Node* node, int output_idx, bool must_be_true, + Predicate** predicate) { + TensorId tensor_id(node->name(), output_idx); + + bool is_boolean_tensor = node->output_type(tensor_id.index()) == DT_BOOL; + TF_RET_CHECK(!must_be_true || is_boolean_tensor); + + if (node->type_string() == "Const" && must_be_true) { + const TensorProto* proto = nullptr; + TF_RETURN_IF_ERROR(GetNodeAttr(node->def(), "value", &proto)); + + Tensor tensor(proto->dtype()); + TF_RET_CHECK(tensor.FromProto(*proto)); + + *predicate = tensor.scalar()() ? MakeTrue() : MakeFalse(); + return Status::OK(); + } + SignatureForSymbol signature = {tensor_id, must_be_true}; auto it = interned_symbol_instances_.find(signature); if (it == interned_symbol_instances_.end()) { @@ -369,20 +391,70 @@ class PredicateFactory { Predicate* new_pred_ptr = new_pred.get(); interned_symbol_instances_.emplace(std::move(signature), std::move(new_pred)); - return new_pred_ptr; + *predicate = new_pred_ptr; } else { - return it->second.get(); + *predicate = it->second.get(); } + + return Status::OK(); } Predicate* MakeTrue() { return MakeAndPredicate({}); } Predicate* MakeFalse() { return MakeOrPredicate({}); } + ~PredicateFactory() { + DCHECK_EQ(stack_depth_, 0) << "Unnested IncrementStackDepth?"; + } + private: + Predicate* MakeNotPredicateImpl(Predicate* pred) { + IncrementStackDepth stack_frame(this); + if (!stack_frame.HasOverflowed()) { + if (Predicate* simplified = SimplifyUsingDeMorgan(pred)) { + return simplified; + } + + // ~~A => A + if (auto* not_pred = dynamic_cast(pred)) { + return not_pred->operand(); + } + } + + SignatureForNot signature = pred; + auto it = interned_not_instances_.find(signature); + if (it == interned_not_instances_.end()) { + std::unique_ptr new_pred = Make(pred); + Predicate* new_pred_ptr = new_pred.get(); + interned_not_instances_.emplace(signature, std::move(new_pred)); + return new_pred_ptr; + } else { + return it->second.get(); + } + } + + Predicate* SimplifyUsingDeMorgan(Predicate* pred) { + // ~(A & B & C & ...) => ~A | ~B | ~C | ~... + // ~(A | B | C | ...) -> ~A & ~B & ~C & ~... + Predicate::Kind kind = pred->kind(); + + if (kind == Predicate::Kind::kAnd || kind == Predicate::Kind::kOr) { + std::vector new_operands; + absl::c_transform(pred->GetOperands(), std::back_inserter(new_operands), + [&](Predicate* p) { return MakeNotPredicate(p); }); + return kind == Predicate::Kind::kOr ? MakeAndPredicate(new_operands) + : MakeOrPredicate(new_operands); + } + + return nullptr; + } + template std::unique_ptr Make(Args&&... args) { + // If we ever expose the Predicate class outside this .cc file then we may + // want to make this hard to misuse (by accidentally passing in an arbitrary + // integer to the Predicate constructor for instance). return std::unique_ptr( - new PredicateT(std::forward(args)...)); + new PredicateT(id_counter_++, std::forward(args)...)); } Predicate* MakeAndOrImpl(absl::Span operands, bool is_and); @@ -402,7 +474,8 @@ class PredicateFactory { using SignatureForAndOr = std::pair>; using SignatureForNot = Predicate*; - using SignatureForAndRec = std::pair; + using SignatureForAndRec = + std::tuple>; using SignatureForSymbol = std::pair; struct HashSignatureForAndOr { @@ -422,6 +495,36 @@ class PredicateFactory { } }; + // Used to limit recursion to avoid blowing up the stack and cap compile time. + class IncrementStackDepth { + public: + explicit IncrementStackDepth(PredicateFactory* parent) : parent_(parent) { + parent_->stack_depth_++; + } + + bool HasOverflowed() const { + const int kMaxStackDepth = 8; + return parent_->stack_depth_ >= kMaxStackDepth; + } + + ~IncrementStackDepth() { parent_->stack_depth_--; } + + private: + PredicateFactory* parent_; + }; + + // A cache for the MakeNotPredicate function. + // + // NB! This is *not* the same as `interned_not_instances_`. + // `interned_not_instances_` maps ensures pointer identity for `NotPredicate` + // instances, i.e., it ensures there at most one instance of Not(predicate) + // for any given predicate whereas `make_not_predicate_cache_` simply caches + // the result of the `MakeNotPredicate` function. The values in + // `interned_not_instances_` are always instance of `NotPredicate` whereas the + // values in `make_not_predicate_cache_` may not be (for instance it will map + // Not(Not(A)) to A). + absl::flat_hash_map make_not_predicate_cache_; + absl::flat_hash_map, HashSignatureForAndOr> interned_and_or_instances_; @@ -432,13 +535,15 @@ class PredicateFactory { absl::flat_hash_map, HashSignatureForSymbol> interned_symbol_instances_; + int64 id_counter_ = 0; + int stack_depth_ = 0; }; Predicate* PredicateFactory::MakeInternedAndOr( std::vector simplified_ops, Predicate::Kind pred_kind) { std::stable_sort( simplified_ops.begin(), simplified_ops.end(), - [](Predicate* a, Predicate* b) { return a->hash() < b->hash(); }); + [](Predicate* a, Predicate* b) { return a->id() < b->id(); }); auto it = interned_and_or_instances_.find({pred_kind, simplified_ops}); if (it != interned_and_or_instances_.end()) { @@ -466,6 +571,13 @@ Predicate* PredicateFactory::MakeAndOrImpl( absl::Span operands, bool is_and) { Predicate::Kind pred_kind = is_and ? Predicate::Kind::kAnd : Predicate::Kind::kOr; + + IncrementStackDepth stack_frame(this); + if (stack_frame.HasOverflowed()) { + return MakeInternedAndOr( + std::vector(operands.begin(), operands.end()), pred_kind); + } + Predicate::Kind other_pred_kind = is_and ? Predicate::Kind::kOr : Predicate::Kind::kAnd; absl::flat_hash_set simplified_ops_set; @@ -494,16 +606,31 @@ Predicate* PredicateFactory::MakeAndOrImpl( // Simplify "A&~A=>False" and "A|~A=>True". absl::flat_hash_set negated_ops; - for (Predicate* op : simplified_ops) { - if (op->kind() == Predicate::Kind::kNot) { - negated_ops.insert(dynamic_cast(*op).operand()); - } - } - for (Predicate* op : simplified_ops) { if (negated_ops.count(op)) { + // Simple case: + // + // A & ~A & ... == False + // A | ~A | ... == True return is_and ? MakeFalse() : MakeTrue(); } + + Predicate* negated_op = MakeNotPredicate(op); + if (negated_op->kind() == pred_kind) { + // Slightly more complicated case: + // + // (~A | ~B | ~C) & A & B & C & ... == + // ~(A & B & C) & (A & B & C) & ... == False + // + // (~A & ~B & ~C) | A | B | C | ... == + // ~(A | B | C) | (A | B | C) | ... == True + if (absl::c_all_of(negated_op->GetOperands(), [&](Predicate* p) { + return simplified_ops_set.contains(p); + })) { + return is_and ? MakeFalse() : MakeTrue(); + } + } + negated_ops.insert(negated_op); } // If all ops contain the same subop, then factor it out thanks to the @@ -619,6 +746,7 @@ class DeadnessAnalysisImpl : public DeadnessAnalysis { const Graph& graph_; absl::flat_hash_map predicate_map_; PredicateFactory predicate_factory_; + std::vector control_flow_info_; bool vlog_; }; @@ -661,9 +789,12 @@ Status DeadnessAnalysisImpl::HandleSwitch(Node* n, TF_RETURN_IF_ERROR(GetInputPreds(n, EdgeKind::kDataAndControl, &input_preds)); const Edge* pred_edge; TF_RETURN_IF_ERROR(n->input_edge(1, &pred_edge)); - Predicate* true_switch = predicate_factory_.MakeSymbolPredicate( - TensorId(pred_edge->src()->name(), pred_edge->src_output()), - /*must_be_true=*/true); + + Predicate* true_switch; + TF_RETURN_IF_ERROR(predicate_factory_.MakeSymbolPredicate( + pred_edge->src(), pred_edge->src_output(), + /*must_be_true=*/true, &true_switch)); + Predicate* false_switch = predicate_factory_.MakeNotPredicate(true_switch); // Output 0 is alive iff all inputs are alive and the condition is false. @@ -761,6 +892,23 @@ Predicate* DeduceStepPredicate(PredicateFactory* predicate_factory, return found_sym ? predicate_factory->MakeAndPredicate(and_ops) : nullptr; } + +Status GetFullFrame(const Node* n, absl::Span cfi_infos, + std::vector* frame) { + int depth = 0; + for (const ControlFlowInfo* cfi_iter = &cfi_infos[n->id()]; !n->IsSource(); + n = cfi_iter->parent_frame, cfi_iter = &cfi_infos[n->id()]) { + frame->push_back(cfi_iter->frame_name); + + if (depth++ > 5000) { + return errors::Internal( + "Frame of depth > 5000: Probably malformed graph or a bug in " + "BuildControlFlowInfo"); + } + } + + return Status::OK(); +} } // namespace Status DeadnessAnalysisImpl::HandleMerge(Node* n, @@ -783,8 +931,10 @@ Status DeadnessAnalysisImpl::HandleMerge(Node* n, if (has_unvisited_backedge) { // We're visiting this merge for the first time and it has an unvisited // backedge. - Predicate* input_data_pred = predicate_factory_.MakeSymbolPredicate( - TensorId(n->name(), 0), /*must_be_true=*/false); + Predicate* input_data_pred; + TF_RETURN_IF_ERROR(predicate_factory_.MakeSymbolPredicate( + n, /*output_idx=*/0, /*must_be_true=*/false, &input_data_pred)); + SetPredicate(n, {0, 1, Graph::kControlSlot}, input_data_pred, should_revisit); return Status::OK(); @@ -825,8 +975,10 @@ Status DeadnessAnalysisImpl::HandleMerge(Node* n, Predicate* start = predicate_factory_.MakeOrPredicate(non_recurrent_inputs); - Predicate* and_rec = - predicate_factory_.MakeAndRecurrencePredicate(start, step); + std::vector frame; + TF_RETURN_IF_ERROR(GetFullFrame(n, control_flow_info_, &frame)); + Predicate* and_rec = predicate_factory_.MakeAndRecurrencePredicate( + start, step, std::move(frame)); SetPredicate(n, {0, 1, Graph::kControlSlot}, and_rec, should_revisit); return Status::OK(); } @@ -841,8 +993,10 @@ Status DeadnessAnalysisImpl::HandleRecv(Node* n, // acquire a dead signal from a _Send. std::vector input_preds; TF_RETURN_IF_ERROR(GetInputPreds(n, EdgeKind::kDataAndControl, &input_preds)); - input_preds.push_back(predicate_factory_.MakeSymbolPredicate( - TensorId(n->name(), 0), /*must_be_true=*/false)); + Predicate* signal_is_alive; + TF_RETURN_IF_ERROR(predicate_factory_.MakeSymbolPredicate( + n, /*output_idx=*/0, /*must_be_true=*/false, &signal_is_alive)); + input_preds.push_back(signal_is_alive); SetPredicate(n, {0, Graph::kControlSlot}, predicate_factory_.MakeAndPredicate(input_preds), should_revisit); @@ -892,6 +1046,24 @@ Status DeadnessAnalysisImpl::Populate() { Status DeadnessAnalysisImpl::PopulateWithReversePostOrder( absl::Span rpo) { + std::vector unreachable_nodes; + // Compute the loop structure of the graph. + TF_RETURN_IF_ERROR( + BuildControlFlowInfo(&graph_, &control_flow_info_, &unreachable_nodes)); + + // Do some opportunistic error checking: + if (!unreachable_nodes.empty()) { + if (unreachable_nodes.size() > 5) { + unreachable_nodes.erase(unreachable_nodes.begin() + 5, + unreachable_nodes.end()); + } + + return errors::InvalidArgument( + "Found unreachable nodes, most likely source and sink nodes not " + "connected: ", + absl::StrJoin(unreachable_nodes, ", ")); + } + // This an abstract interpretation over the deadness propagation semantics of // the graph executor. // diff --git a/tensorflow/compiler/jit/deadness_analysis_test.cc b/tensorflow/compiler/jit/deadness_analysis_test.cc index 8a73101c184e6190921fd7729742922bd96f4bcf..38a5118d9a721b814e1b52ce4202d4fb783e3ac3 100644 --- a/tensorflow/compiler/jit/deadness_analysis_test.cc +++ b/tensorflow/compiler/jit/deadness_analysis_test.cc @@ -123,10 +123,9 @@ InductionVarInfo CreateInductionVariable(const Scope& root, Output increment_by = ops::Const(root.WithOpName(prefix + "/incr"), 1); Output final_value = ops::Const(root.WithOpName(prefix + "/final"), 10); Output loop_cond_expr = - ops::Less(root.WithOpName(prefix + "/less"), iv.output, final_value); - Output loop_cond = - ops::LoopCond(root.WithOpName(prefix + "/cond"), loop_cond_expr); - ops::Switch latch(root.WithOpName(prefix + "/latch"), iv.output, loop_cond); + ops::Less(root.WithOpName(prefix + "/cond"), iv.output, final_value); + ops::Switch latch(root.WithOpName(prefix + "/latch"), iv.output, + loop_cond_expr); ops::internal::Exit exit(root.WithOpName(prefix + "/exit"), latch.output_false); Output iv_next = ops::Add(root.WithOpName(prefix + "/ivnext"), @@ -140,7 +139,7 @@ InductionVarInfo CreateInductionVariable(const Scope& root, root.graph()->AddControlEdge(iv.output.node(), increment_by.node()); root.graph()->AddControlEdge(iv.output.node(), final_value.node()); - return {iv.output, loop_cond}; + return {iv.output, loop_cond_expr}; } InductionVarInfo CreateInductionVariable(const Scope& root, @@ -515,24 +514,27 @@ TEST(DeadnessAnalysisTest, Loop) { // In theory we should be able to tell that iv0/cond:0 and iv1/cond:0 // produce the same deadness. But we're not that smart today. - EXPECT_EQ(predicate_map[ControlOutputFor(iv0)], "{#true,&,*iv0/cond:0}"); - EXPECT_EQ(predicate_map[ControlOutputFor(iv1)], "{#true,&,*iv1/cond:0}"); - EXPECT_EQ(predicate_map[ControlOutputFor(iv2)], "{#true,&,*iv2/cond:0}"); + EXPECT_EQ(predicate_map[ControlOutputFor(iv0)], + "{#true,&,*iv0/cond:0}"); + EXPECT_EQ(predicate_map[ControlOutputFor(iv1)], + "{#true,&,*iv1/cond:0}"); + EXPECT_EQ(predicate_map[ControlOutputFor(iv2)], + "{#true,&,*iv2/cond:0}"); EXPECT_EQ(predicate_map[ControlOutputFor(add0)], - "({#true,&,*iv1/cond:0} & {#true,&,*iv0/cond:0})"); + "({#true,&,*iv0/cond:0} & {#true,&,*iv1/cond:0})"); EXPECT_EQ(predicate_map[ControlOutputFor(add1)], - "({#true,&,*iv1/cond:0} & {#true,&,*iv2/cond:0})"); + "({#true,&,*iv1/cond:0} & {#true,&,*iv2/cond:0})"); } } TEST(DeadnessAnalysisTest, ControlEquivalentLoopBodies) { Scope root = Scope::NewRootScope().ExitOnError(); - InductionVarInfo iv = CreateInductionVariable(root, "iv0", "frame", 0); + InductionVarInfo iv = CreateInductionVariable(root, "iv0", "loop", 0); Output dependent_iv0 = - CreateDependentLoopInvariantValue(root, "div0", "frame", iv.loop_cond, 0) + CreateDependentLoopInvariantValue(root, "div0", "loop", iv.loop_cond, 0) .induction_var; Output dependent_iv1 = - CreateDependentLoopInvariantValue(root, "div1", "frame", iv.loop_cond, 0) + CreateDependentLoopInvariantValue(root, "div1", "loop", iv.loop_cond, 0) .induction_var; Output add0 = ops::Add(root.WithOpName("add0"), dependent_iv0, dependent_iv1); @@ -549,13 +551,13 @@ TEST(DeadnessAnalysisTest, ControlEquivalentLoopBodies) { TF_ASSERT_OK(ComputePredicates(*root.graph(), &predicate_map)); EXPECT_EQ(predicate_map[ControlOutputFor(iv.induction_var)], - "{#true,&,*iv0/cond:0}"); + "{#true,&,*iv0/cond:0}"); EXPECT_EQ(predicate_map[ControlOutputFor(dependent_iv0)], - "{#true,&,(*iv0/cond:0 & iv0/iv:0)}"); + "{#true,&,(iv0/iv:0 & *iv0/cond:0)}"); EXPECT_EQ(predicate_map[ControlOutputFor(dependent_iv1)], - "{#true,&,(*iv0/cond:0 & iv0/iv:0)}"); + "{#true,&,(iv0/iv:0 & *iv0/cond:0)}"); EXPECT_EQ(predicate_map[ControlOutputFor(add0)], - "{#true,&,(*iv0/cond:0 & iv0/iv:0)}"); + "{#true,&,(iv0/iv:0 & *iv0/cond:0)}"); } } @@ -595,32 +597,33 @@ TEST(DeadnessAnalysisTest, LoopInvariantPredicateOnBackedge) { TEST(DeadnessAnalysisTest, ControlEquivalentNestedLoopBodies) { Scope root = Scope::NewRootScope().ExitOnError(); InductionVarInfo iv_outer = - CreateInductionVariable(root, "iv_outer", "frame", 0); + CreateInductionVariable(root, "iv_outer", "outer_loop", 0); + Output enter_constant_outer_loop = ops::internal::Enter( + root.WithOpName("constant_enter_outer_loop"), + ops::Const(root.WithOpName("constant"), 5), "outer_loop", + ops::internal::Enter::Attrs().IsConstant(true)); ops::Switch inner_value(root.WithOpName("outer_is_live"), - ops::Const(root.WithOpName("constant"), 5), - iv_outer.loop_cond); + enter_constant_outer_loop, iv_outer.loop_cond); InductionVarInfo iv_inner = CreateInductionVariable( - root, "iv_inner", "frame", - ops::internal::Enter(root.WithOpName("iv_inner/enter"), - inner_value.output_true, "frame_inner")); + root, "iv_inner", "inner_loop", inner_value.output_true); Output dependent_outer_iv0 = - CreateDependentLoopInvariantValue(root, "dependent_outer_iv0", "frame", - iv_outer.loop_cond, 0) + CreateDependentLoopInvariantValue(root, "dependent_outer_iv0", + "outer_loop", iv_outer.loop_cond, 0) .induction_var; Output dependent_outer_iv1 = - CreateDependentLoopInvariantValue(root, "dependent_outer_iv1", "frame", - iv_outer.loop_cond, 0) + CreateDependentLoopInvariantValue(root, "dependent_outer_iv1", + "outer_loop", iv_outer.loop_cond, 0) .induction_var; - Output dependent_inner_iv0 = - CreateDependentLoopInvariantValue(root, "dependent_inner_iv0", "frame", - iv_inner.loop_cond, dependent_outer_iv0) - .induction_var; - Output dependent_inner_iv1 = - CreateDependentLoopInvariantValue(root, "dependent_inner_iv1", "frame", - iv_inner.loop_cond, dependent_outer_iv1) - .induction_var; + Output dependent_inner_iv0 = CreateDependentLoopInvariantValue( + root, "dependent_inner_iv0", "inner_loop", + iv_inner.loop_cond, dependent_outer_iv0) + .induction_var; + Output dependent_inner_iv1 = CreateDependentLoopInvariantValue( + root, "dependent_inner_iv1", "inner_loop", + iv_inner.loop_cond, dependent_outer_iv1) + .induction_var; Output add0 = ops::Add(root.WithOpName("add0"), dependent_inner_iv0, dependent_inner_iv1); @@ -638,46 +641,51 @@ TEST(DeadnessAnalysisTest, ControlEquivalentNestedLoopBodies) { TF_ASSERT_OK(ComputePredicates(*root.graph(), &predicate_map)); EXPECT_EQ(predicate_map[ControlOutputFor(iv_outer.induction_var)], - "{#true,&,*iv_outer/cond:0}"); + "{#true,&,*iv_outer/cond:0}"); EXPECT_EQ(predicate_map[ControlOutputFor(iv_inner.induction_var)], - "{(*iv_outer/cond:0 & {#true,&,*iv_outer/cond:0}),&," - "*iv_inner/cond:0}"); + "{(*iv_outer/cond:0 & " + "{#true,&,*iv_outer/cond:0}),&,*iv_inner/" + "cond:0}"); EXPECT_EQ(predicate_map[ControlOutputFor(dependent_inner_iv0)], - "{{#true,&,(iv_outer/iv:0 & *iv_outer/cond:0)},&," - "(*iv_inner/cond:0 & iv_inner/iv:0)}"); + "{{#true,&,(iv_outer/iv:0 & " + "*iv_outer/cond:0)},&,(iv_inner/iv:0 & " + "*iv_inner/cond:0)}"); + EXPECT_EQ(predicate_map[ControlOutputFor(dependent_inner_iv1)], - "{{#true,&,(iv_outer/iv:0 & *iv_outer/cond:0)},&," - "(*iv_inner/cond:0 & iv_inner/iv:0)}"); + "{{#true,&,(iv_outer/iv:0 & " + "*iv_outer/cond:0)},&,(iv_inner/iv:0 & " + "*iv_inner/cond:0)}"); EXPECT_EQ(predicate_map[ControlOutputFor(add0)], - "{{#true,&,(iv_outer/iv:0 & *iv_outer/cond:0)},&," - "(*iv_inner/cond:0 & iv_inner/iv:0)}"); + "{{#true,&,(iv_outer/iv:0 & " + "*iv_outer/cond:0)},&,(iv_inner/iv:0 & " + "*iv_inner/cond:0)}"); } } TEST(DeadnessAnalysisTest, ControlNonEquivalentNestedLoopBodies) { Scope root = Scope::NewRootScope().ExitOnError(); - InductionVarInfo iv_outer_0 = - CreateInductionVariable(root, "iv_outer_0", "frame", 0); - ops::Switch inner_value_0(root.WithOpName("outer_0_is_live"), - ops::Const(root.WithOpName("constant"), 5), - iv_outer_0.loop_cond); - InductionVarInfo iv_inner_0 = CreateInductionVariable( - root, "iv_inner_0", "frame", - ops::internal::Enter(root.WithOpName("iv_inner_0/enter"), - inner_value_0.output_true, "frame_inner")); - - InductionVarInfo iv_outer_1 = - CreateInductionVariable(root, "iv_outer_1", "frame", 1); - ops::Switch inner_init_value_1(root.WithOpName("outer_1_is_live"), - ops::Const(root.WithOpName("constant"), 5), - iv_outer_1.loop_cond); - InductionVarInfo iv_inner_1 = CreateInductionVariable( - root, "iv_inner_1", "frame", - ops::internal::Enter(root.WithOpName("iv_inner_1/enter"), - inner_init_value_1.output_true, "frame_inner")); - Output add0 = ops::Add(root.WithOpName("add0"), iv_inner_0.induction_var, - iv_inner_1.induction_var); + + std::array outer_iv; + std::array inner_iv; + + for (int i : {0, 1}) { + InductionVarInfo iv_outer = + CreateInductionVariable(root, "iv_outer", "outer_loop", 0); + Output enter_constant_outer_loop = ops::internal::Enter( + root.WithOpName("constant_enter_outer_loop"), + ops::Const(root.WithOpName("constant"), 5), "outer_loop", + ops::internal::Enter::Attrs().IsConstant(true)); + ops::Switch inner_value(root.WithOpName("outer_is_live"), + enter_constant_outer_loop, iv_outer.loop_cond); + InductionVarInfo iv_inner = CreateInductionVariable( + root, "iv_inner", "inner_loop", inner_value.output_true); + + outer_iv[i] = iv_outer.induction_var; + inner_iv[i] = iv_inner.induction_var; + } + + Output add0 = ops::Add(root.WithOpName("add0"), inner_iv[0], inner_iv[1]); VLogGraphIfAsked(*root.graph()); @@ -692,21 +700,77 @@ TEST(DeadnessAnalysisTest, ControlNonEquivalentNestedLoopBodies) { PredicateMapTy predicate_map; TF_ASSERT_OK(ComputePredicates(*root.graph(), &predicate_map)); - EXPECT_EQ(predicate_map[ControlOutputFor(iv_outer_0.induction_var)], - "{#true,&,*iv_outer_0/cond:0}"); - EXPECT_EQ(predicate_map[ControlOutputFor(iv_inner_0.induction_var)], - "{(*iv_outer_0/cond:0 & {#true,&,*iv_outer_0/cond:0}),&," - "*iv_inner_0/cond:0}"); - EXPECT_EQ(predicate_map[ControlOutputFor(iv_outer_1.induction_var)], - "{#true,&,*iv_outer_1/cond:0}"); - EXPECT_EQ(predicate_map[ControlOutputFor(iv_inner_1.induction_var)], - "{(*iv_outer_1/cond:0 & {#true,&,*iv_outer_1/cond:0}),&," - "*iv_inner_1/cond:0}"); + EXPECT_EQ(predicate_map[ControlOutputFor(outer_iv[0])], + "{#true,&,*iv_outer/cond:0}"); + EXPECT_EQ(predicate_map[ControlOutputFor(inner_iv[0])], + "{(*iv_outer/cond:0 & " + "{#true,&,*iv_outer/cond:0}),&,*iv_inner/" + "cond:0}"); + EXPECT_EQ(predicate_map[ControlOutputFor(outer_iv[1])], + "{#true,&,*iv_outer/cond_1:0}"); + EXPECT_EQ(predicate_map[ControlOutputFor(inner_iv[1])], + "{(*iv_outer/cond_1:0 & " + "{#true,&,*iv_outer/cond_1:0}),&,*iv_inner/" + "cond_1:0}"); EXPECT_EQ(predicate_map[ControlOutputFor(add0)], - "({(*iv_outer_1/cond:0 & {#true,&,*iv_outer_1/cond:0}),&," - "*iv_inner_1/cond:0} & " - "{(*iv_outer_0/cond:0 & {#true,&,*iv_outer_0/cond:0}),&," - "*iv_inner_0/cond:0})"); + "({(*iv_outer/cond:0 & " + "{#true,&,*iv_outer/cond:0}),&,*iv_inner/" + "cond:0} & {(*iv_outer/cond_1:0 & " + "{#true,&,*iv_outer/cond_1:0}),&,*iv_inner/" + "cond_1:0})"); + } +} + +TEST(DeadnessAnalysisTest, AndRecurrenceNeedsFrameName) { + Scope root = Scope::NewRootScope().ExitOnError(); + InductionVarInfo iv_0 = CreateInductionVariable(root, "iv_0", "frame_0", 10); + InductionVarInfo iv_1 = CreateInductionVariable(root, "iv_1", "frame_1", 9); + + Output init = CreateSwitch(root, "init").output_true; + Output step = CreateSwitch(root, "step").output_true; + + std::array exits; + std::array next_iterations; + + for (int i : {0, 1}) { + Output init_enter = ops::internal::Enter( + root.WithOpName(absl::StrCat("init_enter_frame_", i)), init, + absl::StrCat("frame_", i), + ops::internal::Enter::Attrs().IsConstant(true)); + Output step_enter = ops::internal::Enter( + root.WithOpName(absl::StrCat("step_enter_frame_", i)), step, + absl::StrCat("frame_", i), + ops::internal::Enter::Attrs().IsConstant(true)); + + ops::Merge iv(root.WithOpName(absl::StrCat("expr_", i)), + {init_enter, init_enter}); + Output add = ops::Add(root.WithOpName(absl::StrCat("add_", i)), iv.output, + step_enter); + next_iterations[i] = ops::NextIteration( + root.WithOpName(absl::StrCat("expr_", i, "_next_iteration")), add); + EXPECT_TRUE( + root.graph() + ->UpdateEdge(next_iterations[i].node(), 0, iv.output.node(), 1) + .ok()); + exits[i] = ops::internal::Exit(root.WithOpName(absl::StrCat("exit_", i)), + iv.output); + } + + FixupSourceAndSinkEdges(root.graph()); + + { + PredicateMapTy predicate_map; + TF_ASSERT_OK(ComputePredicates(*root.graph(), &predicate_map)); + + EXPECT_NE(predicate_map[ControlOutputFor(exits[0])], + predicate_map[ControlOutputFor(exits[1])]); + EXPECT_NE(predicate_map[ControlOutputFor(exits[0])], ""); + EXPECT_NE(predicate_map[ControlOutputFor(exits[1])], ""); + + EXPECT_NE(predicate_map[ControlOutputFor(next_iterations[0])], + predicate_map[ControlOutputFor(next_iterations[1])]); + EXPECT_NE(predicate_map[ControlOutputFor(next_iterations[0])], ""); + EXPECT_NE(predicate_map[ControlOutputFor(next_iterations[1])], ""); } } @@ -818,5 +882,82 @@ TEST(DeadnessAnalysisTest, RecvVsSwitchText) { EXPECT_EQ(predicate_map[logical_and_output_0], "(recv:0 & *recv:0)"); } +TEST(DeadnessAnalysisTest, DeMorgan) { + Scope root = Scope::NewRootScope().ExitOnError(); + + Output cond_0 = ops::Placeholder(root.WithOpName("cond_0"), DT_BOOL); + Output cond_1 = ops::Placeholder(root.WithOpName("cond_1"), DT_BOOL); + Output value = ops::Placeholder(root.WithOpName("value"), DT_FLOAT); + + ops::Switch sw_0(root.WithOpName("switch_0"), value, cond_0); + ops::Switch sw_1(root.WithOpName("switch_1"), value, cond_1); + + Output and_0_1 = + ops::Add(root.WithOpName("and_0_1"), sw_0.output_true, sw_1.output_true); + + Output or_not0_not1 = ops::Merge(root.WithOpName("or_not0_not1"), + {sw_0.output_false, sw_1.output_false}) + .output; + + // Predicate(should_always_be_dead) = + // (A & B) & (~A | ~B) = (A & B) & ~(A & B) = False + Output should_always_be_dead = + ops::Add(root.WithOpName("should_always_be_dead"), and_0_1, or_not0_not1); + + // Predicate(should_always_be_dead) = + // (A & B) | (~A | ~B) = (A & B) | ~(A & B) = True + Output should_always_be_alive = + ops::Merge(root.WithOpName("should_always_be_alive"), + {and_0_1, or_not0_not1}) + .output; + + std::unique_ptr result; + TF_ASSERT_OK(AnalyzeDeadness(root.graph(), &result)); + + PredicateMapTy predicate_map; + TF_ASSERT_OK(ComputePredicates(*root.graph(), &predicate_map)); + + EXPECT_EQ(predicate_map[ControlOutputFor(should_always_be_dead)], "#false"); + EXPECT_EQ(predicate_map[ControlOutputFor(should_always_be_alive)], "#true"); +} + +TEST(DeadnessAnalysisTest, ConstantTrueSwitchCondition) { + Scope root = Scope::NewRootScope().ExitOnError(); + + Output constant_true = ops::Const(root.WithOpName("const_true"), true); + Output value = ops::Placeholder(root.WithOpName("value"), DT_FLOAT); + ops::Switch sw(root.WithOpName("switch"), value, constant_true); + + Output id_false = ops::Identity(root.WithOpName("id_false"), sw.output_false); + Output id_true = ops::Identity(root.WithOpName("id_true"), sw.output_true); + + FixupSourceAndSinkEdges(root.graph()); + + PredicateMapTy predicate_map; + TF_ASSERT_OK(ComputePredicates(*root.graph(), &predicate_map)); + + EXPECT_EQ(predicate_map[ControlOutputFor(id_false)], "#false"); + EXPECT_EQ(predicate_map[ControlOutputFor(id_true)], "#true"); +} + +TEST(DeadnessAnalysisTest, ConstantFalseSwitchCondition) { + Scope root = Scope::NewRootScope().ExitOnError(); + + Output constant_false = ops::Const(root.WithOpName("const_false"), false); + Output value = ops::Placeholder(root.WithOpName("value"), DT_FLOAT); + ops::Switch sw(root.WithOpName("switch"), value, constant_false); + + Output id_false = ops::Identity(root.WithOpName("id_false"), sw.output_false); + Output id_true = ops::Identity(root.WithOpName("id_true"), sw.output_true); + + FixupSourceAndSinkEdges(root.graph()); + + PredicateMapTy predicate_map; + TF_ASSERT_OK(ComputePredicates(*root.graph(), &predicate_map)); + + EXPECT_EQ(predicate_map[ControlOutputFor(id_false)], "#true"); + EXPECT_EQ(predicate_map[ControlOutputFor(id_true)], "#false"); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc index 03aba97bbe81a11f6366d118ee5bc573d0c6b31b..d0d7a3f3785469acd79a83b6897668f94fc6ea2e 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc @@ -1008,13 +1008,15 @@ Status Encapsulator::Subgraph::AddHostComputes( // subgraph. for (const auto& src_node : oc_subgraph.control_inputs) { Node* src_image = node_images.at(src_node); - graph_->AddControlEdge(src_image, host_compute); + graph_->AddControlEdge(src_image, host_compute, + /* allow_duplicates= */ true); } // Connect the _HostCompute node to its ancestor host compute nodes. for (const auto& ancestor_name : host_compute_ancestors) { Node* ancestor = host_compute_node[ancestor_name]; - graph_->AddControlEdge(ancestor, host_compute); + graph_->AddControlEdge(ancestor, host_compute, + /* allow_duplicates= */ true); } // Connect the consumers in the subgraph to the _HostCompute node. @@ -1031,7 +1033,8 @@ Status Encapsulator::Subgraph::AddHostComputes( // node. for (const auto& dst_node : oc_subgraph.control_outputs) { Node* dst_image = node_images.at(dst_node); - graph_->AddControlEdge(host_compute, dst_image); + graph_->AddControlEdge(host_compute, dst_image, + /* allow_duplicates= */ true); } } } @@ -1059,7 +1062,8 @@ Status Encapsulator::Subgraph::MakeSequencingNode(const string& subgraph_name, void Encapsulator::Subgraph::ConnectSequencerToCallNode(Graph* graph_out) { if (sequencer_ != nullptr) { VLOG(2) << "ConnectSequencerToCallNode"; - graph_out->AddControlEdge(sequencer_, call_node_); + graph_out->AddControlEdge(sequencer_, call_node_, + /* allow_duplicates= */ true); } } @@ -1279,7 +1283,8 @@ Status Encapsulator::Subgraph::AddRecvAtHostNode( // completes. This has no effect on execution order but prevents the // RecvAtHost being pruned. TF_RETURN_IF_ERROR(MakeSequencingNode(subgraph_name, graph_out)); - graph_out->AddControlEdge(oc_subgraph->recv_at_host, sequencer_); + graph_out->AddControlEdge(oc_subgraph->recv_at_host, sequencer_, + true /* skip duplicates check */); return Status::OK(); } @@ -1336,7 +1341,8 @@ Status Encapsulator::Subgraph::AddSendFromHostNode( // subgraph completes. This has no effect on execution order but prevents the // RecvAtHost being pruned. TF_RETURN_IF_ERROR(MakeSequencingNode(subgraph_name, graph_out)); - graph_out->AddControlEdge(oc_subgraph->send_from_host, sequencer_); + graph_out->AddControlEdge(oc_subgraph->send_from_host, sequencer_, + /* allow_duplicates= */ true); return Status::OK(); } @@ -1446,7 +1452,8 @@ Status Encapsulator::CopySubgraphEdges( src_func_id == dst_func_id) { Graph* g = subgraphs_[src_func_id].GetGraph(); if (edge->IsControlEdge()) { - g->AddControlEdge(src_image, dst_image); + g->AddControlEdge(src_image, dst_image, + /* allow_duplicates= */ true); } else { g->AddEdge(src_image, edge->src_output(), dst_image, edge->dst_input()); } @@ -1732,7 +1739,8 @@ Status Encapsulator::CopyEdgeToOutputGraph( if (edges_added ->emplace(OutputTensor(src_image, -1), InputTensor(dst_image, -1)) .second) { - graph_out->AddControlEdge(src_image, dst_image); + graph_out->AddControlEdge(src_image, dst_image, + /* allow_duplicates= */ true); } return Status::OK(); @@ -1761,7 +1769,8 @@ Status Encapsulator::AddCallNodeDependencies(Graph* graph_out) { const string& subgraph = ancestors.first; for (const string& ancestor : ancestors.second) { graph_out->AddControlEdge(subgraphs_[ancestor].GetCallNode(), - subgraphs_[subgraph].GetCallNode()); + subgraphs_[subgraph].GetCallNode(), + /* allow_duplicates= */ true); } } return Status::OK(); @@ -2129,7 +2138,8 @@ Status CheckClusterDependencyForCycles( const string& ancestor, const string& successor, const std::unordered_map>& ancestors, const std::unordered_map& node_ancestors_map, - GraphCycles* cycle_detector, std::map* cycle_detector_map) { + GraphCycles* cycle_detector, + std::unordered_map* cycle_detector_map) { if (cycle_detector_map->find(ancestor) == cycle_detector_map->end()) { (*cycle_detector_map)[ancestor] = cycle_detector->NewNode(); } @@ -2173,7 +2183,7 @@ Status Encapsulator::FindClusterDependencies() { // We check that clusters are acyclic using this cycle detector. GraphCycles cycle_detector; // Map from cluster name to cycle detector node id. - std::map cycle_detector_map; + std::unordered_map cycle_detector_map; // Process the nodes in topologically-sorted order. std::vector nodes; GetReversePostOrder(*graph_in_, &nodes); @@ -2535,7 +2545,33 @@ Status EncapsulateSubgraphsPass::Run( std::vector* input_permutation, std::vector* output_permutation, NodeDef* node) { // Optimize the subgraph. - OptimizeGraph(flr, subgraph); + // Do not constant fold nodes that output DT_VARIANT type tensors. + // XLA does not support Const nodes of Variant type since it needs + // to know the original ops to be able to compile them to the relevant + // XLA form. + // TODO(srbs): This filter is a little conservative. E.g. a subgraph of + // the form: + // Const + // | + // EmptyTensorList -> TensorListPushBack -> TensorListPopBack -> Op + // | + // (Discard popped list) + // + // Would have been reduced to "Const -> Op" without this filter. + // However since we are only allowed to specify the filter at the "Node" + // level there is no good way to allow the above behavior. So we + // disallow any sort of constant folding on Variant nodes for now. + auto cf_consider_fn = [](const Node* n) { + for (const auto& output_arg : n->op_def().output_arg()) { + if (output_arg.type() == DT_VARIANT) { + return false; + } + } + return true; + }; + GraphOptimizer::Options graph_optimizer_options; + graph_optimizer_options.cf_consider_fn = cf_consider_fn; + OptimizeGraph(flr, subgraph, graph_optimizer_options); const int num_args = input_permutation->size(); std::vector const_args(num_args); diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc index 8617beec004d0fe912155f054442c5b6249bb6b5..1f8ec09e19c01d0a8b2a3761135ed53dfb2ad3b0 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc @@ -25,6 +25,8 @@ limitations under the License. #include "tensorflow/compiler/jit/encapsulate_util.h" #include "tensorflow/compiler/jit/extract_outside_compilation_pass.h" #include "tensorflow/compiler/tf2xla/side_effect_util.h" +#include "tensorflow/core/common_runtime/device_factory.h" +#include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/framework/function_testlib.h" #include "tensorflow/core/framework/graph_to_functiondef.h" #include "tensorflow/core/graph/graph_constructor.h" @@ -32,6 +34,8 @@ limitations under the License. #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/test.h" +#include "tensorflow/core/public/session_options.h" +#include "tensorflow/core/public/version.h" #include "tensorflow/core/util/equal_graph_def.h" namespace tensorflow { @@ -513,6 +517,18 @@ Status Encapsulate(GraphDef* graphdef, FunctionDefLibrary* library, s = PerformStaticShapeInferenceBeforeEncapsulation(graph.get()); if (!s.ok()) return s; + // Create FunctionLibraryRuntime. + SessionOptions session_options; + std::vector> devices; + TF_CHECK_OK(DeviceFactory::AddDevices( + session_options, "/job:localhost/replica:0/task:0", &devices)); + OptimizerOptions opts; + auto device_mgr = absl::make_unique(std::move(devices)); + auto pflr = absl::make_unique( + device_mgr.get(), Env::Default(), TF_GRAPH_DEF_VERSION, lib_def.get(), + opts, /*default_thread_pool=*/nullptr, /*cluster_flr=*/nullptr); + auto flr = pflr->GetFLR("/job:localhost/replica:0/task:0/cpu:0"); + std::unique_ptr graph_out; s = EncapsulateSubgraphsInFunctions( "_encapsulate", /*outside_compilation_attribute=*/"", *graph, @@ -538,7 +554,7 @@ Status Encapsulate(GraphDef* graphdef, FunctionDefLibrary* library, std::map{}}); } s = ExtractOutsideCompilation("_encapsulate", "_outside", clusters, - graph_out.get(), lib_def.get()); + graph_out.get(), flr, lib_def.get()); if (!s.ok()) return s; GraphDef graphdef_out; @@ -941,7 +957,9 @@ TEST(EncapsulateSubgraphsTest, OneFunctionOneOutside) { {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", shape_inference_graph}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O1"}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}, {"c"}}, }, {{"f_0_retval_retval", "F:o:0"}}); @@ -1101,7 +1119,9 @@ TEST(EncapsulateSubgraphsTest, OneFunctionTwoOutside) { {"key", "host_compute_channel_F1_O2"}, {"shape_inference_graph", shape_inference_graph2}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O2"}}, + {"_outside_compilation_subgraph", "O2"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}, {"F"}}, {{"outside_compilation_O1_host_compute"}, "XlaHostCompute", @@ -1112,7 +1132,9 @@ TEST(EncapsulateSubgraphsTest, OneFunctionTwoOutside) { {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", shape_inference_graph1}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O1"}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}, {"D"}}, }, {{"g_0_retval_retval", "outside_compilation_O2_host_compute:outputs:0"}, @@ -1244,7 +1266,9 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { {"shape_inference_graph", NameAttrList()}, {"shapes", absl::Span({shape_proto_expected})}, - {"_outside_compilation_subgraph", "O1"}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}, {"D"}}, }, {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"}, @@ -1269,7 +1293,9 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { {"shape_inference_graph", NameAttrList()}, {"shapes", absl::Span({shape_proto_expected})}, - {"_outside_compilation_subgraph", "O1"}}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, }, {{"g_0_retval_retval", "G:o:0"}, {"i_0_retval_retval", "I:o:0"}}); @@ -1397,7 +1423,9 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutsideDependencyFromOutside) { {"shape_inference_graph", NameAttrList()}, {"shapes", absl::Span({shape_proto_expected})}, - {"_outside_compilation_subgraph", "O1"}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}, {"D"}}, }, {{"f_0_retval_retval", "F:o:0"}}); @@ -1419,7 +1447,9 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutsideDependencyFromOutside) { {"shape_inference_graph", NameAttrList()}, {"shapes", absl::Span({shape_proto_expected})}, - {"_outside_compilation_subgraph", "O1"}}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, }, {{"i_0_retval_retval", "I:o:0"}}); @@ -1527,7 +1557,9 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputs) { {"shape_inference_graph", NameAttrList()}, {"shapes", absl::Span({shape_proto_expected})}, - {"_outside_compilation_subgraph", "O1"}}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, }, {{"f_0_retval_retval", "F:o:0"}}); @@ -1615,7 +1647,9 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlInput) { {"shape_inference_graph", NameAttrList()}, {"shapes", absl::Span({shape_proto_expected})}, - {"_outside_compilation_subgraph", "O1"}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}, {"D"}}, }, {{"f_0_retval_retval", "F:o:0"}}); @@ -1716,7 +1750,9 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoOutputs) { {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", shape_inference_graph}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O1"}}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, }, {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"}, {"f_0_retval_retval", "F:o:0"}}); @@ -1821,7 +1857,9 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlOutput) { {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", shape_inference_graph}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O1"}}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, }, {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"}, {"f_0_retval_retval", "F:o:0"}}); @@ -1949,7 +1987,9 @@ TEST(EncapsulateSubgraphsTest, {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", shape_inference_graph1}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O1"}}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, {{"outside_compilation_O2_host_compute"}, "XlaHostCompute", {"F:o:0"}, @@ -1959,7 +1999,9 @@ TEST(EncapsulateSubgraphsTest, {"key", "host_compute_channel_F1_O2"}, {"shape_inference_graph", shape_inference_graph2}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O2"}}}, + {"_outside_compilation_subgraph", "O2"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, }, {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"}, {"h_0_retval_retval", "H:o:0"}}); @@ -2082,7 +2124,9 @@ TEST(EncapsulateSubgraphsTest, {"key", "host_compute_channel_F1_O2"}, {"shape_inference_graph", NameAttrList()}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O2"}}}, + {"_outside_compilation_subgraph", "O2"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, {{"outside_compilation_O1_host_compute"}, "XlaHostCompute", {"D:o:0"}, @@ -2092,7 +2136,9 @@ TEST(EncapsulateSubgraphsTest, {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", shape_inference_graph}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O1"}}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, }, {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"}, {"h_0_retval_retval", "H:o:0"}}); @@ -2214,7 +2260,9 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationClusterDependency) { {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", shape_inference_graph}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O1"}}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, {{"outside_compilation_O2_host_compute"}, "XlaHostCompute", {"D:o:0"}, @@ -2224,7 +2272,9 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationClusterDependency) { {"key", "host_compute_channel_F1_O2"}, {"shape_inference_graph", NameAttrList()}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O2"}}, + {"_outside_compilation_subgraph", "O2"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}, {}}, {{"outside_compilation_O3_host_compute"}, "XlaHostCompute", @@ -2235,7 +2285,9 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationClusterDependency) { {"key", "host_compute_channel_F1_O3"}, {"shape_inference_graph", NameAttrList()}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O3"}}, + {"_outside_compilation_subgraph", "O3"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}, {}}}, {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"}, {"h_0_retval_retval", "H:o:0"}}); @@ -2354,7 +2406,9 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputsOrOutputs) { {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", shape_inference_graph}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O1"}}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, }, {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"}, {"f_0_retval_retval", "F:o:0"}}); @@ -2465,7 +2519,9 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationShapeInference) { {"key", "host_compute_channel_F1_O1"}, {"shape_inference_graph", shape_inference_graph}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O1"}}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}, {"c"}}, }, {{"f_0_retval_retval", "F:o:0"}}); diff --git a/tensorflow/compiler/jit/encapsulate_util_test.cc b/tensorflow/compiler/jit/encapsulate_util_test.cc index 3bb979e0698d2d6be42ed5bae66c25267928192c..6d1661222e3eaf9df4f9f91f2b426c80b55245b2 100644 --- a/tensorflow/compiler/jit/encapsulate_util_test.cc +++ b/tensorflow/compiler/jit/encapsulate_util_test.cc @@ -21,7 +21,6 @@ limitations under the License. #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_shape.pb.h" -#include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/platform/test.h" namespace tensorflow { diff --git a/tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc b/tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc index ec745cdbb7e237f8b4935dd41e9791fc75f5355d..109684be72a2d67d04ac9efda0b17650f6905752 100644 --- a/tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc +++ b/tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc @@ -15,13 +15,17 @@ limitations under the License. #include "tensorflow/compiler/jit/encapsulate_xla_computations_pass.h" +#include "absl/algorithm/container.h" #include "absl/container/flat_hash_set.h" #include "absl/memory/memory.h" +#include "absl/strings/ascii.h" #include "absl/strings/str_cat.h" #include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h" #include "tensorflow/compiler/tf2xla/dump_graph.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/core/framework/node_def.pb.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/lib/core/stringpiece.h" #include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/lib/strings/proto_serialization.h" #include "tensorflow/core/lib/strings/str_util.h" @@ -36,6 +40,25 @@ namespace { const char* const kXlaClusterOutput = "XlaClusterOutput"; +bool IsCpuGpuCompile(const Graph* graph) { + for (Node* n : graph->nodes()) { + string name; + // Only consider nodes being compiled. + if (!GetNodeAttr(n->attrs(), + EncapsulateXlaComputationsPass::kXlaClusterAttr, &name) + .ok()) + continue; + // Early return for any node with a device that is not a CPU or GPU. + DeviceNameUtils::ParsedName parsed; + if (DeviceNameUtils::ParseFullName(n->requested_device(), &parsed)) { + if (parsed.type != DEVICE_CPU && parsed.type != DEVICE_GPU) { + return false; + } + } + } + return true; +} + // Checks if a graph node is marked to be a guaranteed constant. bool is_guaranteed_constant(const Node& n) { bool guaranteed_constant = false; @@ -173,9 +196,10 @@ Status RewriteSubgraph(const std::vector& arg_source_tensors, // Nondeterminism in serialization would not lead to incorrect results, but // may cause spurious cache misses. DeterministicSerialization is a // best-effort deterministic serialization. - string serialized; - TF_RET_CHECK(SerializeToStringDeterministic(gdef, &serialized)); - uint64 fingerprint = Fingerprint64(serialized); + const size_t size = gdef.ByteSizeLong(); + auto serialized = absl::make_unique(size); + TF_RET_CHECK(SerializeToBufferDeterministic(gdef, serialized.get(), size)); + uint64 fingerprint = Fingerprint64(absl::string_view(serialized.get(), size)); LOG(INFO) << "Subgraph fingerprint:" << fingerprint; call_def->set_op(absl::StrCat(call_def->op(), "_", fingerprint)); return Status::OK(); @@ -351,12 +375,19 @@ Status EncapsulateXlaComputationsPass::Run( << dump_graph::DumpGraphToFile("encapsulate_xla_computations_before", **options.graph, options.flib_def); - TF_RETURN_IF_ERROR(Encapsulate(options.graph, options.flib_def)); + const char* additional_help = + IsCpuGpuCompile(options.graph->get()) + ? xla::status_macros::kPossibleAutoJitAlternative + : ""; + + TF_RETURN_WITH_CONTEXT_IF_ERROR(Encapsulate(options.graph, options.flib_def), + additional_help); VLOG(1) << "EncapsulateXlaComputations() half-way: " << dump_graph::DumpGraphToFile("encapsulate_xla_computations_halfway", **options.graph, options.flib_def); - TF_RETURN_IF_ERROR(BuildXlaLaunchOps(options.graph->get())); + TF_RETURN_WITH_CONTEXT_IF_ERROR(BuildXlaLaunchOps(options.graph->get()), + additional_help); VLOG(1) << "EncapsulateXlaComputations() finished: " << dump_graph::DumpGraphToFile("encapsulate_xla_computations_after", **options.graph, options.flib_def); diff --git a/tensorflow/compiler/jit/extract_outside_compilation_pass.cc b/tensorflow/compiler/jit/extract_outside_compilation_pass.cc index 8b01768c49422b331b52a8ba31bade000c95722e..2a770c527b2fae91352fd17dacb13495a3a73f34 100644 --- a/tensorflow/compiler/jit/extract_outside_compilation_pass.cc +++ b/tensorflow/compiler/jit/extract_outside_compilation_pass.cc @@ -30,6 +30,7 @@ limitations under the License. #include "tensorflow/core/framework/tensor_shape.pb.h" #include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/gtl/cleanup.h" namespace tensorflow { @@ -308,6 +309,10 @@ xla::StatusOr BuildXlaHostComputeNodeDef( host_compute_builder.Attr("tpu_core", core); } + // Set input tokens. + host_compute_builder.Attr(kXlaTokenInputNodesAttrName, + std::vector{kXlaTokenArgNodeName}); + // Populate inputs. std::vector input_dtypes; TF_RETURN_IF_ERROR(GetNodeAttr(call_node->attrs(), "Tinputs", &input_dtypes)); @@ -398,8 +403,8 @@ Status ReplaceOrRemoveOutsideCompilationCallNode( } // Resets "device_ordinal" attr to placeholder value for related nodes -// (XlaRecvAtHost nodes; XlaSendFromHost nodes; If nodes containing -// XlaRecvAtHost/XlaSendFromHost). +// (XlaRecvAtHost nodes; XlaSendFromHost nodes; If/While/FuncCall nodes +// containing XlaRecvAtHost/XlaSendFromHost). Status ResetDeviceOrdinalToPlaceholderValue(Graph* g) { AttrValue device_ordinal_value; device_ordinal_value.set_placeholder("device_ordinal"); @@ -429,6 +434,10 @@ Status ResetDeviceOrdinalToPlaceholderValue(Graph* g) { n->ClearAttr(attr_name); n->AddAttr(attr_name, branch_func); } + } else if (HasNodeAttr(n->def(), "device_ordinal")) { + // Function call node containing outside compilation. + n->ClearAttr("device_ordinal"); + n->AddAttr("device_ordinal", device_ordinal_value); } else { return errors::Internal("Unknown node marked with ", kXlaHasHostTransferAttrName, ": ", @@ -1217,20 +1226,129 @@ Status BuildHostGraphForWhileNode( return Status::OK(); } +// Builds host graph for func call nodes. +Status BuildHostGraphForFuncCallNode(const string& func_call_node_name, + const string& xla_cluster_name, + const string& func_call_host_func_name, + const string& host_graph_func_name, + FunctionLibraryDefinition* fld) { + Graph host_graph(fld); + AttrValue device_ordinal_value; + device_ordinal_value.set_placeholder("device_ordinal"); + + // Step 1: add key placeholder node. + TF_ASSIGN_OR_RETURN( + Node * key_placeholder, + AddHostComputeKeyPlaceholder(xla_cluster_name, &host_graph)); + + // Step 2: rewrite `host_func_name`, replace key placeholder with an _Arg + // node. + TF_RETURN_IF_ERROR(ReplaceKeyPlaceholderWithArgNode( + xla_cluster_name, func_call_host_func_name, fld)); + + // Step 3: build a function call node with `host_func_name`, with + // `key_placeholder` as input. + NodeDefBuilder call_builder(absl::StrCat("oc_call_", func_call_node_name), + func_call_host_func_name, fld); + call_builder.Input(key_placeholder->name(), 0, DT_STRING); + call_builder.Attr("device_ordinal", device_ordinal_value); + call_builder.Attr(kXlaHasHostTransferAttrName, true); + NodeDef call_def; + TF_RETURN_IF_ERROR(call_builder.Finalize(&call_def)); + Status s; + Node* call_node = host_graph.AddNode(call_def, &s); + TF_RETURN_IF_ERROR(s); + host_graph.AddEdge(key_placeholder, 0, call_node, 0); + + // Convert `host_graph` to function, and add a "device_ordinal" attr. + FunctionDef oc_host_graph_fdef; + TF_RETURN_IF_ERROR(GraphToFunctionDef(host_graph, host_graph_func_name, + &oc_host_graph_fdef)); + if (fld->Find(host_graph_func_name)) { + TF_RETURN_IF_ERROR( + fld->ReplaceFunction(host_graph_func_name, oc_host_graph_fdef)); + } else { + TF_RETURN_IF_ERROR(fld->AddFunctionDef(oc_host_graph_fdef)); + } + + return Status::OK(); +} + Status ExtractOutsideCompilationForNodesWithAssociatedFunctions( Graph* g, const string& xla_cluster_attr_name, const string& outside_compilation_attr_name, const string& xla_cluster_name, - const std::map& host_compute_core, + const std::map& host_compute_core, FunctionLibraryRuntime* flr, FunctionLibraryDefinition* fld, std::vector* host_graphs, std::vector* shape_inference_graphs, bool* has_outside_compilation) { - std::vector if_nodes, while_nodes; + std::vector if_nodes, while_nodes, func_call_nodes; for (Node* n : g->nodes()) { if (n->type_string() == "If") { if_nodes.push_back(n); } else if (n->type_string() == "While") { while_nodes.push_back(n); + } else if (fld->Contains(n->type_string())) { + func_call_nodes.push_back(n); + } else if (n->type_string() == FunctionLibraryDefinition::kGradientOp) { + // Only gradient for user-defined function should be considered as + // function call node. + NameAttrList original_func; + TF_RETURN_IF_ERROR(GetNodeAttr( + n->def(), FunctionLibraryDefinition::kFuncAttr, &original_func)); + if (fld->Contains(original_func.name())) { + func_call_nodes.push_back(n); + } + } + } + + for (Node* n : func_call_nodes) { + // Extract outside compilation for the function call. + bool func_has_outside_compilation = false; + NameAttrList func; + func.set_name(n->type_string()); + typedef protobuf::Map AttrMap; + *func.mutable_attr() = AttrMap(n->attrs().begin(), n->attrs().end()); + string new_func_name = absl::StrCat(n->name(), "_oc"); + string host_func_name = absl::StrCat("oc_func_call_host_", n->name()); + TF_RETURN_IF_ERROR(ExtractOutsideCompilationForFunction( + xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name, + func, new_func_name, host_func_name, host_compute_core, flr, fld, + shape_inference_graphs, &func_has_outside_compilation)); + + // If the function call does not have outside compilation, nothing to do. + if (!func_has_outside_compilation) { + continue; } + + *has_outside_compilation = true; + + // Change `n` to call the new function directly. + NodeDefBuilder replace_builder(n->name(), new_func_name, fld); + for (const Edge* e : n->in_edges()) { + if (e->IsControlEdge()) { + continue; + } + replace_builder.Input(e->src()->name(), e->src_output(), + e->src()->output_type(e->src_output())); + } + for (const auto& attr : n->attrs()) { + replace_builder.Attr(attr.first, attr.second); + } + NodeDef replace_def; + TF_RETURN_IF_ERROR(replace_builder.Finalize(&replace_def)); + TF_ASSIGN_OR_RETURN(Node * replace, ReplaceNode(g, n, replace_def)); + replace->AddAttr(kXlaTokenInputNodesAttrName, + std::vector{kXlaTokenArgNodeName}); + + // Build host side graph for the function call. + string oc_host_graph_name = + absl::StrCat("oc_func_host_graph_", replace->name()); + TF_RETURN_IF_ERROR( + BuildHostGraphForFuncCallNode(replace->name(), xla_cluster_name, + host_func_name, oc_host_graph_name, fld)); + + // Record the host graph. + host_graphs->push_back(oc_host_graph_name); } for (Node* n : if_nodes) { @@ -1251,12 +1369,12 @@ Status ExtractOutsideCompilationForNodesWithAssociatedFunctions( TF_RETURN_IF_ERROR(ExtractOutsideCompilationForFunction( xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name, then_branch, then_branch_xla_func_name, then_branch_host_func_name, - host_compute_core, fld, shape_inference_graphs, + host_compute_core, flr, fld, shape_inference_graphs, &then_branch_has_outside_compilation)); TF_RETURN_IF_ERROR(ExtractOutsideCompilationForFunction( xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name, else_branch, else_branch_xla_func_name, else_branch_host_func_name, - host_compute_core, fld, shape_inference_graphs, + host_compute_core, flr, fld, shape_inference_graphs, &else_branch_has_outside_compilation)); // If then/else branch do not have outside compilation, nothing to do. @@ -1316,12 +1434,12 @@ Status ExtractOutsideCompilationForNodesWithAssociatedFunctions( body_xla_func_name = absl::StrCat(body.name(), "_oc"); TF_RETURN_IF_ERROR(ExtractOutsideCompilationForFunction( xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name, - cond, cond_xla_func_name, cond_host_func_name, host_compute_core, fld, - shape_inference_graphs, &cond_has_outside_compilation)); + cond, cond_xla_func_name, cond_host_func_name, host_compute_core, flr, + fld, shape_inference_graphs, &cond_has_outside_compilation)); TF_RETURN_IF_ERROR(ExtractOutsideCompilationForFunction( xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name, - body, body_xla_func_name, body_host_func_name, host_compute_core, fld, - shape_inference_graphs, &body_has_outside_compilation)); + body, body_xla_func_name, body_host_func_name, host_compute_core, flr, + fld, shape_inference_graphs, &body_has_outside_compilation)); // If cond/body do not have outside compilation, nothing to do. if (!cond_has_outside_compilation && !body_has_outside_compilation) { @@ -1469,17 +1587,27 @@ Status ExtractOutsideCompilationForFunction( const string& outside_compilation_attr_name, const string& xla_cluster_name, const NameAttrList& func_name_attrs, const string& new_func_name, const string& host_graph_func_name, - const std::map& host_compute_core, + const std::map& host_compute_core, FunctionLibraryRuntime* flr, FunctionLibraryDefinition* fld, std::vector* shape_inference_graphs, bool* has_outside_compilation) { + // Convert the function to graph. const string& func_name = func_name_attrs.name(); - const FunctionDef* fdef = fld->Find(func_name); - if (!fdef) { - return errors::Internal("Cannot find function ", func_name); - } + FunctionLibraryRuntime::Handle handle; + TF_RETURN_IF_ERROR( + flr->Instantiate(func_name, AttrSlice(&func_name_attrs.attr()), &handle)); + Status ret_status = Status::OK(); + auto cleanup_handle = gtl::MakeCleanup([&]() { + auto s = flr->ReleaseHandle(handle); + if (!s.ok()) { + ret_status.Update(s); + } + }); + const FunctionBody* fbody = flr->GetFunctionBody(handle); + + // Check if we have outside compilation nodes. *has_outside_compilation = false; - for (auto& node_def : fdef->node_def()) { - if (HasNodeAttr(node_def, outside_compilation_attr_name)) { + for (Node* n : fbody->graph->nodes()) { + if (HasNodeAttr(n->def(), outside_compilation_attr_name)) { *has_outside_compilation = true; break; } @@ -1487,16 +1615,6 @@ Status ExtractOutsideCompilationForFunction( // We cannot early return here, because we might have outside compilation in // If/While function body. - // Convert the function to graph. - FunctionBody* fbody = nullptr; - TF_RETURN_IF_ERROR(FunctionDefToBodyHelper( - *fld->Find(func_name), AttrSlice(&func_name_attrs.attr()), fld, - [&](const string& op, const OpDef** sig) { - return fld->LookUpOpDef(op, sig); - }, - &fbody)); - std::unique_ptr fbody_deleter(fbody); - // Preprocess edges between different outside compilations. They will be // restored in `ConstructHostGraph()`. TF_RETURN_IF_ERROR(PreprocessEdgesBetweenOutsideCompilations( @@ -1553,16 +1671,11 @@ Status ExtractOutsideCompilationForFunction( TF_RETURN_IF_ERROR(ReplaceOrRemoveOutsideCompilationCallNode( graph_out.get(), n, host_compute_core)); } - if (VLOG_IS_ON(4)) { - dump_graph::DumpGraphToFile( - absl::StrCat("extract_outside_compilation_for_func_after_", func_name), - *graph_out, fld); - } // Handle nodes with associated functions. TF_RETURN_IF_ERROR(ExtractOutsideCompilationForNodesWithAssociatedFunctions( graph_out.get(), xla_cluster_attr_name, outside_compilation_attr_name, - xla_cluster_name, host_compute_core, fld, + xla_cluster_name, host_compute_core, flr, fld, &outside_compilation_host_graphs, shape_inference_graphs, has_outside_compilation)); @@ -1580,20 +1693,31 @@ Status ExtractOutsideCompilationForFunction( FunctionDef updated_fdef; TF_RETURN_IF_ERROR( GraphToFunctionDef(*graph_out, new_func_name, &updated_fdef)); + const FunctionDef* original_fdef = fld->Find(func_name); + if (original_fdef) { + for (const auto& attr : original_fdef->attr()) { + (*updated_fdef.mutable_attr())[attr.first] = attr.second; + } + } if (fld->Find(new_func_name)) { TF_RETURN_IF_ERROR(fld->ReplaceFunction(new_func_name, updated_fdef)); } else { TF_RETURN_IF_ERROR(fld->AddFunctionDef(updated_fdef)); } + if (VLOG_IS_ON(4)) { + dump_graph::DumpGraphToFile( + absl::StrCat("extract_outside_compilation_for_func_after_", func_name), + *graph_out, fld); + } - return Status::OK(); + return ret_status; } Status ExtractOutsideCompilation( const string& xla_cluster_attr_name, const string& outside_compilation_attr_name, const std::unordered_map& clusters, Graph* g, - FunctionLibraryDefinition* fld) { + FunctionLibraryRuntime* flr, FunctionLibraryDefinition* fld) { if (VLOG_IS_ON(4)) { dump_graph::DumpGraphToFile("extract_outside_compilation_before", *g, fld); } @@ -1610,7 +1734,7 @@ Status ExtractOutsideCompilation( TF_RETURN_IF_ERROR(ExtractOutsideCompilationForFunction( xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name, func_name_attrs, func_name_attrs.name(), host_graph_func_name, - host_compute_core, fld, &shape_inference_graphs, + host_compute_core, flr, fld, &shape_inference_graphs, &has_outside_compilation)); TF_RETURN_IF_ERROR( ExpandHostGraphIntoMainGraph(g, fld, host_graph_func_name, n)); diff --git a/tensorflow/compiler/jit/extract_outside_compilation_pass.h b/tensorflow/compiler/jit/extract_outside_compilation_pass.h index e07e7c5dd0cd42ddd4d643d8b36583c82056bbb2..d64cc2a103ed040cbf413ac736f97f84459e869b 100644 --- a/tensorflow/compiler/jit/extract_outside_compilation_pass.h +++ b/tensorflow/compiler/jit/extract_outside_compilation_pass.h @@ -89,7 +89,7 @@ Status ExtractOutsideCompilationForFunction( const string& outside_compilation_attr_name, const string& xla_cluster_name, const NameAttrList& func_name_attrs, const string& new_func_name, const string& host_graph_func_name, - const std::map& host_compute_core, + const std::map& host_compute_core, FunctionLibraryRuntime* flr, FunctionLibraryDefinition* fld, std::vector* shape_inference_graphs, bool* has_outside_compilation); @@ -101,7 +101,7 @@ Status ExtractOutsideCompilation( const string& xla_cluster_attr_name, const string& outside_compilation_attr_name, const std::unordered_map& clusters, Graph* g, - FunctionLibraryDefinition* fld); + FunctionLibraryRuntime* flr, FunctionLibraryDefinition* fld); } // namespace tensorflow diff --git a/tensorflow/compiler/jit/extract_outside_compilation_pass_test.cc b/tensorflow/compiler/jit/extract_outside_compilation_pass_test.cc index e9a89e34e0c7b04b4be34e367b2d0bf627c0061a..7c3a24feff81b21a5d2347d21fb80988bc3e6065 100644 --- a/tensorflow/compiler/jit/extract_outside_compilation_pass_test.cc +++ b/tensorflow/compiler/jit/extract_outside_compilation_pass_test.cc @@ -23,6 +23,7 @@ limitations under the License. #include "tensorflow/cc/ops/standard_ops.h" #include "tensorflow/compiler/jit/encapsulate_util.h" #include "tensorflow/compiler/xla/test.h" +#include "tensorflow/core/common_runtime/device_factory.h" #include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/framework/common_shape_fns.h" #include "tensorflow/core/framework/function.h" @@ -31,6 +32,8 @@ limitations under the License. #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_shape.pb.h" #include "tensorflow/core/platform/test.h" +#include "tensorflow/core/public/session_options.h" +#include "tensorflow/core/public/version.h" namespace tensorflow { @@ -222,7 +225,42 @@ TEST(RewriteOutsideCompilationSubgraphFnTest, ShapesInferred) { EXPECT_EQ(shapes[0].dim_size(), 1); } -TEST(ExtractOutsideCompilationForFunctionTest, Basic) { +class ExtractOutsideCompilationForFunctionTest : public ::testing::Test { + public: + void SetUp() override { + SessionOptions session_options; + std::vector> devices; + TF_CHECK_OK(DeviceFactory::AddDevices( + session_options, "/job:localhost/replica:0/task:0", &devices)); + device_mgr_ = absl::make_unique(std::move(devices)); + } + + Status ExtractOutsideCompilationTest( + const string &xla_cluster_attr_name, + const string &outside_compilation_attr_name, + const string &xla_cluster_name, const NameAttrList &func_name_attrs, + const string &new_func_name, const string &host_graph_func_name, + const std::map &host_compute_core, + FunctionLibraryDefinition *fld, + std::vector *shape_inference_graphs, + bool *has_outside_compilation) { + OptimizerOptions opts; + pflr_ = absl::make_unique( + device_mgr_.get(), Env::Default(), TF_GRAPH_DEF_VERSION, fld, opts, + /*default_thread_pool=*/nullptr, /*cluster_flr=*/nullptr); + auto flr = pflr_->GetFLR("/job:localhost/replica:0/task:0/cpu:0"); + return ExtractOutsideCompilationForFunction( + xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name, + func_name_attrs, new_func_name, host_graph_func_name, host_compute_core, + flr, fld, shape_inference_graphs, has_outside_compilation); + } + + private: + std::unique_ptr device_mgr_; + std::unique_ptr pflr_; +}; + +TEST_F(ExtractOutsideCompilationForFunctionTest, Basic) { // Build the XLA computation func. // "const0" // "identity0" = "const0" (outside compilation cluster "0") @@ -256,7 +294,7 @@ TEST(ExtractOutsideCompilationForFunctionTest, Basic) { NameAttrList name_attrs; name_attrs.set_name("cluster"); *name_attrs.mutable_attr() = attrs; - TF_CHECK_OK(ExtractOutsideCompilationForFunction( + TF_CHECK_OK(ExtractOutsideCompilationTest( "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten", "host_graph", host_compute_core, &fld, &shape_inference_graphs, &has_outside_compilation)); @@ -362,7 +400,7 @@ TEST(ExtractOutsideCompilationForFunctionTest, Basic) { } } -TEST(ExtractOutsideCompilationForFunctionTest, NoHostGraph) { +TEST_F(ExtractOutsideCompilationForFunctionTest, NoHostGraph) { // Build the XLA computation func. // "const0" FunctionDefLibrary fdl; @@ -384,7 +422,7 @@ TEST(ExtractOutsideCompilationForFunctionTest, NoHostGraph) { NameAttrList name_attrs; name_attrs.set_name("cluster"); *name_attrs.mutable_attr() = attrs; - TF_CHECK_OK(ExtractOutsideCompilationForFunction( + TF_CHECK_OK(ExtractOutsideCompilationTest( "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten", "host_graph", host_compute_core, &fld, &shape_inference_graphs, &has_outside_compilation)); @@ -406,7 +444,7 @@ TEST(ExtractOutsideCompilationForFunctionTest, NoHostGraph) { EXPECT_EQ(host_graph->num_nodes(), 2); } -TEST(ExtractOutsideCompilationForFunctionTest, XlaHostComputeRemoved) { +TEST_F(ExtractOutsideCompilationForFunctionTest, XlaHostComputeRemoved) { // Build the XLA computation func. // "const0" // "const1" (outside compilation cluster "0") @@ -432,7 +470,7 @@ TEST(ExtractOutsideCompilationForFunctionTest, XlaHostComputeRemoved) { NameAttrList name_attrs; name_attrs.set_name("cluster"); *name_attrs.mutable_attr() = attrs; - TF_CHECK_OK(ExtractOutsideCompilationForFunction( + TF_CHECK_OK(ExtractOutsideCompilationTest( "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten", "host_graph", host_compute_core, &fld, &shape_inference_graphs, &has_outside_compilation)); @@ -489,7 +527,7 @@ REGISTER_OP("XlaRecvFromHost") .Attr("key: string") .SetIsStateful(); -TEST(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInIf) { +TEST_F(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInIf) { // Build the XLA computation func. // "const0" (bool) // "const1" (int32) @@ -555,7 +593,7 @@ TEST(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInIf) { NameAttrList name_attrs; name_attrs.set_name("cluster"); *name_attrs.mutable_attr() = attrs; - TF_CHECK_OK(ExtractOutsideCompilationForFunction( + TF_CHECK_OK(ExtractOutsideCompilationTest( "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten", "host_graph", host_compute_core, &fld, &shape_inference_graphs, &has_outside_compilation)); @@ -651,7 +689,7 @@ TEST(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInIf) { } } -TEST(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInWhile) { +TEST_F(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInWhile) { // Build the XLA computation func. // "const0" (bool) // "while0" (input = "const0", cond = "cond_fn", body = "body_fn") @@ -714,7 +752,7 @@ TEST(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInWhile) { NameAttrList name_attrs; name_attrs.set_name("cluster"); *name_attrs.mutable_attr() = attrs; - TF_CHECK_OK(ExtractOutsideCompilationForFunction( + TF_CHECK_OK(ExtractOutsideCompilationTest( "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten", "host_graph", host_compute_core, &fld, &shape_inference_graphs, &has_outside_compilation)); @@ -782,4 +820,162 @@ TEST(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInWhile) { } } +TEST_F(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInFunction) { + // Build the XLA computation func. + // "const0" (int32) + // "fn" (input = "const0") + FunctionDefLibrary fdl; + { + tensorflow::Scope s = tensorflow::Scope::NewRootScope(); + Output arg = ops::_Arg(s.WithOpName("arg"), DT_INT32, 0); + Output identity = ops::Identity(s.WithOpName("identity"), arg); + ops::_Retval retval(s.WithOpName("retval"), identity, 0); + std::unique_ptr g(new Graph(OpRegistry::Global())); + TF_CHECK_OK(s.ToGraph(g.get())); + auto node_name_image = g->BuildNodeNameIndex(); + node_name_image["identity"]->AddAttr("_oc", "0"); + PartialTensorShape shape({2}); + node_name_image["identity"]->AddAttr( + kXlaInferredShapesAttrName, std::vector{shape}); + + FunctionDef *true_fn_fdef = fdl.add_function(); + TF_CHECK_OK(GraphToFunctionDef(*g, "fn", true_fn_fdef)); + } + FunctionLibraryDefinition fld(OpRegistry::Global(), fdl); + { + std::unique_ptr g(new Graph(&fld)); + + tensorflow::TensorProto tensor_proto; + tensor_proto.set_dtype(tensorflow::DT_INT32); + tensorflow::TensorShapeProto shape; + shape.add_dim()->set_size(2); + *tensor_proto.mutable_tensor_shape() = shape; + for (int i = 0; i < 2; ++i) { + tensor_proto.add_int_val(1); + } + NodeDef const_def; + TF_CHECK_OK(NodeDefBuilder("const", "Const") + .Attr("dtype", DT_INT32) + .Attr("value", tensor_proto) + .Finalize(&const_def)); + Status s; + Node *const_node = g->AddNode(const_def, &s); + TF_CHECK_OK(s); + + NodeDef fn_def; + TF_CHECK_OK(NodeDefBuilder("fn", "fn", &fld) + .Input("const", 0, DT_INT32) + .Finalize(&fn_def)); + Node *fn_node = g->AddNode(fn_def, &s); + TF_CHECK_OK(s); + g->AddEdge(const_node, 0, fn_node, 0); + + NodeDef ret_def; + TF_CHECK_OK(NodeDefBuilder("ret", "_Retval") + .Attr("index", 0) + .Attr("T", DT_INT32) + .Input("fn", 0, DT_INT32) + .Finalize(&ret_def)); + Node *ret_node = g->AddNode(ret_def, &s); + TF_CHECK_OK(s); + g->AddEdge(fn_node, 0, ret_node, 0); + + FunctionDef *xla_fdef = fdl.add_function(); + TF_CHECK_OK(GraphToFunctionDef(*g, "cluster", xla_fdef)); + TF_CHECK_OK(fld.AddFunctionDef(*xla_fdef)); + } + + protobuf::Map attrs; + std::map host_compute_core; + std::vector shape_inference_graphs; + bool has_outside_compilation; + NameAttrList name_attrs; + name_attrs.set_name("cluster"); + *name_attrs.mutable_attr() = attrs; + TF_CHECK_OK(ExtractOutsideCompilationTest( + "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten", "host_graph", + host_compute_core, &fld, &shape_inference_graphs, + &has_outside_compilation)); + + // Check host graph. + { + FunctionBody *host_fbody = nullptr; + AttrValue device_ordinal_temp_value; + device_ordinal_temp_value.set_i(0); + protobuf::Map host_func_attrs; + host_func_attrs["device_ordinal"] = device_ordinal_temp_value; + TF_CHECK_OK(FunctionDefToBodyHelper( + *fld.Find("host_graph"), AttrSlice(&host_func_attrs), &fld, + [&](const string &op, const OpDef **sig) { + return fld.LookUpOpDef(op, sig); + }, + &host_fbody)); + std::unique_ptr host_fbody_deleter(host_fbody); + Graph *host_graph = host_fbody->graph; + auto node_name_index = host_graph->BuildNodeNameIndex(); + + // Verify we have call node for outside compilation in `fn`. + Node *call_node = node_name_index["oc_call_fn"]; + EXPECT_NE(call_node, nullptr); + + FunctionBody *call_fbody = nullptr; + TF_CHECK_OK(FunctionDefToBodyHelper( + *fld.Find("oc_func_call_host_fn"), AttrSlice(&host_func_attrs), &fld, + [&](const string &op, const OpDef **sig) { + return fld.LookUpOpDef(op, sig); + }, + &call_fbody)); + std::unique_ptr call_fbody_deleter(call_fbody); + + // Verify we have _XlaRecvAtHost and _XlaSendFromHost nodes. + bool has_recv = false, has_send = false; + for (Node *n : call_fbody->graph->nodes()) { + if (n->type_string() == "_XlaRecvAtHost") { + has_recv = true; + } else if (n->type_string() == "_XlaSendFromHost") { + has_send = true; + } + } + EXPECT_TRUE(has_recv); + EXPECT_TRUE(has_send); + } + + // Check XLA graph. + { + FunctionBody *xla_fbody = nullptr; + TF_CHECK_OK(FunctionDefToBodyHelper( + *fld.Find("cluster_rewritten"), AttrSlice(), &fld, + [&](const string &op, const OpDef **sig) { + return fld.LookUpOpDef(op, sig); + }, + &xla_fbody)); + std::unique_ptr xla_fbody_deleter(xla_fbody); + Graph *xla_graph = xla_fbody->graph; + auto node_name_index = xla_graph->BuildNodeNameIndex(); + + // Check that we have call node. + Node *fn_node = node_name_index["fn"]; + EXPECT_NE(fn_node, nullptr); + EXPECT_EQ(fn_node->type_string(), "fn_oc"); + + FunctionBody *call_fbody = nullptr; + TF_CHECK_OK(FunctionDefToBodyHelper( + *fld.Find("fn_oc"), AttrSlice(), &fld, + [&](const string &op, const OpDef **sig) { + return fld.LookUpOpDef(op, sig); + }, + &call_fbody)); + std::unique_ptr call_fbody_deleter(call_fbody); + + // Verify we have XlaHostCompute nodes. + bool has_hc = false; + for (Node *n : call_fbody->graph->nodes()) { + if (n->type_string() == "XlaHostCompute") { + has_hc = true; + } + } + EXPECT_TRUE(has_hc); + } +} + } // namespace tensorflow diff --git a/tensorflow/compiler/jit/flags.cc b/tensorflow/compiler/jit/flags.cc index 98e344b3a080aa8aab27cd41564a90427bac151e..fba69dfccc31e01e73d8f86006b41ce5e3283f15 100644 --- a/tensorflow/compiler/jit/flags.cc +++ b/tensorflow/compiler/jit/flags.cc @@ -68,7 +68,12 @@ void AppendMarkForCompilationPassFlagsInternal(std::vector* flag_list) { Flag("tf_xla_fusion_only", &mark_for_compilation_flags->tf_xla_fusion_only, "enable fusion of element-wise operations only using XLA when " - "global_jit_level is ON*.")}; + "global_jit_level is ON*."), + Flag("tf_xla_disable_deadness_safety_checks_for_debugging", + &mark_for_compilation_flags + ->tf_xla_disable_deadness_safety_checks_for_debugging, + "Disable deadness related safety checks when clustering (this is " + "unsound).")}; flag_list->insert(flag_list->end(), new_flags.begin(), new_flags.end()); } @@ -89,6 +94,8 @@ void AllocateAndParseFlags() { mark_for_compilation_flags->tf_xla_clustering_fuel = std::numeric_limits::max(); mark_for_compilation_flags->tf_xla_fusion_only = false; + mark_for_compilation_flags + ->tf_xla_disable_deadness_safety_checks_for_debugging = false; device_flags = new XlaDeviceFlags; device_flags->tf_xla_compile_on_demand = false; diff --git a/tensorflow/compiler/jit/flags.h b/tensorflow/compiler/jit/flags.h index 5ddea588eef5270880d91623dc05893da265960a..ed7810fcfd85c17db70d42e691446b60dc696939 100644 --- a/tensorflow/compiler/jit/flags.h +++ b/tensorflow/compiler/jit/flags.h @@ -25,27 +25,39 @@ namespace tensorflow { // Flags associated with the XLA bridge's mark_for_compilation_pass module. struct MarkForCompilationPassFlags { - int32 tf_xla_auto_jit; // Control compilation of operators into XLA - // computations on CPU and GPU devices. 0 = use - // ConfigProto setting; -1 = off; 1 = on for things - // very likely to be improved; 2 = on for everything. - // Experimental. - int32 tf_xla_min_cluster_size; // Minimum number of operators in an XLA - // compilation. Ignored for operators placed - // on an XLA device or operators explicitly - // marked for compilation. - int32 tf_xla_max_cluster_size; // Maximum number of operators in an XLA - // compilation. - bool tf_xla_clustering_debug; // Dump graphs during XLA compilation. - bool tf_xla_cpu_global_jit; // Enables global JIT compilation for CPU - // via SessionOptions. - int64 tf_xla_clustering_fuel; // "Compiler fuel" for clustering. Only this - // many ops will be marked as eligible for - // clustering. - bool tf_xla_fusion_only; // This flag is effective only when global_jit_level - // is set to ON* and overrides its behavior. If - // true, enable fusion of element-wise operations - // only using XLA. + // Control compilation of operators into XLA computations on CPU and GPU + // devices. 0 = use ConfigProto setting; -1 = off; 1 = on for things very + // likely to be improved; 2 = on for everything. + // + // Experimental. + int32 tf_xla_auto_jit; + + // Minimum number of operators in an XLA compilation. Ignored for operators + // placed on an XLA device or operators explicitly marked for compilation. + int32 tf_xla_min_cluster_size; + + // Maximum number of operators in an XLA compilation. + int32 tf_xla_max_cluster_size; + + // Dump graphs during XLA compilation. + bool tf_xla_clustering_debug; + + // Enables global JIT compilation for CPU via SessionOptions. + bool tf_xla_cpu_global_jit; + + // "Compiler fuel" for clustering. Only this many ops will be marked as + // eligible for clustering. + int64 tf_xla_clustering_fuel; + + // tf_xla_fusion_only is effective only when global_jit_level is set to ON* + // and overrides its behavior. If true, enable fusion of element-wise + // operations only using XLA. + bool tf_xla_fusion_only; + + // If tf_xla_disable_deadness_safety_checks_for_debugging is set to true then + // we do not do deadness related safety checks. This is unsound in general, + // but can be used as a debugging aid. + bool tf_xla_disable_deadness_safety_checks_for_debugging; }; // Flags associated with the XLA bridge's xla_device module. diff --git a/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.cc b/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.cc index ce53f70b79d97ab087fefe542920b33f883632a2..5287fd175df206970b9fa73bc6b0176eddcdcaa9 100644 --- a/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.cc +++ b/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.h" +#include #include "absl/algorithm/container.h" #include "absl/container/inlined_vector.h" #include "absl/strings/str_cat.h" @@ -144,7 +145,9 @@ SliceInputs MakeSliceIndexAndSizeInt64(const Scope& host_scope, // same constant value. This helps make the generated GraphDef more readable. class ConstantCache { public: - explicit ConstantCache(const Scope& s) : scope_(s) {} + explicit ConstantCache(const Scope& s, + const std::vector& control_deps) + : scope_(s), control_deps_(control_deps) {} Output Get1DHostConstant(int64 constant) { auto it = cache_.find(constant); @@ -152,6 +155,9 @@ class ConstantCache { Output new_const = ops::Const(scope_.WithOpName("const_", constant), {constant}); it = cache_.insert({constant, new_const}).first; + for (const Edge* e : control_deps_) { + scope_.graph()->AddControlEdge(e->src(), new_const.node()); + } } return it->second; } @@ -159,11 +165,13 @@ class ConstantCache { private: Scope scope_; std::unordered_map cache_; + std::vector control_deps_; }; // Returns a node computing the size of the Slice op with inputs `slice_inputs`. Status ComputeSliceSize(const Scope& host_scope, - const SliceInputs& slice_inputs, Output* size) { + const SliceInputs& slice_inputs, + std::vector control_deps, Output* size) { // If slice_size[i] >= 0 then slice_size[i] = slice_size[i]. // // If slice_size[i] == -1 then slice_size[i] = input_size[i] - @@ -183,7 +191,7 @@ Status ComputeSliceSize(const Scope& host_scope, ops::Shape(host_scope.WithOpName("input_shape"), slice_inputs.input, ops::Shape::OutType(DT_INT64)); - ConstantCache constant_pool(host_scope); + ConstantCache constant_pool(host_scope, control_deps); std::vector slice_size; for (int i = 0; i < slice_inputs.size_as_vector.size(); i++) { @@ -209,11 +217,16 @@ Status ComputeSliceSize(const Scope& host_scope, } // Trivial ConcatV2 nodes (with exactly one input) are disallowed. - *size = - slice_size.size() == 1 - ? slice_size[0] - : ops::Concat(host_scope.WithOpName("slice_size"), slice_size, - ops::Const(host_scope.WithOpName("concat_axis"), 0)); + if (slice_size.size() == 1) { + *size = slice_size[0]; + } else { + auto concat_axis = ops::Const(host_scope.WithOpName("concat_axis"), 0); + for (const Edge* e : control_deps) { + host_scope.graph()->AddControlEdge(e->src(), concat_axis.node()); + } + *size = ops::Concat(host_scope.WithOpName("slice_size"), slice_size, + concat_axis); + } return Status::OK(); } @@ -234,12 +247,21 @@ Status ConvertTensorFlowSliceToStaticShapedSlice( .NewSubScope(absl::StrCat(slice->name(), "/static_shaped_slice")); Scope host_scope = main_scope.WithAssignedDevice(host_name); + // In the future we may want to be clever here and avoid the extra Cast ops. SliceInputs slice_inputs_int64 = MakeSliceIndexAndSizeInt64(host_scope, slice_inputs); + // Create a list of all control dependencies to be copied when possibly + // replacing nodes related to slice_size. + Node* old_size; + std::vector old_size_ctrl_deps; + TF_RETURN_IF_ERROR(slice->input_node(2, &old_size)); + absl::c_copy_if(old_size->in_edges(), std::back_inserter(old_size_ctrl_deps), + [](const Edge* e) { return e->IsControlEdge(); }); + Output slice_size; - TF_RETURN_IF_ERROR( - ComputeSliceSize(host_scope, slice_inputs_int64, &slice_size)); + TF_RETURN_IF_ERROR(ComputeSliceSize(host_scope, slice_inputs_int64, + old_size_ctrl_deps, &slice_size)); *result = ops::Slice(main_scope.WithAssignedDevice(slice->assigned_device_name()) @@ -291,9 +313,9 @@ Status RewriteSlice(Graph* g, Node* slice, const SliceInputs& slice_inputs, return Status::OK(); } -// Return true if `n` is a slice we can rewrite to have a static shape +// Return true if `n` is a slice we should rewrite to have a static shape // (i.e. have the output shape only depend on the "size" input). -xla::StatusOr IsRewritableSlice(Node* n) { +xla::StatusOr ShouldRewriteSlice(Node* n) { if (n->type_string() != "Slice") { return false; } @@ -311,14 +333,20 @@ xla::StatusOr IsRewritableSlice(Node* n) { // If slice_size[i] < -1 for any i then executing the slice will throw an // error, and we don't do anything here. - return absl::c_all_of(slice_inputs->size_as_vector, - [](int64 size_i) { return size_i >= -1; }); + bool slice_size_has_error = absl::c_all_of( + slice_inputs->size_as_vector, [](int64 size_i) { return size_i >= -1; }); + if (!slice_size_has_error) { + return false; + } + + // No point in rewriting slices that have both size and begin as constants. + return !slice_inputs->begin.node()->IsConstant(); } Status FindAndRewriteSlices(Graph* g, bool* changed) { std::vector slices_to_rewrite; for (Node* n : g->nodes()) { - TF_ASSIGN_OR_RETURN(bool is_rewritable, IsRewritableSlice(n)); + TF_ASSIGN_OR_RETURN(bool is_rewritable, ShouldRewriteSlice(n)); if (is_rewritable) { slices_to_rewrite.push_back(n); } diff --git a/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass_test.cc b/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass_test.cc index a2f1b831ad7605237e23c15cc43b337e06265553..2add2c13f92f561904163012ee16cc17ce5badce 100644 --- a/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass_test.cc +++ b/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass_test.cc @@ -401,5 +401,57 @@ TEST(SliceToDynamicSliceRewriteTest, SliceWithSliceBegin) { Name("begin/static_shaped_slice/static_shaped_slice"))), _))); } + +// New constants being created need to have control dependencies copied to +// ensure correct control flow analysis in TF V2. +TEST(SliceToDynamicSliceRewriteTest, WithControlDepsToConstant) { + Scope root = Scope::NewRootScope() + .ExitOnError() + .WithAssignedDevice(kDeviceName) + .WithXlaCluster("cluster_0"); + + Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT); + Output begin = ops::Placeholder(root.WithOpName("begin"), DT_INT32); + Output size = ops::Const(root.WithOpName("size"), {-1}); + Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size); + + // Add an additional dependency that should still exist in with the new size + // variables. + Output dependency = ops::Placeholder(root.WithOpName("dependency"), DT_BOOL); + root.graph()->AddControlEdge(dependency.node(), size.node()); + + std::unique_ptr result; + TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result)); + + // Check that the new constants have control dependencies. + Node* const_0 = testing::FindNodeByName(result.get(), + "slice/static_shaped_slice/const_0"); + EXPECT_NE(const_0, nullptr); + EXPECT_THAT(const_0, + NodeWith(Op("Const"), CtrlDeps(NodeWith(Op("Placeholder"), + Name("dependency"))))); +} + +TEST(SliceToDynamicSliceRewriteTest, DontRewriteSliceWithConstBegin) { + Scope root = Scope::NewRootScope() + .ExitOnError() + .WithAssignedDevice(kDeviceName) + .WithXlaCluster("cluster_0"); + + Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT); + Output begin = ops::Const(root.WithOpName("begin"), {10, 10}); + Output size = ops::Const(root.WithOpName("size"), {-1, 500}); + Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size); + + std::unique_ptr result; + TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result)); + + Node* slice_node = testing::FindNodeByName(result.get(), "slice"); + EXPECT_THAT(slice_node, + NodeWith(Op("Slice"), Inputs(Out(NodeWith(Op("Placeholder"))), + Out(NodeWith(Op("Const"))), + Out(NodeWith(Op("Const")))))); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/jit/kernels/BUILD b/tensorflow/compiler/jit/kernels/BUILD index 0583774714c6db7a2fa515fc8a0d304e1898db97..d0fa2c40be9d6b13ec736a9d6483dae0b4f0f45e 100644 --- a/tensorflow/compiler/jit/kernels/BUILD +++ b/tensorflow/compiler/jit/kernels/BUILD @@ -19,12 +19,14 @@ cc_library( "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:tf2xla_util", "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//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", diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc index ad71df5a694a5f8da94675049df1062a7edb6253..997ef6e14bb9bd16ddac13eaf67368966818b29e 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.cc +++ b/tensorflow/compiler/jit/kernels/xla_ops.cc @@ -25,6 +25,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/core/common_runtime/dma_helper.h" #include "tensorflow/core/common_runtime/function.h" @@ -35,6 +36,8 @@ limitations under the License. #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" #include "tensorflow/core/platform/stream_executor_no_cuda.h" #include "tensorflow/core/util/stream_executor_util.h" @@ -304,10 +307,19 @@ void XlaLocalLaunchBase::Compute(OpKernelContext* ctx) { xla::LocalExecutable* executable; std::map variables; - OP_REQUIRES_OK( - ctx, CompileToLocalExecutable(ctx, function_, platform_info_, resources_, - constants_, /*lazy=*/false, &client, - &variables, &kernel, &executable)); + { + Status s = CompileToLocalExecutable( + ctx, function_, platform_info_, resources_, constants_, /*lazy=*/false, + &client, &variables, &kernel, &executable); + if (!s.ok() && (platform_info_.device_type().type_string() == DEVICE_CPU || + platform_info_.device_type().type_string() == DEVICE_GPU)) { + // Suggest auto jit if the failure was with GPU or CPU. + errors::AppendToMessage(&s, + xla::status_macros::kPossibleAutoJitAlternative); + } + + OP_REQUIRES_OK(ctx, s); + } se::Stream* stream = ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr; diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index 6618e3a58ab7b6374ed775cd6e4e18a6a4975588..20c2cd7e0561f92a01486102c4d2c572fd80c957 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -34,6 +34,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/common_runtime/function.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/graph_def_util.h" #include "tensorflow/core/framework/memory_types.h" #include "tensorflow/core/framework/node_def.pb.h" @@ -41,7 +42,7 @@ limitations under the License. #include "tensorflow/core/framework/types.h" #include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/graph/control_flow.h" -#include "tensorflow/core/kernels/bounds_check.h" +#include "tensorflow/core/graph/graph_constructor.h" #include "tensorflow/core/lib/gtl/cleanup.h" #include "tensorflow/core/lib/strings/stringprintf.h" #include "tensorflow/core/public/version.h" @@ -677,12 +678,28 @@ Status MarkForCompilationPass::Run( VLOG(1) << "flags->tf_xla_auto_jit = " << flags->tf_xla_auto_jit; const FunctionLibraryDefinition* fld = options.flib_def; + // Deadness analysis expects a graph with source and sink edges properly + // connected but sometimes the incoming graph does not follow this invariant. + // So fix up the source and sink edges before calling into deadness analysis. + FixupSourceAndSinkEdges(options.graph->get()); + std::unique_ptr deadness; { XLA_SCOPED_LOGGING_TIMER_LEVEL("DeadnessAnalysis", 1); TF_RETURN_IF_ERROR(DeadnessAnalysis::Run(**options.graph, &deadness)); } + bool deadness_analysis_disabled = + GetMarkForCompilationPassFlags() + ->tf_xla_disable_deadness_safety_checks_for_debugging; + + if (deadness_analysis_disabled) { + LOG(WARNING) << "Deadness analysis was manually disabled via " + "--tf_xla_disable_deadness_safety_checks_for_debugging; " + "auto-clustering " + "is unsound!"; + } + auto is_compilable = [&](const Node* node, const DeviceType& device_type) { const XlaOpRegistry::DeviceRegistration* registration; if (!XlaOpRegistry::GetCompilationDevice(device_type.type(), @@ -715,9 +732,12 @@ Status MarkForCompilationPass::Run( // and some are dead) then don't compile it. XLA cannot represent the // deadness semantics of these nodes correctly and auto-clustering these // nodes can cause deadness to propagate to nodes that should be live. - if (node->IsMerge() || deadness->HasInputsWithMismatchingDeadness(*node)) { - VLOG(2) << "Rejecting " << node->name() << ": mismatching deadness."; - return false; + if (!deadness_analysis_disabled) { + if (node->IsMerge() || + deadness->HasInputsWithMismatchingDeadness(*node)) { + VLOG(2) << "Rejecting " << node->name() << ": mismatching deadness."; + return false; + } } // Check for fusable ops only if requested. @@ -1145,6 +1165,27 @@ Status MarkForCompilationPass::RunImpl( if (flags->tf_xla_clustering_debug) { dump_graph::DumpGraphToFile("mark_for_compilation", **options.graph, options.flib_def); + + // We also dump out an annoated version of the TF graph where the nodes + // names are prefixed with the cluster names. This can help visualizing the + // clustering decisions on TensorBoard. + Graph new_graph((*options.graph)->op_registry()); + CopyGraph(**options.graph, &new_graph); + + for (Node* n : new_graph.nodes()) { + if (absl::optional cluster_name = + GetXlaClusterForNode(*n)) { + n->set_name(absl::StrCat(*cluster_name, "/", 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 + // specific class has edges to and from the same set of clusters. + n->set_name(absl::StrCat("unclustered/", n->name())); + } + } + + dump_graph::DumpGraphToFile("mark_for_compilation_annotated", new_graph, + options.flib_def); } VLogClusteringSummary(*graph); diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc index bf2c5508ea9e987e80093f4c2e15d3ff5191126f..c2b6250f738fafa35b2c5f79e97cf1281b50a316 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc @@ -151,7 +151,7 @@ TEST(XlaCompilationTest, CompilableCycles) { EXPECT_EQ(clusters["A"], clusters["C"]); } -TEST(XlaCompilationTest, Complex128Unsupported) { +TEST(XlaCompilationTest, StringUnsupported) { std::unique_ptr graph(new Graph(OpRegistry::Global())); GraphDef graphdef; { @@ -159,10 +159,10 @@ TEST(XlaCompilationTest, Complex128Unsupported) { Node* a = ops::SourceOp( "Const", builder.opts() .WithName("A") - .WithAttr("dtype", DT_COMPLEX128) - .WithAttr("value", Tensor(DT_COMPLEX128, TensorShape()))); - Node* b = ops::UnaryOp("Neg", a, builder.opts().WithName("B")); - ops::BinaryOp("MatMul", a, b, builder.opts().WithName("C")); + .WithAttr("dtype", DT_STRING) + .WithAttr("value", Tensor(DT_STRING, TensorShape()))); + Node* b = ops::UnaryOp("EncodeBase64", a, builder.opts().WithName("B")); + ops::BinaryOp("StringSplit", a, b, builder.opts().WithName("C")); TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get())); } diff --git a/tensorflow/compiler/jit/partially_decluster_pass_test.cc b/tensorflow/compiler/jit/partially_decluster_pass_test.cc index 38a54cc5efae35ad77b6dc8039c653e920cfc071..1d81a8f4fcbf050663626b1f7660afd71f4027bc 100644 --- a/tensorflow/compiler/jit/partially_decluster_pass_test.cc +++ b/tensorflow/compiler/jit/partially_decluster_pass_test.cc @@ -33,7 +33,6 @@ limitations under the License. #include "tensorflow/core/graph/graph_constructor.h" #include "tensorflow/core/graph/graph_def_builder.h" #include "tensorflow/core/graph/graph_def_builder_util.h" -#include "tensorflow/core/grappler/optimizers/data/graph_utils.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/test.h" diff --git a/tensorflow/compiler/jit/xla_cluster_util.cc b/tensorflow/compiler/jit/xla_cluster_util.cc index fef28fc810cb4e544fe3f271f0b96cebd8a96779..3adcfef4dacecb343812cefc3a893a65c74ca101 100644 --- a/tensorflow/compiler/jit/xla_cluster_util.cc +++ b/tensorflow/compiler/jit/xla_cluster_util.cc @@ -19,9 +19,9 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "tensorflow/compiler/jit/resource_operation_safety_analysis.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/graph/control_flow.h" -#include "tensorflow/core/kernels/bounds_check.h" #include "tensorflow/core/util/device_name_utils.h" namespace tensorflow { @@ -43,7 +43,7 @@ string DescribeCycle(const GraphCycles* cycles, const Graph& graph, int src, return ""; } - auto node_name = [cycles, &graph](int node_id) { + auto node_name = [&graph](int node_id) { if (!FastBoundsCheck(node_id, graph.num_node_ids())) { return string("(null)"); } diff --git a/tensorflow/compiler/jit/xla_compilation_cache.cc b/tensorflow/compiler/jit/xla_compilation_cache.cc index d5a1ab0850c0af556995eddb04e0ab3a6c542cc2..611515cf33bc1abe21e06eb7f1513800276e095b 100644 --- a/tensorflow/compiler/jit/xla_compilation_cache.cc +++ b/tensorflow/compiler/jit/xla_compilation_cache.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include "absl/strings/str_cat.h" +#include "absl/strings/str_join.h" #include "tensorflow/compiler/tf2xla/dump_graph.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/type_util.h" @@ -62,7 +63,7 @@ XlaCompilationCache::~XlaCompilationCache() { // about? } -string XlaCompilationCache::DebugString() { +string XlaCompilationCache::DebugString() const { return "XLA JIT compilation cache"; } @@ -70,9 +71,9 @@ string XlaCompilationCache::DebugString() { // arguments in the supplied list. string XlaCompilationCache::Signature::HumanString() const { string result = name; - for (const auto& a : arg_types) { - absl::StrAppend(&result, ",", DataTypeString(a.first), - a.second.DebugString()); + for (const auto& a : arg_shapes) { + absl::StrAppend(&result, ",", DataTypeString(a.first)); + absl::StrAppend(&result, " [", absl::StrJoin(a.second, ","), "]"); } for (const auto& v : arg_values) { @@ -83,7 +84,7 @@ string XlaCompilationCache::Signature::HumanString() const { bool XlaCompilationCache::Signature::operator==(const Signature& other) const { if (name != other.name) return false; - if (arg_types != other.arg_types) return false; + if (arg_shapes != other.arg_shapes) return false; if (arg_values.size() != other.arg_values.size()) return false; for (int i = 0; i < arg_values.size(); ++i) { @@ -99,10 +100,10 @@ bool XlaCompilationCache::Signature::operator==(const Signature& other) const { uint64 XlaCompilationCache::Signature::Hash::operator()( const XlaCompilationCache::Signature& signature) const { uint64 h = std::hash()(signature.name); - for (const auto& arg : signature.arg_types) { + for (const auto& arg : signature.arg_shapes) { h = Hash64Combine(h, std::hash()(static_cast(arg.first))); - h = Hash64Combine(h, std::hash()(arg.second.dims())); - for (int dim : arg.second.dim_sizes()) { + h = Hash64Combine(h, std::hash()(arg.second.size())); + for (int dim : arg.second) { h = Hash64Combine(h, std::hash()(dim)); } } @@ -126,7 +127,7 @@ XlaCompilationCache::BuildSignature( break; case XlaCompiler::Argument::kParameter: case XlaCompiler::Argument::kResource: - signature.arg_types.emplace_back(arg.type, arg.shape); + signature.arg_shapes.emplace_back(arg.type, arg.DimensionSizes()); break; default: return errors::InvalidArgument( diff --git a/tensorflow/compiler/jit/xla_compilation_cache.h b/tensorflow/compiler/jit/xla_compilation_cache.h index 846d0c963dbfdf55f51120f2f138d12f5f63839b..7748b4700f39da4f952278ca6c6d2cadff4d3fb8 100644 --- a/tensorflow/compiler/jit/xla_compilation_cache.h +++ b/tensorflow/compiler/jit/xla_compilation_cache.h @@ -88,14 +88,16 @@ class XlaCompilationCache : public ResourceBase { xla::LocalClient* client() const { return client_; } const DeviceType& device_type() const { return device_type_; } - string DebugString() override; + string DebugString() const override; // Describes the types, shapes and any compile-time constant arguments // to a kernel. Key that uniquely identifies a compilation output. struct Signature { string name; - std::vector> arg_types; + // List of Tensor types & shapes for compile-time constant arguments to the + // compilation, ordered by argument number. + std::vector>> arg_shapes; // List of Tensor values for compile-time constant arguments to the // compilation, ordered by argument number. Tensors must be in host memory. diff --git a/tensorflow/compiler/jit/xla_cpu_device.cc b/tensorflow/compiler/jit/xla_cpu_device.cc index e9770647e7ba96cc1db026d12d5f11f52ce98d35..94dc61d55fb047c0ea81d98fde24cb55387c27d7 100644 --- a/tensorflow/compiler/jit/xla_cpu_device.cc +++ b/tensorflow/compiler/jit/xla_cpu_device.cc @@ -83,9 +83,9 @@ 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_BOOL}}; + DT_HALF, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64, DT_COMPLEX128, DT_BOOL}}; 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 77cd2f44628677942da9e576070d1d295194cead..92229842bdbced6431fd5b3e158f275a41819728 100644 --- a/tensorflow/compiler/jit/xla_device.cc +++ b/tensorflow/compiler/jit/xla_device.cc @@ -219,9 +219,6 @@ XlaDevice::XlaDevice(const SessionOptions& session_options, XlaDevice::~XlaDevice() { VLOG(1) << "Destroying XLA device " << jit_device_name_ << " " << this; mutex_lock lock(mu_); - while (outstanding_asynchronous_operations_ > 0) { - outstanding_asynchronous_operations_cv_.wait(lock); - } if (device_context_) { device_context_->Unref(); } @@ -398,12 +395,6 @@ Status XlaDevice::Sync() { if (!stream) return Status::OK(); Status status = stream->BlockHostUntilDone(); - { - mutex_lock lock(mu_); - while (outstanding_asynchronous_operations_ > 0) { - outstanding_asynchronous_operations_cv_.wait(lock); - } - } TF_RETURN_IF_ERROR(status); if (!stream->ok()) { return errors::Internal("XlaDevice::Sync() failed."); @@ -412,6 +403,8 @@ Status XlaDevice::Sync() { return Status::OK(); } +// TODO(b/112409994): This is no longer necessary. Consolidate it with the +// synchronous version. void XlaDevice::Sync(const DoneCallback& done) { VLOG(1) << "XlaDevice::Sync (asynchronous)"; std::shared_ptr stream; @@ -424,14 +417,20 @@ void XlaDevice::Sync(const DoneCallback& done) { return; } + // The call to ThenEnqueueOnBackgroundThread below enqueues a host callback at + // the end of the stream, after everything that has already been enqueued + // there at this moment. When the host callback is called, everything before + // it must have already finished, and the host callback will then place the + // task below onto a background thread. (See the implementation of + // ThenEnqueueOnBackgroundThread for details.) Therefore, when the done + // callback is finally called from that background thread, we know for sure + // that everything enqueued onto the stream (i.e., the device) at this very + // moment--when ThenEnqueueOnBackgroundThread is called--will have finished. + // This achieves a device-wide sync. stream->ThenEnqueueOnBackgroundThread( - [this, stream, done](se::StreamExecutor*) { + [stream, done](se::StreamExecutor*) { tracing::ScopedActivity activity("XlaDevice::Sync::Callback", /*is_expensive=*/true); - mutex_lock lock(mu_); - while (outstanding_asynchronous_operations_ > 0) { - outstanding_asynchronous_operations_cv_.wait(lock); - } done(stream->ok() ? Status::OK() : errors::Internal("XlaDevice::Sync() failed.")); }); @@ -470,57 +469,27 @@ Status XlaDevice::MakeTensorFromProto(const TensorProto& tensor_proto, return status; } -void XlaDevice::SetRequiresSyncOnCompletion(bool sync_on_completion) { +void XlaDevice::SetAllowsSyncOnCompletion(bool sync_on_completion) { mutex_lock lock(mu_); sync_on_completion_ = sync_on_completion; } -bool XlaDevice::RequiresSyncOnCompletion() const { +bool XlaDevice::AllowsSyncOnCompletion() const { mutex_lock lock(mu_); return sync_on_completion_; } -XlaDevice::AsynchronousOperationHandle::AsynchronousOperationHandle( - XlaDevice* device) - : device_(device) { - mutex_lock lock(device_->mu_); - ++device_->outstanding_asynchronous_operations_; -} - -XlaDevice::AsynchronousOperationHandle::~AsynchronousOperationHandle() { - if (device_) { - mutex_lock lock(device_->mu_); - --device_->outstanding_asynchronous_operations_; - device_->outstanding_asynchronous_operations_cv_.notify_all(); +Status XlaDevice::RefreshStatus() { + std::shared_ptr stream; + { + mutex_lock lock(mu_); + stream = stream_; } -} - -XlaDevice::AsynchronousOperationHandle::AsynchronousOperationHandle( - const XlaDevice::AsynchronousOperationHandle& other) - : device_(other.device_) { - mutex_lock lock(device_->mu_); - ++device_->outstanding_asynchronous_operations_; -} - -XlaDevice::AsynchronousOperationHandle::AsynchronousOperationHandle( - XlaDevice::AsynchronousOperationHandle&& other) - : device_(other.device_) { - other.device_ = nullptr; -} - -XlaDevice::AsynchronousOperationHandle& XlaDevice::AsynchronousOperationHandle:: -operator=(const XlaDevice::AsynchronousOperationHandle& other) { - device_ = other.device_; - mutex_lock lock(device_->mu_); - ++device_->outstanding_asynchronous_operations_; - return *this; -} - -XlaDevice::AsynchronousOperationHandle& XlaDevice::AsynchronousOperationHandle:: -operator=(XlaDevice::AsynchronousOperationHandle&& other) { - device_ = other.device_; - other.device_ = nullptr; - return *this; + if (!stream) { + return Status::OK(); + } + // Stream status is XlaDevice status, no extra operations needed. + return stream->RefreshStatus(); } XlaDeviceOpRegistrations* RegisterXlaDeviceKernels(const char* device, diff --git a/tensorflow/compiler/jit/xla_device.h b/tensorflow/compiler/jit/xla_device.h index 45f18ac9ee6d403c192bd421d7823f2d408d994b..5fe1290fa03f2b1f9d90e36dbc5769b3c2728c8d 100644 --- a/tensorflow/compiler/jit/xla_device.h +++ b/tensorflow/compiler/jit/xla_device.h @@ -167,35 +167,13 @@ class XlaDevice : public LocalDevice { Status UseGpuDeviceInfo() LOCKS_EXCLUDED(mu_); // Instructs this XlaDevice to return 'sync_on_completion' for - // RequiresSyncOnCompletion(). - void SetRequiresSyncOnCompletion(bool sync_on_completion) LOCKS_EXCLUDED(mu_); + // AllowsSyncOnCompletion(). + void SetAllowsSyncOnCompletion(bool sync_on_completion) LOCKS_EXCLUDED(mu_); + bool AllowsSyncOnCompletion() const override LOCKS_EXCLUDED(mu_); - bool RequiresSyncOnCompletion() const override LOCKS_EXCLUDED(mu_); - - // A simple RAII handle. On construction the device's - // outstanding_asynchronous_operations_ field is incremented; on destruction - // it is decremented. - class AsynchronousOperationHandle { - public: - AsynchronousOperationHandle(XlaDevice* device); - ~AsynchronousOperationHandle(); - AsynchronousOperationHandle(const AsynchronousOperationHandle& other); - AsynchronousOperationHandle(AsynchronousOperationHandle&& other); - AsynchronousOperationHandle& operator=( - const AsynchronousOperationHandle& other); - AsynchronousOperationHandle& operator=(AsynchronousOperationHandle&& other); - - private: - XlaDevice* device_ = nullptr; - }; - - AsynchronousOperationHandle CreateAsynchronousOperationHandle() { - return AsynchronousOperationHandle(this); - } + Status RefreshStatus() override LOCKS_EXCLUDED(mu_); private: - friend class AsynchronousOperationHandle; - xla::LocalClient* client() const; Allocator* GetAllocatorLocked(AllocatorAttributes attr) EXCLUSIVE_LOCKS_REQUIRED(mu_); @@ -255,14 +233,9 @@ class XlaDevice : public LocalDevice { // Thread pool used for running closures std::unique_ptr thread_pool_; - // True if the device requires XlaDevice::Sync to be called on completion + // True if the device allows XlaDevice::Sync to be called on completion // regardless of status. - bool sync_on_completion_ GUARDED_BY(mu_) = false; - - // Count of outstanding asynchronous operations which must be zero on Sync() - // completion. - int64 outstanding_asynchronous_operations_ GUARDED_BY(mu_) = 0; - condition_variable outstanding_asynchronous_operations_cv_; + bool sync_on_completion_ GUARDED_BY(mu_) = true; // Set of devices to use. This controls which of the devices on the given // platform will have resources allocated. For GPUs this will be diff --git a/tensorflow/compiler/jit/xla_device_context.cc b/tensorflow/compiler/jit/xla_device_context.cc index 1f3afe8822d441a5ce37617fe18d7767e9bc72e4..28681bb8b03dbf97e8145972f9a04b5855fafdae 100644 --- a/tensorflow/compiler/jit/xla_device_context.cc +++ b/tensorflow/compiler/jit/xla_device_context.cc @@ -131,7 +131,7 @@ void XlaDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor, xla::ShapeUtil::MakeShape(shape.element_type(), xla::AsInt64Slice(shape.dimensions()))); - VLOG(1) << "Transfer to device as literal: " << literal.ToString() << " " + VLOG(2) << "Transfer to device as literal: " << literal.ToString() << " " << xla_tensor->shaped_buffer().ToString(); if (UseMultipleStreams() && !transfer_manager_->CanShapedBufferBeAccessedNow( @@ -214,7 +214,7 @@ void XlaDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor, device_to_host_stream_.get(), xla_tensor->shaped_buffer(), literal, [ref, xla_tensor, done](xla::Status status) { done([&]() -> Status { - VLOG(1) << "Transfer from device as literal: " + VLOG(2) << "Transfer from device as literal: " << xla_tensor->shaped_buffer().ToString(); return status; }()); diff --git a/tensorflow/compiler/jit/xla_device_ops.h b/tensorflow/compiler/jit/xla_device_ops.h index 927f983ba9ef23c8509523f42366c0c89c29db9f..f201f62a78ce9d9599b2397a5f108e335469445a 100644 --- a/tensorflow/compiler/jit/xla_device_ops.h +++ b/tensorflow/compiler/jit/xla_device_ops.h @@ -241,6 +241,8 @@ class XlaAssignVariableOp : public OpKernel { data::AnonymousIteratorHandleOp); \ REGISTER_KERNEL_BUILDER(Name("IteratorGetNext").Device(DEVICE), \ data::IteratorGetNextOp); \ + REGISTER_KERNEL_BUILDER(Name("IteratorGetNextAsOptional").Device(DEVICE), \ + data::IteratorGetNextAsOptionalOp); \ REGISTER_KERNEL_BUILDER(Name("IteratorGetNextSync").Device(DEVICE), \ data::IteratorGetNextSyncOp); \ REGISTER_KERNEL_BUILDER(Name("IteratorToStringHandle") \ diff --git a/tensorflow/compiler/jit/xla_interpreter_device.cc b/tensorflow/compiler/jit/xla_interpreter_device.cc index 4007309ed1c57b663dca5bac0df11260bf1327f3..e1a582406153d2af447fa9d4ebcaf0bf0842b132 100644 --- a/tensorflow/compiler/jit/xla_interpreter_device.cc +++ b/tensorflow/compiler/jit/xla_interpreter_device.cc @@ -26,9 +26,9 @@ namespace tensorflow { const char* const DEVICE_XLA_INTERPRETER = "XLA_INTERPRETER"; const char* const DEVICE_INTERPRETER_XLA_JIT = "XLA_INTERPRETER_JIT"; -constexpr std::array kExecAllTypes = { +constexpr std::array kExecAllTypes = { {DT_INT8, DT_INT32, DT_INT64, DT_HALF, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64, - DT_BOOL, DT_BFLOAT16}}; + DT_COMPLEX128, DT_BOOL, DT_BFLOAT16}}; class XlaInterpreterDeviceFactory : public DeviceFactory { public: diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc index 4e3cb0238f73c47f5f3c76c2f3835a6a6d381d26..c64981053fad2dbf1e8bcd623a940ded8b4d9150 100644 --- a/tensorflow/compiler/jit/xla_launch_util.cc +++ b/tensorflow/compiler/jit/xla_launch_util.cc @@ -377,7 +377,7 @@ Status XlaComputationLaunchContext::PopulateOutputs( } if (VLOG_IS_ON(3)) { - VLOG(3) << ctx->mutable_output(i)->DebugString(); + VLOG(3) << ctx->mutable_output(i)->DeviceSafeDebugString(); } } diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index f80cb1812f00d36ddb7c28ae0e77c58498058ef3..9b6ca4092c3177ac26503add13bce25d2c0bb820 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -72,7 +72,7 @@ py_test( tf_xla_py_test( name = "adadelta_test", - size = "large", + size = "medium", srcs = ["adadelta_test.py"], deps = [ ":xla_test", @@ -243,6 +243,7 @@ tf_xla_py_test( ":xla_test", "//tensorflow/python:array_ops", "//tensorflow/python:framework", + "//tensorflow/python:map_fn", "//tensorflow/python:math_ops", "//tensorflow/python:platform_test", "//tensorflow/python:training", @@ -278,10 +279,9 @@ tf_xla_py_test( ], ) -# This test is large because occasionally the cpu test is long for testConcatLargeNumberOfTensors tf_xla_py_test( name = "concat_ops_test", - size = "large", + size = "medium", srcs = ["concat_ops_test.py"], deps = [ ":xla_test", @@ -407,7 +407,7 @@ tf_xla_py_test( tf_xla_py_test( name = "eager_test", - size = "large", + size = "medium", srcs = ["eager_test.py"], deps = [ ":xla_test", diff --git a/tensorflow/compiler/tests/adadelta_test.py b/tensorflow/compiler/tests/adadelta_test.py index b7b7fda293b69d6f0cec61d0d234277636a3670d..6cf16cc07ff503c4f3e008cfb720224abe5e9166 100644 --- a/tensorflow/compiler/tests/adadelta_test.py +++ b/tensorflow/compiler/tests/adadelta_test.py @@ -32,10 +32,18 @@ class AdadeltaOptimizerTest(xla_test.XLATestCase): def testBasic(self): num_updates = 4 # number of ADADELTA steps to perform + if "CPU" in self.device: + # To avoid timeout on CPU. + all_grad = [0.2, 0.01] + all_lr = [1.0, 0.1] + else: + all_grad = [0.2, 0.1, 0.01] + all_lr = [1.0, 0.5, 0.1] + for dtype in self.float_types: with self.cached_session(), self.test_scope(): - for grad in [0.2, 0.1, 0.01]: - for lr in [1.0, 0.5, 0.1]: + for grad in all_grad: + for lr in all_lr: var0_init = [1.0, 2.0] var1_init = [3.0, 4.0] var0 = resource_variable_ops.ResourceVariable( diff --git a/tensorflow/compiler/tests/binary_ops_test.py b/tensorflow/compiler/tests/binary_ops_test.py index 9a5423c1b2a5df7880453cbb328f6a8174066255..c829c50b5518b29c96c0b0117a6cd143911bd1fc 100644 --- a/tensorflow/compiler/tests/binary_ops_test.py +++ b/tensorflow/compiler/tests/binary_ops_test.py @@ -311,6 +311,30 @@ class BinaryOpsTest(xla_test.XLATestCase): dtype(7), expected=np.array([[-6], [-5]], dtype=dtype)) + if dtype in [np.float32, np.float64]: + x = np.array([ + -0.0, 0.0, -0.0, +0.0, np.inf, np.inf, -np.inf, -np.inf, 2.0, 2.0, + 1.0 + ], + dtype=dtype) + y = np.array( + [-0.0, 0.0, +0.0, -0.0, 1.0, -1.0, 1.0, -1.0, 2.0, 1.0, 2.0], + dtype=dtype) + expected = np.nextafter(x, y) + + # We use assertAllEqual to expose any bugs hidden by relative or + # absolute error tolerances. + def NextAfterEqualityTest(result, expected, rtol): + del rtol + return self.assertAllEqual(result, expected) + + self._testBinary( + math_ops.nextafter, + x, + y, + expected=expected, + equality_test=NextAfterEqualityTest) + # min/max not supported for complex if dtype not in self.complex_types | {np.uint8, np.int8}: self._testBinary( @@ -400,7 +424,7 @@ class BinaryOpsTest(xla_test.XLATestCase): def testComplexOps(self): for dtype in self.complex_types: - ctypes = {np.complex64: np.float32} + ctypes = {np.complex64: np.float32, np.complex128: np.float64} self._testBinary( math_ops.complex, np.array([[[[-1, 2], [2, 0]]]], dtype=ctypes[dtype]), diff --git a/tensorflow/compiler/tests/build_defs.bzl b/tensorflow/compiler/tests/build_defs.bzl index 447a7de2cb6526a5dcf7789d4f2bffb5e733e8c0..ed580f95b6c2f57dfdf46cfcd64cabb452980c5d 100644 --- a/tensorflow/compiler/tests/build_defs.bzl +++ b/tensorflow/compiler/tests/build_defs.bzl @@ -5,6 +5,7 @@ load("//tensorflow/compiler/tests:plugin.bzl", "plugins") load( "//tensorflow/core:platform/default/build_config_root.bzl", "tf_cuda_tests_tags", + "tf_exec_compatible_with", ) def all_backends(): @@ -64,7 +65,7 @@ def tf_xla_py_test( if backend == "cpu": backend_args += [ "--test_device=XLA_CPU", - "--types=DT_HALF,DT_FLOAT,DT_DOUBLE,DT_UINT8,DT_QUINT8,DT_INT8,DT_QINT8,DT_INT32,DT_QINT32,DT_INT64,DT_BOOL,DT_COMPLEX64", + "--types=DT_HALF,DT_FLOAT,DT_DOUBLE,DT_UINT8,DT_QUINT8,DT_INT8,DT_QINT8,DT_INT32,DT_QINT32,DT_INT64,DT_BOOL,DT_COMPLEX64,DT_COMPLEX128", ] elif backend == "gpu": backend_args += [ @@ -84,6 +85,7 @@ def tf_xla_py_test( else: fail("Unknown backend {}".format(backend)) + test_tags = tags + backend_tags native.py_test( name = test_name, srcs = srcs, @@ -92,7 +94,8 @@ def tf_xla_py_test( main = "{}.py".format(name) if main == None else main, data = data + backend_data, deps = deps + backend_deps, - tags = tags + backend_tags, + tags = test_tags, + exec_compatible_with = tf_exec_compatible_with({"tags": test_tags}), **kwargs ) test_names.append(test_name) diff --git a/tensorflow/compiler/tests/categorical_op_test.py b/tensorflow/compiler/tests/categorical_op_test.py index 5d5e486f616937601214aa169a4c329ab78932c8..eec69ea7d2d9af9ff570f927fb25b668ccce2b97 100644 --- a/tensorflow/compiler/tests/categorical_op_test.py +++ b/tensorflow/compiler/tests/categorical_op_test.py @@ -119,7 +119,7 @@ class CategoricalTest(xla_test.XLATestCase): def testSamplingCorrectness(self): np.random.seed(1618) # Make it reproducible. - num_samples = 21000 + num_samples = 40000 rand_probs = np.random.dirichlet([1., 1., 2., 3.]) rand_probs2 = np.random.dirichlet([1., 4., 5.], size=3) # batched diff --git a/tensorflow/compiler/tests/concat_ops_test.py b/tensorflow/compiler/tests/concat_ops_test.py index 2187f57960f80300d631bdc7eb8fe5e9c8dddeea..76750decd2963ea12680a46d7340f48e8b011fa9 100644 --- a/tensorflow/compiler/tests/concat_ops_test.py +++ b/tensorflow/compiler/tests/concat_ops_test.py @@ -294,6 +294,9 @@ class ConcatTest(xla_test.XLATestCase): # The purpose of this is to ensure that XLA on GPU will not run out of memory # with too many arguments. def testConcatLargeNumberOfTensors(self): + if "CPU" in self.device: + self.skipTest("This test can time out on CPU, so we will just allow " + "other backends to catch this specific error.") with self.cached_session(): with self.test_scope(): for concat_dim in range(2): diff --git a/tensorflow/compiler/tests/dense_layer_test.py b/tensorflow/compiler/tests/dense_layer_test.py index bf5ea7b1fb6fb3c774c4db20d059f131990d20d3..b7d08df9f7d144b71fd0b09535e10b8f596ea6ca 100644 --- a/tensorflow/compiler/tests/dense_layer_test.py +++ b/tensorflow/compiler/tests/dense_layer_test.py @@ -72,7 +72,7 @@ class DenseLayerTest(test.TestCase): x = array_ops.placeholder(shape=[None, None, 3], dtype=np.float32) y = layers.dense(x, 3) - self.evaluate(variables.initialize_all_variables()) + self.evaluate(variables.global_variables_initializer()) run_metadata = config_pb2.RunMetadata() test_utils.RunWithWarmup( sess, @@ -97,7 +97,7 @@ class DenseLayerTest(test.TestCase): with jit_scope(): y = layers.dense(x, 3) - self.evaluate(variables.initialize_all_variables()) + self.evaluate(variables.global_variables_initializer()) run_metadata = config_pb2.RunMetadata() test_utils.RunWithWarmup( sess, @@ -126,7 +126,7 @@ class DenseLayerTest(test.TestCase): with jit_scope(): y = layers.dense(x, 3) - self.evaluate(variables.initialize_all_variables()) + self.evaluate(variables.global_variables_initializer()) run_metadata = config_pb2.RunMetadata() test_utils.RunWithWarmup( sess, diff --git a/tensorflow/compiler/tests/eager_test.py b/tensorflow/compiler/tests/eager_test.py index 2af32b537ba53723370faf81aebf308a465718c7..632eccbb097b4e84f10f926e89d7fa439c8a38cd 100644 --- a/tensorflow/compiler/tests/eager_test.py +++ b/tensorflow/compiler/tests/eager_test.py @@ -24,6 +24,7 @@ from tensorflow.compiler.tests import xla_test from tensorflow.core.protobuf import config_pb2 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.framework import constant_op from tensorflow.python.framework import dtypes @@ -31,7 +32,9 @@ from tensorflow.python.framework import ops from tensorflow.python.layers import convolutional from tensorflow.python.layers import pooling from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import embedding_ops +from tensorflow.python.ops import functional_ops from tensorflow.python.ops import gen_random_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops @@ -463,7 +466,7 @@ class EagerFunctionTest(xla_test.XLATestCase): def f(x, y): return x[0::2, y:, ...] - x = array_ops.ones([2, 3, 4]) + x = array_ops.ones([2, 3, 4], dtype=dtypes.float32) y = array_ops.ones([], dtype=dtypes.int32) with backprop.GradientTape() as tape: tape.watch(x) @@ -479,15 +482,15 @@ class EagerFunctionTest(xla_test.XLATestCase): @function.defun def times_two(x): - return 2 * x + return 2. * x @function.defun def two_x_plus_1(x): - return times_two(x) + 1 + return times_two(x) + 1. - x = constant_op.constant([2, 3, 4]) + x = constant_op.constant([2., 3., 4.]) y = two_x_plus_1(x) - self.assertAllEqual([5, 7, 9], y.numpy()) + self.assertAllEqual([5., 7., 9.], y.numpy()) def testNestedDefunWithVariable(self): with self.test_scope(): @@ -506,7 +509,7 @@ class EagerFunctionTest(xla_test.XLATestCase): x = constant_op.constant(3.0) y = f(x) - self.assertEqual(75, y.numpy()) + self.assertEqual(75.0, y.numpy()) def testNestedDefunInGradientTape(self): with self.test_scope(): @@ -555,6 +558,71 @@ class EagerFunctionTest(xla_test.XLATestCase): self.assertEqual(9, dy_v0.numpy()) self.assertEqual(15, dy_v1.numpy()) + def testWhileInDefun(self): + with self.test_scope(): + @def_function.function + def f(start): + c = lambda x: math_ops.less(x, 13.0) + b = lambda x: math_ops.add(x, 1.0) + return control_flow_ops.while_loop(c, b, [start]) + + y = f(constant_op.constant(3.0)) + self.assertEqual(13.0, y.numpy()) + + def testAutoGraphWhileInDefun(self): + with self.test_scope(): + @def_function.function + def f(start): + x = start + while x < 13.0: + x += 1.0 + return x + + y = f(constant_op.constant(3.0)) + self.assertEqual(13.0, y.numpy()) + + def testCondInDefun(self): + with self.test_scope(): + @def_function.function + def f(pred, value): + fn1 = lambda: math_ops.add(value, 1.0) + fn2 = lambda: math_ops.subtract(value, 1.0) + return control_flow_ops.cond(pred, fn1, fn2) + + plus_one = f(constant_op.constant(True), constant_op.constant(10.0)) + minus_one = f(constant_op.constant(False), constant_op.constant(10.0)) + self.assertEqual(11.0, plus_one.numpy()) + self.assertEqual(9.0, minus_one.numpy()) + + def testAutoGraphCondInDefun(self): + with self.test_scope(): + @def_function.function + def f(pred, value): + if pred: + return value + 1.0 + else: + return value - 1.0 + + plus_one = f(constant_op.constant(True), constant_op.constant(10.0)) + minus_one = f(constant_op.constant(False), constant_op.constant(10.0)) + self.assertEqual(11.0, plus_one.numpy()) + self.assertEqual(9.0, minus_one.numpy()) + + def testScanInDefun(self): + with self.test_scope(): + elems = constant_op.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name='data') + v = constant_op.constant(2.0, name='v') + + @def_function.function + def f(y): + # pylint: disable=unnecessary-lambda + return functional_ops.scan( + lambda a, x: math_ops.multiply(a, x), y, initializer=v) + # pylint: enable=unnecessary-lambda + + r = f(elems) + self.assertAllEqual([2., 4., 12., 48., 240., 1440.], self.evaluate(r)) + class ExcessivePaddingTest(xla_test.XLATestCase): """Test that eager execution works with TPU flattened tensors. diff --git a/tensorflow/compiler/tests/fused_batchnorm_test.py b/tensorflow/compiler/tests/fused_batchnorm_test.py index 374942a0b339b816944ea5529e4f84134b60017b..56a8e1b1667f154f6cec475ee0f4f8b308121c09 100644 --- a/tensorflow/compiler/tests/fused_batchnorm_test.py +++ b/tensorflow/compiler/tests/fused_batchnorm_test.py @@ -191,6 +191,20 @@ class FusedBatchNormTest(xla_test.XLATestCase, parameterized.TestCase): mean_val = np.random.random_sample(scale_shape).astype(np.float32) var_val = np.random.random_sample(scale_shape).astype(np.float32) epsilon = 0.001 + + # The TensorFlow FusedBatchNormGrad training operation takes two inputs with + # implementation defined values. In theory the only correct value these + # inputs are the corresponding reserve_space_{1|2} outputs from the + # FusedBatchNorm training operation. However, in practice, we rely on the + # first one being mean on {C|G}PU, and the second one being variance on CPU + # and inverse(sqrt(variance + epsilon)) on GPU (we test this assumption + # separately). + reserve_space_1_val = mean_val + if self.device == "XLA_GPU": + reserve_space_2_val = np.reciprocal(np.sqrt(var_val + epsilon)) + else: + reserve_space_2_val = var_val + data_format_src = "NHWC" grad_x_ref, grad_scale_ref, grad_offset_ref = self._reference_grad( x_val, grad_val, scale_val, mean_val, var_val, epsilon, data_format_src) @@ -207,18 +221,26 @@ class FusedBatchNormTest(xla_test.XLATestCase, parameterized.TestCase): np.float32, shape=x_val_converted.shape, name="grad") x = array_ops.placeholder( np.float32, shape=x_val_converted.shape, name="x") - mean = array_ops.placeholder(np.float32, shape=scale_shape, name="mean") - var = array_ops.placeholder(np.float32, shape=scale_shape, name="var") + reserve_space_1 = array_ops.placeholder( + np.float32, shape=scale_shape, name="reserve_space_1") + reserve_space_2 = array_ops.placeholder( + np.float32, shape=scale_shape, name="reserve_space_2") scale = array_ops.placeholder(np.float32, shape=scale_shape, name="scale") grad_x, grad_scale, grad_offset, _, _ = gen_nn_ops.fused_batch_norm_grad( - grad, x, scale, mean, var, data_format=data_format, is_training=True) + grad, + x, + scale, + reserve_space_1, + reserve_space_2, + data_format=data_format, + is_training=True) grad_x_val, grad_scale_val, grad_offset_val = sess.run( [grad_x, grad_scale, grad_offset], { grad: grad_val_converted, x: x_val_converted, - mean: mean_val, - var: var_val, + reserve_space_1: reserve_space_1_val, + reserve_space_2: reserve_space_2_val, scale: scale_val }) diff --git a/tensorflow/compiler/tests/image_ops_test.py b/tensorflow/compiler/tests/image_ops_test.py index 0e2d840418156d825e2d141018e49f42374c8fee..42e688174fce9e939feb09e1767ebab31e30a6ee 100644 --- a/tensorflow/compiler/tests/image_ops_test.py +++ b/tensorflow/compiler/tests/image_ops_test.py @@ -403,6 +403,117 @@ class AdjustSaturationTest(xla_test.XLATestCase): self.assertAllClose(y_fused, y_baseline, rtol=2e-5, atol=1e-5) +class ResizeNearestNeighborTest(xla_test.XLATestCase): + # TODO(ilch): Wrap each test with `for dtype in self.float_types:` + # Some work to understand how that should be done was presented here: + # cl/227850213 + + def _assertForwardOpMatchesExpected(self, + image_np, + target_shape, + expected=None, + large_tolerance=False, + align_corners=True): + if expected is None: + self.fail("expected must be specified") + with self.cached_session() as sess, self.test_scope(): + image = array_ops.placeholder(image_np.dtype) + resized = gen_image_ops.resize_nearest_neighbor( + image, target_shape, align_corners=align_corners) + out = sess.run(resized, {image: image_np[np.newaxis, :, :, np.newaxis]}) + if large_tolerance: + self.assertAllClose( + expected[np.newaxis, :, :, np.newaxis], out, rtol=2e-4, atol=2e-4) + else: + self.assertAllClose(expected[np.newaxis, :, :, np.newaxis], out) + + def testAlignCorners2x2To1x1(self): + self._assertForwardOpMatchesExpected( + np.array([[1, 2], [3, 4]], dtype=np.float32), [1, 1], + expected=np.array([[1]], dtype=np.float32)) + + def testAlignCorners1x1To2x2(self): + self._assertForwardOpMatchesExpected( + np.array([[1]], dtype=np.float32), [2, 2], + expected=np.array([[1, 1], [1, 1]], dtype=np.float32)) + + def testAlignCorners1x1To3x3(self): + self._assertForwardOpMatchesExpected( + np.array([[1]], dtype=np.float32), [3, 3], + expected=np.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]], dtype=np.float32)) + + def testAlignCorners2x2To3x3(self): + self._assertForwardOpMatchesExpected( + np.array([[1, 2], [3, 4]], dtype=np.float32), [3, 3], + expected=np.array([[1, 2, 2], [3, 4, 4], [3, 4, 4]], dtype=np.float32)) + + def testAlignCorners2x2To4x4(self): + self._assertForwardOpMatchesExpected( + np.array([[1, 2], [3, 4]], dtype=np.float32), [4, 4], + expected=np.array( + [[1, 1, 2, 2], [1, 1, 2, 2], [3, 3, 4, 4], [3, 3, 4, 4]], + dtype=np.float32), large_tolerance=True) + + def testAlignCorners3x3To2x2(self): + self._assertForwardOpMatchesExpected( + np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32), [2, 2], + expected=np.array([[1, 3], [7, 9]], dtype=np.float32)) + + def testAlignCorners4x4To3x3(self): + self._assertForwardOpMatchesExpected( + np.array( + [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]], + dtype=np.float32), [3, 3], + expected=np.array([[1, 3, 4], [9, 11, 12], [13, 15, 16]], + dtype=np.float32)) + + def testAlignCorners3x3To4x4(self): + self._assertForwardOpMatchesExpected( + np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32), [4, 4], + expected=np.array( + [[1, 2, 2, 3], [4, 5, 5, 6], [4, 5, 5, 6], [7, 8, 8, 9]], + dtype=np.float32)) + + def testAlignCorners3x3To6x6(self): + self._assertForwardOpMatchesExpected( + np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32), [6, 6], + expected=np.array( + [[1, 1, 2, 2, 3, 3], [1, 1, 2, 2, 3, 3], [4, 4, 5, 5, 6, 6], + [4, 4, 5, 5, 6, 6], [7, 7, 8, 8, 9, 9], [7, 7, 8, 8, 9, 9]], + dtype=np.float32)) + + def testAlignCorners3x3To9x9(self): + # The expected matrix might look uneven in terms of how many of each number + # there is, but this is an artifact of doing the dilation and convolution + # iteratively. The behavior is less esoteric in the 3x3To12x12 case below. + self._assertForwardOpMatchesExpected( + np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32), [9, 9], + expected=np.array( + [[1, 2, 2, 2, 2, 3, 3, 3, 3], [4, 5, 5, 5, 5, 6, 6, 6, 6], + [4, 5, 5, 5, 5, 6, 6, 6, 6], [4, 5, 5, 5, 5, 6, 6, 6, 6], + [4, 5, 5, 5, 5, 6, 6, 6, 6], [7, 8, 8, 8, 8, 9, 9, 9, 9], + [7, 8, 8, 8, 8, 9, 9, 9, 9], [7, 8, 8, 8, 8, 9, 9, 9, 9], + [7, 8, 8, 8, 8, 9, 9, 9, 9]], + dtype=np.float32)) + + def testAlignCorners3x3To12x12(self): + self._assertForwardOpMatchesExpected( + np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32), [12, 12], + expected=np.array([[1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3], + [1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3], + [1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3], + [4, 4, 4, 5, 5, 5, 5, 5, 5, 6, 6, 6], + [4, 4, 4, 5, 5, 5, 5, 5, 5, 6, 6, 6], + [4, 4, 4, 5, 5, 5, 5, 5, 5, 6, 6, 6], + [4, 4, 4, 5, 5, 5, 5, 5, 5, 6, 6, 6], + [4, 4, 4, 5, 5, 5, 5, 5, 5, 6, 6, 6], + [4, 4, 4, 5, 5, 5, 5, 5, 5, 6, 6, 6], + [7, 7, 7, 8, 8, 8, 8, 8, 8, 9, 9, 9], + [7, 7, 7, 8, 8, 8, 8, 8, 8, 9, 9, 9], + [7, 7, 7, 8, 8, 8, 8, 8, 8, 9, 9, 9]], + dtype=np.float32)) + + class ResizeBilinearTest(xla_test.XLATestCase): def _assertForwardOpMatchesExpected(self, @@ -444,14 +555,14 @@ class ResizeBilinearTest(xla_test.XLATestCase): self.assertAllCloseAccordingToType(expected[np.newaxis, :, :, np.newaxis], out) - def testAlignCorners1x2To3x2(self): + def testAlignCorners1x2To3x3(self): for dtype in self.float_types: self._assertForwardOpMatchesExpected( np.array([[1, 2]], dtype=dtype), [3, 3], expected=np.array([[1, 1.5, 2], [1, 1.5, 2], [1, 1.5, 2]], dtype=np.float32)) - def testAlignCorners1x2To3x2Grad(self): + def testAlignCorners1x2To3x3Grad(self): for dtype in self.float_types: self._assertBackwardOpMatchesExpected( np.array([[1, 2], [3, 4], [5, 6]], dtype=np.float32), diff --git a/tensorflow/compiler/tests/randomized_tests.cc b/tensorflow/compiler/tests/randomized_tests.cc index d23fd125163d1afe8c7fd5e008d4b617ff4b2874..1521cc760b85b176acb27c1489640e92ef90e247 100644 --- a/tensorflow/compiler/tests/randomized_tests.cc +++ b/tensorflow/compiler/tests/randomized_tests.cc @@ -63,6 +63,7 @@ limitations under the License. #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/graph/graph_constructor.h" #include "tensorflow/core/kernels/ops_util.h" +#include "tensorflow/core/lib/bfloat16/bfloat16.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/test.h" @@ -80,6 +81,7 @@ int64 tf_xla_random_seed = 0; int32 tf_xla_test_repetitions = 20; int64 tf_xla_max_tensor_size = 10000LL; string* tf_xla_test_device_ptr; // initial value set in main() +string* tf_xla_reference_device_ptr; // initial value set in main() bool tf_xla_test_use_jit = true; string LocalDeviceToFullDeviceName(const string& device) { @@ -321,6 +323,9 @@ class OpTest : public ::testing::Test { // for use as reduction indices. Tensor RandomReductionIndices(int rank); + // Returns a random bit. + bool RandomBool(); + struct WindowedSpatialDims { Padding padding; std::vector kernel_dims; @@ -453,6 +458,11 @@ std::vector OpTest::RandomDims(int min_rank, int max_rank, return dims; } +bool OpTest::RandomBool() { + std::bernoulli_distribution d(0.5); + return d(generator()); +} + Tensor OpTest::RandomTensor(DataType dtype, bool needs_unique_values, absl::Span shape) { Tensor tensor(dtype, TensorShape(shape)); @@ -760,8 +770,22 @@ Status TensorsAreEqualImpl(const Tensor& x, const Tensor& y) { for (int i = 0; i < Tx.size(); ++i) { if (Tx(i) != Ty(i)) { return errors::InvalidArgument(absl::StrCat( - i, "-th tensor element isn't equal: ", Tx(i), " vs. ", Ty(i), - ". x = ", x.DebugString(), "y = ", y.DebugString())); + i, "-th tensor element isn't equal: ", Str(Tx(i)), " vs. ", + Str(Ty(i)), ". x = ", x.DebugString(), "y = ", y.DebugString())); + } + } + return Status::OK(); +} + +Status TensorsAreEqualImplBfloat16(const Tensor& x, const Tensor& y) { + auto Tx = x.flat(); + auto Ty = y.flat(); + for (int i = 0; i < Tx.size(); ++i) { + if (Tx(i) != Ty(i)) { + return errors::InvalidArgument(absl::StrCat( + i, "-th tensor element isn't equal: ", static_cast(Tx(i)), + " vs. ", static_cast(Ty(i)), ". x = ", x.DebugString(), + "y = ", y.DebugString())); } } return Status::OK(); @@ -797,6 +821,8 @@ Status TensorsAreClose(const Tensor& a, const Tensor& b, double atol, return TensorsAreEqualImpl(a, b); case DT_BOOL: return TensorsAreEqualImpl(a, b); + case DT_BFLOAT16: + return TensorsAreEqualImplBfloat16(a, b); default: LOG(FATAL) << "Unexpected type : " << DataTypeString(a.dtype()); } @@ -829,8 +855,8 @@ OpTest::TestResult OpTest::ExpectTfAndXlaOutputsAreClose( VLOG(1) << "Input: " << input_tensors.back().DebugString(); } - string cpu_device = - LocalDeviceToFullDeviceName(absl::StrCat(DEVICE_CPU, ":0")); + string reference_device = + LocalDeviceToFullDeviceName(*tf_xla_reference_device_ptr); string test_device = LocalDeviceToFullDeviceName(*tf_xla_test_device_ptr); DeviceNameUtils::ParsedName parsed_name; @@ -845,9 +871,9 @@ OpTest::TestResult OpTest::ExpectTfAndXlaOutputsAreClose( std::vector expected_inputs, test_inputs; std::vector expected_fetches, test_fetches; Status status = builder.BuildGraph( - absl::StrCat("test", num_tests_, "_expected"), cpu_device, - /* use_jit= */ false, &graph, /* test_node_def= */ nullptr, - &expected_inputs, &expected_fetches); + absl::StrCat("test", num_tests_, "_expected"), reference_device, + /*use_jit=*/false, &graph, /*test_node_def=*/nullptr, &expected_inputs, + &expected_fetches); if (!status.ok()) { LOG(ERROR) << "Expected graph construction failed: " << status; return kFatalError; @@ -1371,6 +1397,19 @@ TEST_F(OpTest, Cast) { }); } +TEST_F(OpTest, CastBF16) { + Repeatedly([this]() { + DataType src_type, dst_type; + src_type = Choose({DT_FLOAT}); + dst_type = Choose({DT_BFLOAT16}); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Cast") + .RandomInput(src_type) + .Attr("SrcT", src_type) + .Attr("DstT", dst_type) + .Attr("Truncate", true)); + }); +} + TEST_F(OpTest, Ceil) { Repeatedly([this]() { return ExpectTfAndXlaOutputsAreClose( @@ -3346,11 +3385,41 @@ TEST_F(OpTest, ZerosLike) { }); } +// Example failing run: +// --tf_xla_reference_device=GPU:0 +// --tf_xla_test_use_jit=true --tf_xla_test_device=GPU:0 +// --tf_xla_test_repetitions=2 +// --gunit_filter='OpTest.FusedBatchNormTraining' +// --tf_xla_random_seed=2838146746 +TEST_F(OpTest, FusedBatchNormTraining) { + bool is_nhwc = RandomBool(); + std::vector x_dims = RandomDims(/*min_rank=*/4, /*max_rank=*/4, + /*min_size=*/5, /*max_size=*/20); + std::vector scale_dims = {x_dims[is_nhwc ? 3 : 1]}; + std::vector offset_dims = {x_dims[is_nhwc ? 3 : 1]}; + std::vector mean_dims = {0}; + std::vector variance_dims = {0}; + DataType type = DT_FLOAT; + Repeatedly([&] { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("FusedBatchNorm") + .RandomInput(type, x_dims) + .RandomInput(type, scale_dims) + .RandomInput(type, offset_dims) + .RandomInput(type, mean_dims) + .RandomInput(type, variance_dims) + .Attr("T", type) + .Attr("data_format", is_nhwc ? "NHWC" : "NCHW") + .Attr("epsilon", static_cast(1.001e-05)) + .Attr("is_training", true)); + }); +} } // anonymous namespace } // namespace tensorflow int main(int argc, char** argv) { tensorflow::tf_xla_test_device_ptr = new tensorflow::string("GPU:0"); + tensorflow::tf_xla_reference_device_ptr = new tensorflow::string("CPU:0"); std::vector flag_list = { tensorflow::Flag( "tf_xla_random_seed", &tensorflow::tf_xla_random_seed, @@ -3366,6 +3435,9 @@ int main(int argc, char** argv) { "Maximum number of elements for random input tensors."), tensorflow::Flag("tf_xla_test_device", tensorflow::tf_xla_test_device_ptr, "Tensorflow device type to use for test"), + tensorflow::Flag("tf_xla_reference_device", + tensorflow::tf_xla_reference_device_ptr, + "Tensorflow device type to use for reference"), tensorflow::Flag("tf_xla_test_use_jit", &tensorflow::tf_xla_test_use_jit, "Use JIT compilation for the operator under test"), }; diff --git a/tensorflow/compiler/tests/tensor_list_ops_test.py b/tensorflow/compiler/tests/tensor_list_ops_test.py index 5c079d595c440cac644f5461154509abe7b1d1ed..47e0f384a4f1e46ccc35584aaff3a0aceff8a985 100644 --- a/tensorflow/compiler/tests/tensor_list_ops_test.py +++ b/tensorflow/compiler/tests/tensor_list_ops_test.py @@ -23,24 +23,20 @@ from tensorflow.compiler.tests import xla_test from tensorflow.python.framework import constant_op 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.ops import list_ops from tensorflow.python.platform import test -def scalar_shape(): - return ops.convert_to_tensor([], dtype=dtypes.int32) - - class ListOpsTest(xla_test.XLATestCase): def testElementShape(self): with self.cached_session() as sess, self.test_scope(): dim = array_ops.placeholder(dtypes.int32) - l = list_ops.tensor_list_reserve( - element_shape=(dim, 15), num_elements=20, - element_dtype=dtypes.float32) + l = list_ops.empty_tensor_list( + element_shape=(dim, 15), + element_dtype=dtypes.float32, + max_num_elements=20) e32 = list_ops.tensor_list_element_shape(l, shape_type=dtypes.int32) e64 = list_ops.tensor_list_element_shape(l, shape_type=dtypes.int64) self.assertAllEqual(sess.run(e32, {dim: 10}), (10, 15)) @@ -48,25 +44,44 @@ class ListOpsTest(xla_test.XLATestCase): def testPushPop(self): with self.cached_session() as sess, self.test_scope(): - num = array_ops.placeholder(dtypes.int32) - l = list_ops.tensor_list_reserve( - element_shape=(7, 15), num_elements=num, element_dtype=dtypes.float32) + l = list_ops.empty_tensor_list( + element_shape=(7, 15), + element_dtype=dtypes.float32, + max_num_elements=10) l = list_ops.tensor_list_push_back( l, constant_op.constant(1.0, shape=(7, 15))) l = list_ops.tensor_list_push_back( l, constant_op.constant(2.0, shape=(7, 15))) l, e2 = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32) _, e1 = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32) - self.assertAllEqual(sess.run(e2, {num: 10}), 2.0 * np.ones((7, 15))) - self.assertAllEqual(sess.run(e1, {num: 10}), 1.0 * np.ones((7, 15))) + self.assertAllEqual(sess.run(e2), 2.0 * np.ones((7, 15))) + self.assertAllEqual(sess.run(e1), 1.0 * np.ones((7, 15))) + + def testDoNotConstantFoldVariants(self): + with self.cached_session() as sess, self.test_scope(): + val = array_ops.placeholder(dtype=dtypes.float32) + l = list_ops.empty_tensor_list( + element_shape=(7, 15), + element_dtype=dtypes.float32, + max_num_elements=10) + # Note: Pushing a Placeholder will force the constant folding code + # to build a Const node with a DT_VARIANT output. This tests that XLA + # passes a cf_consider_fn which prevent folding such nodes. + l = list_ops.tensor_list_push_back( + l, array_ops.fill(value=val, dims=(7, 15))) + l = list_ops.tensor_list_push_back( + l, constant_op.constant(2.0, shape=(7, 15))) + l, e2 = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32) + _, e1 = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32) + self.assertAllEqual(sess.run(e2, {val: 1.0}), 2.0 * np.ones((7, 15))) + self.assertAllEqual(sess.run(e1, {val: 1.0}), 1.0 * np.ones((7, 15))) def testPushPopSeparateLists(self): with self.cached_session() as sess, self.test_scope(): - num = array_ops.placeholder(dtypes.int32) - l = list_ops.tensor_list_reserve( - element_shape=scalar_shape(), - num_elements=num, - element_dtype=dtypes.float32) + l = list_ops.empty_tensor_list( + element_shape=[], + element_dtype=dtypes.float32, + max_num_elements=20) l = list_ops.tensor_list_push_back(l, constant_op.constant(1.0)) l2 = list_ops.tensor_list_push_back(l, constant_op.constant(2.0)) l3 = list_ops.tensor_list_push_back(l, constant_op.constant(3.0)) @@ -75,22 +90,95 @@ class ListOpsTest(xla_test.XLATestCase): l2, e22 = list_ops.tensor_list_pop_back(l2, element_dtype=dtypes.float32) l3, e31 = list_ops.tensor_list_pop_back(l3, element_dtype=dtypes.float32) l3, e32 = list_ops.tensor_list_pop_back(l3, element_dtype=dtypes.float32) - result = sess.run([e11, [e21, e22], [e31, e32]], {num: 20}) + result = sess.run([e11, [e21, e22], [e31, e32]]) self.assertEqual(result, [1.0, [2.0, 1.0], [3.0, 1.0]]) - def testEmptyTensorList(self): - dim = 7 + def testEmptyTensorListNoMax(self): with self.cached_session() as sess, self.test_scope(): - p = array_ops.placeholder(dtypes.int32) l = list_ops.empty_tensor_list( - element_shape=(p, 15), element_dtype=dtypes.float32) + element_shape=(7, 15), element_dtype=dtypes.float32) l = list_ops.tensor_list_push_back( - l, constant_op.constant(1.0, shape=(dim, 15))) + l, constant_op.constant(1.0, shape=(7, 15))) _, e = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32) with self.assertRaisesRegexp(errors.InvalidArgumentError, - "Use TensorListReserve instead"): - self.assertEqual(sess.run(e, {p: dim}), 1.0 * np.ones((dim, 15))) + "Set the max number of elements"): + self.assertEqual(sess.run(e), 1.0 * np.ones((7, 15))) + def testEmptyTensorListMax(self): + with self.cached_session() as sess, self.test_scope(): + l = list_ops.empty_tensor_list( + element_shape=(10, 15), element_dtype=dtypes.float32, + max_num_elements=2) + l = list_ops.tensor_list_push_back( + l, array_ops.fill(value=3.0, dims=(10, 15))) + _, e = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32) + self.assertAllEqual(sess.run(e), 3.0 * np.ones((10, 15))) + + def testListFromTensor(self): + with self.cached_session(), self.test_scope(): + t = constant_op.constant([1.0, 2.0]) + l = list_ops.tensor_list_from_tensor(t, element_shape=[]) + e = list_ops.tensor_list_get_item(l, 0, element_dtype=dtypes.float32) + self.assertAllEqual(e, 1.0) + l, e0 = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32) + self.assertAllEqual(e0, 2.0) + l, e1 = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32) + self.assertAllEqual(e1, 1.0) + self.assertAllEqual(list_ops.tensor_list_length(l), 0) + + def testGetSet(self): + with self.cached_session(), self.test_scope(): + t = constant_op.constant([1.0, 2.0]) + l = list_ops.tensor_list_from_tensor(t, element_shape=[]) + e0 = list_ops.tensor_list_get_item(l, 0, element_dtype=dtypes.float32) + self.assertAllEqual(e0, 1.0) + l = list_ops.tensor_list_set_item(l, 0, 3.0) + t = list_ops.tensor_list_stack(l, element_dtype=dtypes.float32) + self.assertAllEqual(t, [3.0, 2.0]) + + def testGetSetReserved(self): + with self.cached_session(), self.test_scope(): + l = list_ops.tensor_list_reserve( + element_dtype=dtypes.float32, element_shape=[], num_elements=2) + e0 = list_ops.tensor_list_get_item(l, 0, element_dtype=dtypes.float32) + self.assertAllEqual(e0, 0.0) + l = list_ops.tensor_list_set_item(l, 0, 3.0) + t = list_ops.tensor_list_stack(l, element_dtype=dtypes.float32) + self.assertAllEqual(t, [3.0, 0.0]) + + def testGetSetReservedNonScalar(self): + with self.cached_session() as sess, self.test_scope(): + l = list_ops.tensor_list_reserve( + element_dtype=dtypes.float32, + element_shape=(7, 15), + num_elements=2) + l = list_ops.tensor_list_set_item( + l, 0, constant_op.constant(1.0, shape=(7, 15))) + e1 = list_ops.tensor_list_get_item(l, 0, element_dtype=dtypes.float32) + e2 = list_ops.tensor_list_get_item(l, 1, element_dtype=dtypes.float32) + self.assertAllEqual(sess.run(e1), np.ones((7, 15))) + self.assertAllEqual(sess.run(e2), np.zeros((7, 15))) + + def testStack(self): + with self.cached_session(), self.test_scope(): + l = list_ops.empty_tensor_list( + element_dtype=dtypes.float32, + element_shape=[], + max_num_elements=2) + l = list_ops.tensor_list_push_back(l, constant_op.constant(1.0)) + e = list_ops.tensor_list_get_item(l, 0, element_dtype=dtypes.float32) + self.assertAllEqual(e, 1.0) + l = list_ops.tensor_list_push_back(l, constant_op.constant(2.0)) + t = list_ops.tensor_list_stack(l, element_dtype=dtypes.float32) + self.assertAllEqual(t.shape.as_list(), [None]) + self.assertAllEqual(t, [1.0, 2.0]) + + def testStackWithUninitializedTensors(self): + with self.cached_session(), self.test_scope(): + l = list_ops.tensor_list_reserve( + element_dtype=dtypes.float32, element_shape=[], num_elements=3) + t = list_ops.tensor_list_stack(l, element_dtype=dtypes.float32) + self.assertAllEqual(t, [0., 0., 0.]) if __name__ == "__main__": test.main() diff --git a/tensorflow/compiler/tests/unary_ops_test.py b/tensorflow/compiler/tests/unary_ops_test.py index 95c9e7ffd4651642781143c2c1940b0e51e1e470..831e203f49d4aacefe967bc0df3b2a6ada3b67ac 100644 --- a/tensorflow/compiler/tests/unary_ops_test.py +++ b/tensorflow/compiler/tests/unary_ops_test.py @@ -391,6 +391,11 @@ class UnaryOpsTest(xla_test.XLATestCase): expected=np.array( [[-0.66666669, -0.5, 0, 0.5, 0.66666669]], dtype=dtype)) + self._assertOpOutputMatchesExpected( + math_ops.sign, + np.array([[-2.0, -1.0, -0.0, +0.0, 1.0, 2.0]], dtype=dtype), + expected=np.array([[-1.0, -1.0, -0.0, +0.0, 1.0, 1.0]], dtype=dtype)) + self._assertOpOutputMatchesExpected( math_ops.is_finite, np.array( @@ -647,7 +652,7 @@ class UnaryOpsTest(xla_test.XLATestCase): np.array([1, 2j, 2 - 3j, 4 + 5j], dtype=dtype), expected=np.tan(np.array([1, 2j, 2 - 3j, 4 + 5j], dtype=dtype))) - ctypes = {np.complex64: np.float32} + ctypes = {np.complex64: np.float32, np.complex128: np.float64} self._assertOpOutputMatchesExpected( math_ops.abs, np.array([[3 - 4j, -1j, np.inf]], dtype=dtype), @@ -743,6 +748,10 @@ class UnaryOpsTest(xla_test.XLATestCase): np.array( [[np.NINF, -2, -1, 0, 0.5, 1, 2, np.inf, np.nan]], dtype=dtype), expected=np.array([[0, 0, 0, 0, 0, 0, 0, 0, 1]], dtype=np.bool)) + self._assertOpOutputMatchesExpected( + math_ops.sign, + np.array([[np.nan]], dtype=dtype), + expected=np.array([[0.0]], dtype=dtype)) def testLogicalOps(self): self._assertOpOutputMatchesExpected( @@ -760,7 +769,7 @@ class UnaryOpsTest(xla_test.XLATestCase): lambda x: gen_nn_ops.bias_add_grad(x, data_format="NCHW"), np.array( [[[1., 2.], [3., 4.]], [[5., 6.], [7., 8.]]], dtype=np.float32), - expected=np.array([10., 26.], dtype=np.float32)) + expected=np.array([14., 22.], dtype=np.float32)) def testCast(self): shapes = [[], [4], [2, 3], [2, 0, 4]] @@ -811,6 +820,12 @@ class UnaryOpsTest(xla_test.XLATestCase): np.array([1, 2, 0], np.int32), expected=np.array([2, 0, 1], dtype=np.int32)) + def testInvertPermutationTwiceIsNoop(self): + self._assertOpOutputMatchesExpected( + lambda x: array_ops.invert_permutation(array_ops.invert_permutation(x)), + np.array([1, 2, 0], np.int32), + expected=np.array([1, 2, 0], dtype=np.int32)) + def testRank(self): rank_op = lambda x: array_ops.rank_internal(x, optimize=False) for dtype in self.numeric_types: diff --git a/tensorflow/compiler/tests/variable_ops_test.py b/tensorflow/compiler/tests/variable_ops_test.py index fcd7ac5ba1ca5049246e93e6f5f76746fb28c6b8..18c5870e0decb686f4df1c16bbb4a340c93ad21d 100644 --- a/tensorflow/compiler/tests/variable_ops_test.py +++ b/tensorflow/compiler/tests/variable_ops_test.py @@ -485,7 +485,7 @@ class SliceAssignTest(xla_test.XLATestCase): checker2[None] = [6] # new axis def testUninitialized(self): - with self.assertRaisesRegexp(errors.InvalidArgumentError, + with self.assertRaisesRegexp(errors.FailedPreconditionError, "uninitialized variable"): with self.test_session() as sess, self.test_scope(): v = resource_variable_ops.ResourceVariable([1, 2]) diff --git a/tensorflow/compiler/tf2tensorrt/BUILD b/tensorflow/compiler/tf2tensorrt/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..eca533095ac431aca15e2c3390f6cec4fecb9978 --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/BUILD @@ -0,0 +1,447 @@ +# Description: +# Wrap NVIDIA TensorRT (http://developer.nvidia.com/tensorrt) with tensorflow +# and provide TensorRT operators and converter package. +# APIs are meant to change over time. + +package(default_visibility = ["//visibility:public"]) + +licenses(["notice"]) # Apache 2.0 + +exports_files(["LICENSE"]) + +load( + "//tensorflow:tensorflow.bzl", + "tf_cc_test", + "tf_copts", + "tf_cuda_library", + "tf_custom_op_library", + "tf_custom_op_library_additional_deps", + "tf_gen_op_libs", + "tf_gen_op_wrapper_py", +) +load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test") +load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") +load( + "@local_config_tensorrt//:build_defs.bzl", + "if_tensorrt", +) + +tf_cuda_cc_test( + name = "tensorrt_test_cc", + size = "small", + srcs = ["tensorrt_test.cc"], + tags = [ + "no_windows", + "nomac", + ], + deps = [ + "//tensorflow/core:gpu_init", + "//tensorflow/core:lib", + "//tensorflow/core:stream_executor", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ] + if_tensorrt([ + "@local_config_cuda//cuda:cuda_headers", + "@local_config_tensorrt//:tensorrt", + ]), +) + +tf_custom_op_library( + name = "python/ops/_trt_ops.so", + srcs = [ + "ops/get_serialized_resource_op.cc", + "ops/trt_engine_op.cc", + ], + deps = [ + "//tensorflow/core:lib_proto_parsing", + ] + if_tensorrt([ + "@local_config_tensorrt//:tensorrt", + ]), +) + +cc_library( + name = "trt_op_kernels", + srcs = [ + "kernels/get_serialized_resource_op.cc", + "kernels/trt_engine_op.cc", + ], + copts = tf_copts(), + visibility = ["//visibility:public"], + deps = [ + ":test_utils", + ":trt_allocator", + ":trt_conversion", + ":trt_logging", + ":trt_plugins", + ":trt_resources", + ":utils", + "@com_google_absl//absl/memory", + "@com_google_absl//absl/strings", + "//tensorflow/core:gpu_headers_lib", + "//tensorflow/core:lib_proto_parsing", + "//tensorflow/core:stream_executor_headers_lib", + "//tensorflow/core/grappler/costs:graph_properties", + ] + if_tensorrt([ + "@local_config_tensorrt//:tensorrt", + ]) + tf_custom_op_library_additional_deps(), + alwayslink = 1, +) + +tf_cuda_cc_test( + name = "get_serialized_resource_op_test", + size = "small", + srcs = ["kernels/get_serialized_resource_op_test.cc"], + tags = [ + "no_cuda_on_cpu_tap", + "no_windows", + "nomac", + ], + deps = [ + # TODO(laigd): consider splitting get_serialized_resource_op out from + # TF-TRT. + ":trt_op_kernels", + ":trt_op_libs", + ":trt_resources", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + "//tensorflow/core/kernels:ops_testutil", + ], +) + +tf_gen_op_libs( + op_lib_names = [ + "trt_engine_op", + "get_serialized_resource_op", + ], +) + +cc_library( + name = "trt_op_libs", + deps = [ + ":get_serialized_resource_op_op_lib", + ":trt_engine_op_op_lib", + ], +) + +tf_cuda_library( + name = "trt_logging", + srcs = ["utils/trt_logger.cc"], + hdrs = ["utils/trt_logger.h"], + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/core:lib_proto_parsing", + ] + if_tensorrt([ + "@local_config_tensorrt//:tensorrt", + ]), +) + +tf_gen_op_wrapper_py( + name = "trt_ops", + deps = [ + ":trt_op_libs", + ], +) + +tf_custom_op_py_library( + name = "trt_ops_loader", + srcs = ["python/ops/trt_ops.py"], + dso = [ + "python/ops/_trt_ops.so", + ] + if_tensorrt([ + "@local_config_tensorrt//:tensorrt", + ]), + kernels = [ + ":trt_op_kernels", + ":trt_op_libs", + ], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:errors", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform", + "//tensorflow/python:resources", + ], +) + +tf_cuda_library( + name = "trt_resources", + srcs = [ + "utils/trt_int8_calibrator.cc", + "utils/trt_resource_manager.cc", + "utils/trt_resources.cc", + ], + hdrs = [ + "utils/trt_int8_calibrator.h", + "utils/trt_lru_cache.h", + "utils/trt_resource_manager.h", + "utils/trt_resources.h", + ], + deps = [ + ":trt_allocator", + ":trt_logging", + ":utils", + "//tensorflow/core:framework_headers_lib", + "//tensorflow/core:framework_lite", + "//tensorflow/core:lib_proto_parsing", + ] + if_tensorrt([ + "@local_config_tensorrt//:tensorrt", + ]), +) + +tf_cuda_library( + name = "trt_allocator", + srcs = ["utils/trt_allocator.cc"], + hdrs = ["utils/trt_allocator.h"], + deps = [ + "//tensorflow/core:framework_headers_lib", + "//tensorflow/core:framework_lite", + "//tensorflow/core:lib_proto_parsing", + ] + if_tensorrt([ + "@local_config_tensorrt//:tensorrt", + ]), +) + +tf_cc_test( + name = "trt_allocator_test", + size = "small", + srcs = ["utils/trt_allocator_test.cc"], + tags = [ + "no_windows", + "nomac", + ], + deps = [ + ":trt_allocator", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + +tf_cc_test( + name = "trt_lru_cache_test", + size = "small", + srcs = ["utils/trt_lru_cache_test.cc"], + tags = [ + "no_windows", + "nomac", + ], + deps = [ + ":trt_resources", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + +# Library for the node-level conversion portion of TensorRT operation creation +tf_cuda_library( + name = "trt_conversion", + srcs = [ + "convert/convert_graph.cc", + "convert/convert_nodes.cc", + "convert/trt_optimization_pass.cc", + ], + hdrs = [ + "convert/convert_graph.h", + "convert/convert_nodes.h", + "convert/trt_optimization_pass.h", + ], + deps = [ + ":segment", + ":test_utils", + ":trt_allocator", + ":trt_plugins", + ":trt_logging", + ":trt_resources", + ":utils", + "@com_google_absl//absl/strings", + "//tensorflow/core/grappler/clusters:cluster", + "//tensorflow/core/grappler/optimizers:custom_graph_optimizer", + "//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry", + "//tensorflow/core/grappler:grappler_item", + "//tensorflow/core/grappler:utils", + "//tensorflow/core:framework", + "//tensorflow/core:framework_lite", + "//tensorflow/core:gpu_runtime", + "//tensorflow/core:graph", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core/grappler:devices", + "//tensorflow/core/grappler/clusters:virtual_cluster", + "//tensorflow/core/grappler/costs:graph_properties", + "//tensorflow/core/grappler/optimizers:meta_optimizer", + ] + if_tensorrt([ + "@local_config_tensorrt//:tensorrt", + ]) + tf_custom_op_library_additional_deps(), +) + +tf_cuda_cc_test( + name = "convert_graph_test", + size = "medium", + srcs = ["convert/convert_graph_test.cc"], + tags = [ + "no_cuda_on_cpu_tap", + "no_windows", + "nomac", + ], + deps = [ + ":trt_conversion", + "@com_google_googletest//:gtest", + "@com_google_absl//absl/strings", + "//tensorflow/cc:cc_ops", + "//tensorflow/cc:scope", + "//tensorflow/core/grappler:grappler_item", + "//tensorflow/core/grappler/clusters:cluster", + "//tensorflow/core:core_cpu", + "//tensorflow/core:core_cpu_base", + "//tensorflow/core:direct_session", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + ] + if_tensorrt([ + "@local_config_tensorrt//:tensorrt", + ]), +) + +tf_cuda_cc_test( + name = "convert_nodes_test", + size = "medium", + srcs = ["convert/convert_nodes_test.cc"], + tags = [ + "no_cuda_on_cpu_tap", + "no_windows", + "nomac", + ], + deps = [ + ":trt_logging", + ":trt_conversion", + ":trt_plugins", + "@com_google_googletest//:gtest", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", + "//tensorflow/cc:cc_ops", + "//tensorflow/cc:ops", + "//tensorflow/cc:scope", + "//tensorflow/core/grappler/costs:graph_properties", + "//tensorflow/core:core_cpu", + "//tensorflow/core:core_cpu_base", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:tensor_testutil", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + ] + if_tensorrt([ + "@local_config_cuda//cuda:cuda_headers", + "@local_config_tensorrt//:tensorrt", + ]), +) + +# Library for the segmenting portion of TensorRT operation creation +cc_library( + name = "segment", + srcs = ["segment/segment.cc"], + hdrs = [ + "segment/segment.h", + "segment/union_find.h", + ], + copts = tf_copts(), + deps = [ + "//tensorflow/core:graph", + "//tensorflow/core:lib_proto_parsing", + "//tensorflow/core:protos_all_cc", + "@com_google_absl//absl/strings", + "@protobuf_archive//:protobuf_headers", + ], +) + +tf_cc_test( + name = "segment_test", + size = "small", + srcs = ["segment/segment_test.cc"], + tags = [ + "no_windows", + "nomac", + ], + deps = [ + ":segment", + "//tensorflow/cc:cc_ops", + "//tensorflow/cc:scope", + "//tensorflow/core:core_cpu", + "//tensorflow/core:lib", + "//tensorflow/core:ops", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + ], +) + +# Library for the plugin factory +tf_cuda_library( + name = "trt_plugins", + srcs = [ + "plugin/trt_plugin.cc", + "plugin/trt_plugin_factory.cc", + "plugin/trt_plugin_utils.cc", + ], + hdrs = [ + "plugin/trt_plugin.h", + "plugin/trt_plugin_factory.h", + "plugin/trt_plugin_utils.h", + ], + deps = [ + "//tensorflow/core:framework_lite", + "//tensorflow/core:lib_proto_parsing", + ] + if_tensorrt([ + "@local_config_tensorrt//:tensorrt", + ]), +) + +tf_cuda_cc_test( + name = "trt_plugin_factory_test", + size = "small", + srcs = ["plugin/trt_plugin_factory_test.cc"], + tags = [ + "no_cuda_on_cpu_tap", + "no_windows", + "nomac", + ], + deps = [ + ":trt_plugins", + "//tensorflow/core:lib", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ] + if_tensorrt([ + "@local_config_cuda//cuda:cuda_headers", + "@local_config_tensorrt//:tensorrt", + ]), +) + +cc_library( + name = "utils", + srcs = ["convert/utils.cc"], + hdrs = ["convert/utils.h"], + copts = tf_copts(), + deps = [ + "//tensorflow/core:framework", + "//tensorflow/core:lib", + ], +) + +cc_library( + name = "test_utils", + srcs = ["utils/test_utils.cc"], + hdrs = ["utils/test_utils.h"], + deps = [ + "//tensorflow/core:lib", + "@com_googlesource_code_re2//:re2", + ], +) diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc similarity index 92% rename from tensorflow/contrib/tensorrt/convert/convert_graph.cc rename to tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc index bf2de94e04ae3f6817f7a679ce9fd88e750827dd..d6080c02d435fc149f679ebe9c9bacc8d0a0c144 100644 --- a/tensorflow/contrib/tensorrt/convert/convert_graph.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/convert/convert_graph.h" +#include "tensorflow/compiler/tf2tensorrt/convert/convert_graph.h" #include #include @@ -24,13 +24,14 @@ limitations under the License. #include #include -#include "tensorflow/contrib/tensorrt/convert/convert_nodes.h" -#include "tensorflow/contrib/tensorrt/convert/utils.h" -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.h" -#include "tensorflow/contrib/tensorrt/resources/trt_resource_manager.h" -#include "tensorflow/contrib/tensorrt/resources/trt_resources.h" -#include "tensorflow/contrib/tensorrt/segment/segment.h" -#include "tensorflow/contrib/tensorrt/test/utils.h" +#include "absl/strings/str_cat.h" +#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h" +#include "tensorflow/compiler/tf2tensorrt/convert/utils.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h" +#include "tensorflow/compiler/tf2tensorrt/segment/segment.h" +#include "tensorflow/compiler/tf2tensorrt/utils/test_utils.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_resource_manager.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_resources.h" #include "tensorflow/core/common_runtime/gpu/gpu_id.h" #include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h" #include "tensorflow/core/common_runtime/gpu/gpu_process_state.h" @@ -63,8 +64,8 @@ limitations under the License. namespace tensorflow { namespace tensorrt { namespace convert { -using ::tensorflow::strings::StrAppend; -using ::tensorflow::strings::StrCat; +using absl::StrAppend; +using absl::StrCat; // Returns compiled TRT version information {Maj, Min, Patch} std::vector GetLinkedTensorRTVersion() { @@ -82,7 +83,8 @@ std::vector GetLoadedTensorRTVersion() { } TrtCandidateSelector::TrtCandidateSelector( - const grappler::GraphProperties& graph_properties, int precision_mode) + const grappler::GraphProperties& graph_properties, + TrtPrecisionMode precision_mode) : graph_properties_(graph_properties), precision_mode_(precision_mode) {} Status TrtCandidateSelector::IsTensorRTCandidate(const tensorflow::Node* node) { @@ -97,6 +99,7 @@ Status TrtCandidateSelector::IsTensorRTCandidate(const tensorflow::Node* node) { "ConcatV2", "Const", "Conv2D", + "Conv2DBackpropInput", "DepthwiseConv2dNative", "Div", "Exp", @@ -104,6 +107,7 @@ Status TrtCandidateSelector::IsTensorRTCandidate(const tensorflow::Node* node) { "FusedBatchNorm", "FusedBatchNormV2", "Identity", + "LeakyRelu", "Log", "MatMul", "Max", @@ -148,10 +152,11 @@ Status TrtCandidateSelector::IsTensorRTCandidate(const tensorflow::Node* node) { // 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_ == INT8MODE && quantize_ops.count(node->type_string())) { + if (precision_mode_ == TrtPrecisionMode::INT8 && + quantize_ops.count(node->type_string())) { is_supported_op_type = true; } - // LINT.ThenChange(//tensorflow/contrib/tensorrt/convert/convert_nodes.cc) + // 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"); @@ -190,7 +195,7 @@ tensorflow::Status ConvertCalibGraphToInferGraph( const tensorflow::GraphDef& graph_def, tensorflow::GraphDef* infer_graph, bool is_dyn_op) { LOG(INFO) << "Starting Calib Conversion"; - infer_graph->CopyFrom(graph_def); + *infer_graph = graph_def; auto trt_rm = TRTResourceManager::instance(); auto calib_rm = trt_rm->getManager("TRTCalibration"); int num_nodes = infer_graph->node_size(); @@ -211,12 +216,13 @@ tensorflow::Status ConvertCalibGraphToInferGraph( return tensorflow::errors::FailedPrecondition( "Need to run graph with calibration data first!"); } + tensorflow::core::ScopedUnref calib_sc(cres); if (cres->calibrator_) { cres->calibrator_->waitAndSetDone(); cres->thr_->join(); const auto& calibration_table = cres->calibrator_->getCalibrationTableAsString(); - if (!calibration_table.size()) { + if (calibration_table.empty()) { LOG(ERROR) << "Calibration table is empty"; return tensorflow::errors::Unknown( "Calibration table is missing. This shouldn't have happened!"); @@ -227,7 +233,6 @@ tensorflow::Status ConvertCalibGraphToInferGraph( return tensorflow::errors::Unknown( "Can't get TRTCalibrator from resource manager!"); } - cres->Unref(); TF_RETURN_IF_ERROR(calib_rm->Cleanup(container_name)); } } @@ -238,7 +243,7 @@ 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, - int precision_mode, int minimum_segment_size, bool is_dyn_op, + 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. @@ -298,7 +303,7 @@ tensorflow::Status ConvertGraphDefToTensorRT( parameters["max_batch_size"].set_i(max_batch_size); parameters["is_dynamic_op"].set_b(is_dyn_op); parameters["max_workspace_size_bytes"].set_i(max_workspace_size_bytes); - TF_RETURN_IF_ERROR(GetPrecisionModeName( + TF_RETURN_IF_ERROR(TrtPrecisionModeToName( precision_mode, parameters["precision_mode"].mutable_s())); parameters["maximum_cached_engines"].set_i(max_cached_engines); if (!cached_engine_batches.empty()) { @@ -334,13 +339,12 @@ struct EdgePtrCompare { tensorflow::Status GetEngineInfo( const tensorflow::Graph* g, const tensorflow::grappler::GraphProperties& graph_properties, - const std::set& segment_nodes, + const std::set& segment_nodes, const std::unordered_map& node_map, const std::vector& reverse_topo_order, EngineInfo* info) { - std::vector subgraph_node_ids; // Topologically sorted node ids. - std::set subgraph_node_names = segment_nodes; - std::set added_const_node_ids; // Used to prevent double insertion. + std::vector subgraph_nodes; // Topologically sorted nodes. + std::set added_const_nodes; // Used to prevent double insertion. std::set segment_devices; // Map from src_node_name+port to the unique port numbers of the TRT op, where @@ -352,22 +356,37 @@ tensorflow::Status GetEngineInfo( std::unordered_map input_to_engine_port, output_to_engine_port; for (auto it = reverse_topo_order.rbegin(); it != reverse_topo_order.rend(); ++it) { - const auto& node_name = (*it)->name(); - if (segment_nodes.count(node_name) == 0) continue; - auto node = *it; + const Node* node = *it; + if (segment_nodes.count(node) == 0) continue; auto node_device = node->requested_device(); if (!node_device.empty()) { - segment_devices.insert(node_device); + // If device is CPU, treat as if no device was assigned. Don't add CPU to + // segment_device because that would cause a segfault in + // GetDeviceAndAllocator. This is because GetDeviceAndAllocator assumes + // any already set device is a GPU. + DeviceNameUtils::ParsedName parsed_name; + DeviceNameUtils::ParseFullName(node_device, &parsed_name); + if (parsed_name.type == "CPU") { + VLOG(1) << "Node " << node->name() << " was assigned to the CPU. " + << "Attempting to place on GPU."; + } else { + segment_devices.insert(node_device); + } } else { if (node->has_assigned_device_name()) { + // It appears that nodes will not have assigned devices at this point in + // execution. segment_devices.insert(node->assigned_device_name()); } else { VLOG(2) << "Node " << node->name() << " neither have requested device nor assigned device"; } } + subgraph_nodes.push_back(node); + const int node_id = node->id(); - subgraph_node_ids.push_back(node_id); + const string& node_name = node->name(); + // 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(), @@ -375,7 +394,7 @@ tensorflow::Status GetEngineInfo( std::sort(in_edges.begin(), in_edges.end(), EdgePtrCompare()); for (const auto edge : in_edges) { auto input_node = edge->src(); - if (input_node->IsSource() || segment_nodes.count(input_node->name())) { + if (input_node->IsSource() || segment_nodes.count(input_node)) { continue; } if (edge->IsControlEdge()) { @@ -392,12 +411,11 @@ tensorflow::Status GetEngineInfo( // // Note that the segmenter already ensure that the constant data input // is valid and suppported by the engine. - if (!added_const_node_ids.insert(input_node->id()).second) { + if (!added_const_nodes.insert(input_node).second) { // Already added before. continue; } VLOG(1) << "Adding const node " << input_node->name(); - QCHECK(subgraph_node_names.insert(input_node->name()).second); // Since we already add (duplicate) the const input node to the segment // graphdef, it's now not a data dependency any more, but to make the // dependency correct we still add a control dependency. @@ -428,7 +446,7 @@ tensorflow::Status GetEngineInfo( std::sort(out_edges.begin(), out_edges.end(), EdgePtrCompare()); for (const auto edge : out_edges) { auto output_node = edge->dst(); - if (output_node->IsSink() || segment_nodes.count(output_node->name())) { + if (output_node->IsSink() || segment_nodes.count(output_node)) { continue; } if (edge->IsControlEdge()) { @@ -456,12 +474,11 @@ tensorflow::Status GetEngineInfo( } // For each segment node in topological order. // Construct the const nodes first. - subgraph_node_ids.insert(subgraph_node_ids.begin(), - added_const_node_ids.begin(), - added_const_node_ids.end()); + subgraph_nodes.insert(subgraph_nodes.begin(), added_const_nodes.begin(), + added_const_nodes.end()); TF_RETURN_IF_ERROR(ConvertSegmentToGraphDef( - g, graph_properties, subgraph_node_names, subgraph_node_ids, - &info->connections, &info->segment_graph_def, &info->engine_name)); + g, graph_properties, subgraph_nodes, &info->connections, + &info->segment_graph_def, &info->engine_name)); // TODO(sami): This should not happen once segmenter is updated. if (segment_devices.size() == 1) { info->device = *segment_devices.begin(); @@ -625,7 +642,7 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, } const bool calibrate_int8 = - (info.precision_mode == INT8MODE && info.use_calibration); + (info.precision_mode == TrtPrecisionMode::INT8 && info.use_calibration); // Build the engine and get its serialized representation. string segment_string; if (info.engine_type == EngineInfo::EngineType::TRTStatic || calibrate_int8) { @@ -638,14 +655,15 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, TrtUniquePtrType engine; // TODO(sami): What happens if 1st dim is not batch? TF_RETURN_IF_ERROR(ConvertGraphDefToEngine( - info.segment_graph_def, calibrate_int8 ? FP32MODE : info.precision_mode, + info.segment_graph_def, + calibrate_int8 ? TrtPrecisionMode::FP32 : info.precision_mode, max_batch_size, info.max_workspace_size_bytes, input_shapes, &trt_logger, alloc, /*calibrator=*/nullptr, &engine, info.use_calibration, /*convert_successfully=*/nullptr)); TrtUniquePtrType engine_data(engine->serialize()); - segment_string = - string((const char*)engine_data->data(), engine_data->size()); + segment_string = string(static_cast(engine_data->data()), + engine_data->size()); if (calibrate_int8) { // See above comment about why not putting this inside the 'else' branch. segment_string = info.segment_graph_def.SerializeAsString(); @@ -654,14 +672,8 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, segment_string = info.segment_graph_def.SerializeAsString(); } - // TODO(aaroey): use enum instead, and add a helper method to do the - // conversion. string prec_string; - TF_RETURN_IF_ERROR(GetPrecisionModeName(info.precision_mode, &prec_string)); - if (info.precision_mode == INT8MODE && calibrate_int8 && - !TRTResourceManager::instance()->getManager("TRTCalibration")) { - LOG(ERROR) << "Failed to construct calibration storage"; - } + TF_RETURN_IF_ERROR(TrtPrecisionModeToName(info.precision_mode, &prec_string)); tensorflow::NodeDefBuilder node_builder(info.engine_name, "TRTEngineOp"); if (!info.device.empty()) node_builder.Device(info.device); if (VLOG_IS_ON(1)) { @@ -677,7 +689,7 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, } if (info.engine_type == EngineInfo::EngineType::TRTStatic && - info.cached_engine_batches.size()) { + !info.cached_engine_batches.empty()) { LOG(WARNING) << "Cached engine batches are ignored for static engines"; } tensorflow::NodeDef trt_node; @@ -691,7 +703,6 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, .Attr("serialized_segment", segment_string) .Attr("calibration_data", "") .Attr("max_cached_engines_count", info.maximum_cached_engines) - .Attr("cached_engine_batches", {max_batch_size}) .Attr("workspace_size_bytes", info.max_workspace_size_bytes) .Attr("precision_mode", prec_string) .Attr("use_calibration", info.use_calibration) @@ -843,6 +854,12 @@ tensorflow::Status RegisterSegmentFunctionToFunctionLibrary( auto native_segment = fdeflib.add_function(); TF_RETURN_IF_ERROR(tensorflow::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 + // would be on host if the op generating the tensor has host memory tag set. + (*native_segment + ->mutable_attr())[FunctionLibraryDefinition::kIntsOnDeviceAttr] + .set_b(true); if (VLOG_IS_ON(7)) { VLOG(7) << engine_name << " Function_Def "; VLOG(7) << native_segment->DebugString(); @@ -964,7 +981,8 @@ tensorflow::Status ConvertAfterShapes(ConversionParams& params) { continue; } curr_engine.precision_mode = params.precision_mode; - if (params.use_calibration && params.precision_mode != INT8MODE) { + if (params.use_calibration && + params.precision_mode != TrtPrecisionMode::INT8) { return errors::InvalidArgument( "Calibration with FP32 or FP16 is not supported."); } @@ -1033,27 +1051,31 @@ tensorflow::Status ConvertAfterShapes(ConversionParams& params) { cudaSetDevice(cuda_device_id); auto status = CreateTRTNode(engine_segments, i, params.max_batch_size, &graph, alloc.get(), &engine_nodes); - // If status is ok, we successfully added the node to the graph and can - // remove segment ops. Otherwise graph is not modified. + string msg = StrCat("TensorRT node ", engine.engine_name, " added for segment ", i, " consisting of ", converted_segments.at(i).first.size(), " nodes"); if (status.ok()) { LOG(INFO) << msg << " succeeded."; - for (auto node_name : converted_segments.at(i).first) { - graph.RemoveNode(node_map.at(node_name)); - } } else { // Graph is not modified. LOG(WARNING) << msg << " failed: " << status << ". Fallback to TF..."; } if (VLOG_IS_ON(1)) { msg = "Segment consists of nodes: "; - for (const string& node_name : converted_segments.at(i).first) { - StrAppend(&msg, node_name, ", "); + for (const Node* node : converted_segments.at(i).first) { + StrAppend(&msg, node->name(), ", "); } VLOG(1) << msg; } + + // If status is ok, we successfully added the node to the graph and can + // remove segment ops. Otherwise graph is not modified. + if (status.ok()) { + for (const Node* node : converted_segments.at(i).first) { + graph.RemoveNode(const_cast(node)); + } + } } cudaSetDevice(old_cuda_device); graph.ToGraphDef(params.output_graph_def); diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.h b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h similarity index 86% rename from tensorflow/contrib/tensorrt/convert/convert_graph.h rename to tensorflow/compiler/tf2tensorrt/convert/convert_graph.h index 1f39f56f6392ba33af3d74fec12c326ed4451cb6..95cf0227dcf84396b9de52194ae3a750f4acca66 100644 --- a/tensorflow/contrib/tensorrt/convert/convert_graph.h +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h @@ -12,12 +12,12 @@ 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_CONTRIB_TENSORRT_CONVERT_CONVERT_GRAPH_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_CONVERT_CONVERT_GRAPH_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_GRAPH_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_GRAPH_H_ #include -#include "tensorflow/contrib/tensorrt/convert/convert_nodes.h" +#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/grappler/clusters/cluster.h" #include "tensorflow/core/grappler/costs/graph_properties.h" @@ -36,7 +36,7 @@ namespace convert { class TrtCandidateSelector { public: TrtCandidateSelector(const grappler::GraphProperties& graph_properties, - int precision_mode); + TrtPrecisionMode precision_mode); // Returns OK iff 'node' is a TF-TRT conversion candidate, which will be added // to TRT subgraph and later converted into TRT engine. @@ -52,7 +52,7 @@ class TrtCandidateSelector { const grappler::GraphProperties& graph_properties_; // Quantization ops are only converted when using quantized precisions. - const int precision_mode_; + const TrtPrecisionMode precision_mode_; }; struct ConversionParams { @@ -61,7 +61,7 @@ struct ConversionParams { max_batch_size(1), max_workspace_size_bytes(1 << 30), output_graph_def(nullptr), - precision_mode(1), + precision_mode(TrtPrecisionMode::FP32), minimum_segment_size(3), graph_properties(nullptr), cluster(nullptr), @@ -74,7 +74,7 @@ struct ConversionParams { size_t max_batch_size; size_t max_workspace_size_bytes; tensorflow::GraphDef* output_graph_def; - int precision_mode; + TrtPrecisionMode precision_mode; int minimum_segment_size; const tensorflow::grappler::GraphProperties* graph_properties; const tensorflow::grappler::Cluster* cluster; @@ -99,9 +99,10 @@ 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, - int precision_mode = 1, int minimum_segment_size = 3, - bool is_dyn_op = false, int max_cached_engines = 1, - std::vector cached_engine_batches = {}, bool use_calibration = true); + 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); @@ -123,4 +124,4 @@ std::pair GetDeviceAndAllocator( #endif // GOOGLE_TENSORRT #endif // GOOGLE_CUDA -#endif // TENSORFLOW_CONTRIB_TENSORRT_CONVERT_CONVERT_GRAPH_H_ +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_GRAPH_H_ diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc similarity index 96% rename from tensorflow/contrib/tensorrt/convert/convert_graph_test.cc rename to tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc index 2d2bfeb192c1893824c7b30bfad593c62c203392..cabc6ccfa13df77b3bd26f51f35284816141423a 100644 --- a/tensorflow/contrib/tensorrt/convert/convert_graph_test.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc @@ -13,13 +13,13 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/convert/convert_graph.h" +#include "tensorflow/compiler/tf2tensorrt/convert/convert_graph.h" #include #include #include "tensorflow/cc/framework/scope.h" #include "tensorflow/cc/ops/standard_ops.h" -#include "tensorflow/contrib/tensorrt/convert/convert_nodes.h" +#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h" #include "tensorflow/core/common_runtime/device_mgr.h" #include "tensorflow/core/common_runtime/device_set.h" #include "tensorflow/core/framework/tensor_shape.h" @@ -98,7 +98,8 @@ TEST(TrtCandidateSelector, Basics) { grappler::GraphProperties graph_properties(item); TF_EXPECT_OK(graph_properties.InferStatically(true)); - for (const int precision_mode : {FP32MODE, INT8MODE}) { + for (const TrtPrecisionMode precision_mode : + {TrtPrecisionMode::FP32, TrtPrecisionMode::INT8}) { TrtCandidateSelector selector(graph_properties, precision_mode); TF_EXPECT_OK(selector.IsTensorRTCandidate(matmul.operation.node())); ExpectStatus( @@ -113,7 +114,7 @@ TEST(TrtCandidateSelector, Basics) { matmul_with_incompatible_input.operation.node()), error::INTERNAL, "Failed to convert input with index 0 to a TRT_TensorOrWeights"); - if (precision_mode == INT8MODE) { + if (precision_mode == TrtPrecisionMode::INT8) { TF_EXPECT_OK(selector.IsTensorRTCandidate(quantize.operation.node())); } else { ExpectStatus(selector.IsTensorRTCandidate(quantize.operation.node()), diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc similarity index 86% rename from tensorflow/contrib/tensorrt/convert/convert_nodes.cc rename to tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc index adf8831b960172fc29b5d631e5b0533318d4764d..0d5b9851f79e97279ec0680986efe13e56dbd7c5 100644 --- a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/convert/convert_nodes.h" +#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h" #include #include @@ -24,11 +24,14 @@ limitations under the License. #include #include -#include "tensorflow/contrib/tensorrt/convert/utils.h" -#include "tensorflow/contrib/tensorrt/log/trt_logger.h" -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.h" -#include "tensorflow/contrib/tensorrt/resources/trt_resource_manager.h" -#include "tensorflow/contrib/tensorrt/resources/trt_resources.h" +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "tensorflow/compiler/tf2tensorrt/convert/utils.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_resource_manager.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_resources.h" #include "tensorflow/core/framework/node_def.pb.h" // NOLINT #include "tensorflow/core/framework/node_def_builder.h" #include "tensorflow/core/framework/tensor.pb.h" // NOLINT @@ -43,6 +46,7 @@ limitations under the License. #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/protobuf.h" #include "tensorflow/core/platform/tensor_coding.h" #include "tensorflow/core/platform/types.h" @@ -80,10 +84,16 @@ namespace tensorrt { const char* const kInputPHName = "TensorRTInputPH_"; const char* const kOutputPHName = "TensorRTOutputPH_"; +bool IsEngineInput(absl::string_view name) { + return absl::StartsWith(name, kInputPHName); +} +bool IsEngineOutput(absl::string_view name) { + return absl::StartsWith(name, kOutputPHName); +} + namespace convert { -using ::tensorflow::str_util::Split; -using ::tensorflow::strings::StrAppend; -using ::tensorflow::strings::StrCat; +using absl::StrAppend; +using absl::StrCat; inline tensorflow::Status ConvertDType(tensorflow::DataType tf_dtype, nvinfer1::DataType* trt_dtype) { @@ -286,8 +296,8 @@ Status Converter::GetTrtBroadcastShape( const int max_nb_dims = nvinfer1::Dims::MAX_DIMS + 1; auto compute_output_dims = - [max_nb_dims](const TRT_TensorOrWeights& input, int broadcast_num_dims, - int* output_dims_array, nvinfer1::Dims* output_dims) { + [](const TRT_TensorOrWeights& input, int broadcast_num_dims, + int* output_dims_array, nvinfer1::Dims* output_dims) { const nvinfer1::Dims input_dims = input.GetTrtDims(); std::fill(output_dims_array, output_dims_array + max_nb_dims, 1); std::copy(input_dims.d, input_dims.d + input_dims.nbDims, @@ -334,6 +344,41 @@ Status Converter::GetTrtBroadcastShape( return Status::OK(); } +nvinfer1::ITensor* Converter::CreateConstantLayer( + const TRT_ShapedWeights& weights, const nvinfer1::Dims& dims) { + nvinfer1::Weights trt_weights = weights.GetTrtWeights(); + nvinfer1::IConstantLayer* layer = network()->addConstant(dims, trt_weights); + if (!layer) return nullptr; + const nvinfer1::DataType trt_dtype = trt_weights.type; + nvinfer1::ITensor* trt_tensor = layer->getOutput(0); + // TODO(laigd): there is a bug in TensorRT 5.0 library that, if we don't set + // the data type below, it will always be kFLOAT regardless what the data type + // of the weights is. Once NVIDIA fixes this bug, we should remove the data + // type setting logic below and test should still pass. + trt_tensor->setType(trt_dtype); + return trt_tensor; +} + +tensorflow::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); + auto weights_ptr = + static_cast(const_cast(weights.GetValues())); + weights_ptr[0] = value; + *tensor = params->converter->CreateConstantLayer(weights, broadcastable_dims); + TFTRT_RETURN_ERROR_IF_NULLPTR(*tensor, params->node_def.name()); + params->converter->ProvideQuantizationRange( + const_cast(*tensor), value, value); + return Status::OK(); +} + inline bool DimsEqual(const nvinfer1::Dims& dim_l, const nvinfer1::Dims& dim_r) { if (dim_l.nbDims != dim_r.nbDims) { @@ -848,7 +893,7 @@ Status TrtNodeValidator::ConvertConstToWeights( } Converter::Converter(nvinfer1::INetworkDefinition* trt_network, - int precision_mode, bool use_calibration) + TrtPrecisionMode precision_mode, bool use_calibration) : trt_network_(trt_network), precision_mode_(precision_mode), use_calibration_(use_calibration) { @@ -875,13 +920,15 @@ Status Converter::ConvertNode(const NodeDef& node_def) { for (size_t i = 0; i < outputs.size(); ++i) { TRT_TensorOrWeights& output = outputs[i]; string output_name = node_def.name(); - if (i != 0) output_name = StrCat(output_name, ":", i); + if (i != 0) absl::StrAppend(&output_name, ":", i); // We need to check the name before setting it. If the input is one of the // engine input, setting the name here will overwrite engine input // bindings which will cause runtime error. + // TODO(tmorris): Remove this work-around once we use TRT's IIdentityLayer + // in ConvertIdentity. if (output.is_tensor()) { const char* tensor_name = output.tensor()->getName(); - if (!tensorflow::str_util::StartsWith(tensor_name, kInputPHName)) { + if (!IsEngineInput(tensor_name)) { // TRT initializes tensor names as "(Unnamed ITensor* N)". We rename // them to match their corresponding TensorFlow name. // Note: ITensors that we create internally within TF-TRT which are @@ -927,22 +974,45 @@ Status Converter::AddInputTensor(const string& name, nvinfer1::DataType dtype, } Status Converter::RenameAndMarkOutputTensors( - const std::vector>& output_tensors) { + const std::vector& output_tensors) { for (const auto& output : output_tensors) { TRT_TensorOrWeights tensor_or_weights; - TF_RETURN_IF_ERROR(GetTensorOrWeights(output.first, &tensor_or_weights)); + TF_RETURN_IF_ERROR( + GetTensorOrWeights(output.source_tensor_name, &tensor_or_weights)); if (!tensor_or_weights.is_tensor()) { - return errors::InvalidArgument("Output ", output.first, + return errors::InvalidArgument("Output ", output.source_tensor_name, " is weights not tensor"); } nvinfer1::ITensor* tensor = tensor_or_weights.tensor(); if (tensor == nullptr) { - return errors::NotFound("Output tensor not found: ", output.first); + return errors::NotFound("Output tensor not found: ", + output.source_tensor_name); + } + // Check if this tensor has already been marked as an input or output. + // + // ConvertIdentity can cause the same tensor to be repeated in + // output_tensors, which can cause us to overwrite the name of the output + // tensor binding. For example, if we rename OutputPH_0 to OutputPH_1 then + // we won't be able to locate OutputPH_0 during runtime. To fix this, + // duplicate the tensor using no-op shuffle. + // + // TODO(tmorris): Remove this work-around once we use TRT's IIdentityLayer + // in ConvertIdentity. + if (IsEngineInput(tensor->getName()) || IsEngineOutput(tensor->getName())) { + // Using shuffle layer for identity by not setting reshape or transpose. + nvinfer1::IShuffleLayer* layer = network()->addShuffle(*tensor); + TFTRT_RETURN_ERROR_IF_NULLPTR( + layer, StrCat("Output Copy for ", tensor->getName())); + MarkQuantizationRangesAsInferrable(tensor, layer->getOutput(0)); + tensor = layer->getOutput(0); } - tensor->setName(output.second.c_str()); - VLOG(1) << "Marking output tensor " << output.first << ", as output tensor " - << output.second; + tensor->setName(output.dest_node_name.c_str()); network()->markOutput(*tensor); + // Set type after marking as output. TRT only supports setType for engine + // outputs and inputs (type is inferred otherwise). + tensor->setType(output.trt_dtype); + VLOG(1) << "Marking output TRT tensor " << output.source_tensor_name + << ", which feeds TF node " << output.dest_node_name; } return Status::OK(); } @@ -1086,11 +1156,9 @@ Status Converter::PrepareTensorForShape(const TRT_TensorOrWeights& input, *tensor = layer->getOutput(0); } } else { - nvinfer1::IConstantLayer* layer = - this->network()->addConstant(dims, input.weights().GetTrtWeights()); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, "TF-TRT Internal Reshape"); - *tensor = layer->getOutput(0); - if (precision_mode() == INT8MODE && !use_calibration()) { + *tensor = CreateConstantLayer(input.weights(), dims); + TFTRT_RETURN_ERROR_IF_NULLPTR(*tensor, "TF-TRT Internal Reshape"); + if (precision_mode() == TrtPrecisionMode::INT8 && !use_calibration()) { // If we are in int8 mode and not calibrating, we need to explicitly set a // quantization range for the output tensor of the IConstantLayer. Here we // set the range to [min(weights), max(weights)]. @@ -1125,7 +1193,7 @@ void Converter::ProvideQuantizationRange(nvinfer1::ITensor* tensor, } void Converter::MaybeApplyQuantizationRanges() { - if (precision_mode() != INT8MODE) return; + if (precision_mode() != TrtPrecisionMode::INT8) return; // Infer ranges across marked ops. PropagateQuantizationRanges(); @@ -1248,6 +1316,39 @@ Status Converter::GetInputs(const tensorflow::NodeDef& node_def, return tensorflow::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( + 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( + 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()); + } + // 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 tensorflow::Status::OK(); +} + TRT_ShapedWeights ConvertFP32ToFP16(TrtWeightStore* store, const TRT_ShapedWeights& weights_src) { auto dtype_new = tensorflow::DataType::DT_HALF; @@ -1440,7 +1541,7 @@ Status BinaryTensorOpWeight(OpConverterParams* params, const_cast(tensor), permutation, &tensor)); } - if (params->converter->precision_mode() == FP16MODE) { + if (params->converter->precision_mode() == TrtPrecisionMode::FP16) { weights = ConvertFP32ToFP16(params->weight_store, weights); } @@ -1483,7 +1584,7 @@ Status BinaryTensorOpWeight(OpConverterParams* params, // Because of this issue, fall back to BinaryTensorOpTensor if we are // doing INT8 with no calibration. There is most likely no performance // penalty by falling back here. - if (params->converter->precision_mode() == INT8MODE && + if (params->converter->precision_mode() == TrtPrecisionMode::INT8 && !params->converter->use_calibration()) { return errors::Unimplemented( "Intermediate quantization range cannot be determined without" @@ -1533,92 +1634,126 @@ Status BinaryTensorOpWeight(OpConverterParams* params, return tensorflow::Status::OK(); } -enum class ConvolutionType { DEFAULT, DEPTHWISE_CONV }; - -tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, int group) { +tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, int group, + bool is_conv2d_backprop_input) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; - if (inputs.at(0).is_weights()) { + TRT_TensorOrWeights backprop_output_size; + const nvinfer1::ITensor* tensor = nullptr; + if (is_conv2d_backprop_input) { + // In the case when Conv2dBackpropInput is used for conv2d_transpose, these + // inputs correspond to: output size, filter, and input. + TF_RETURN_IF_ERROR(CheckInputsWeights( + *params, + {{"input_sizes", true}, {"filter", true}, {"out_backprop", false}})); + backprop_output_size = inputs.at(0); + tensor = inputs.at(2).tensor(); + } else { + TF_RETURN_IF_ERROR( + CheckInputsWeights(*params, {{"input", false}, {"filter", true}})); + tensor = inputs.at(0).tensor(); + } + 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()); + } + 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"); + if (tf_dilations.size() != 4) { + return tensorflow::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( - node_def.op(), " is only implemented for tensors, not weights, at ", + "Dilation rate must be 1 for batch and channel dimensions, at ", node_def.name()); } - if (inputs.at(1).is_tensor()) { - return tensorflow::errors::Unimplemented("Kernel for ", node_def.op(), - " must be constant weights, 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( + "Dilation with Conv2DBackpropInput (conv2d_transpose) is not supported", + ", at ", node_def.name()); } - TRT_ShapedWeights weights_rsck = inputs.at(1).weights(); - VLOG(2) << "weight shape: " << weights_rsck.DebugString(); - if (weights_rsck.shape_.nbDims != 4) { - return tensorflow::errors::Internal( - "Conv2D expects kernel of dimension 4, at: " + node_def.name()); + + const auto tf_stride = attrs.get>("strides"); + if (tf_stride.size() != 4) { + return tensorflow::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( + "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(); - const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); - TFAttrs attrs(node_def); - - int h_index = 2; - int w_index = 3; - auto data_format = attrs.get("data_format"); - if (data_format == "NHWC") { + // Transpose to NCHW (NCHW is required for IConvLayer). + const bool need_transpose = (data_format == "NHWC"); + if (need_transpose) { TF_RETURN_IF_ERROR(params->converter->TransposeTensor( const_cast(tensor), {0, 3, 1, 2}, &tensor)); - h_index = 1; - w_index = 2; - // TODO(jie): transpose it } - - // tensor after transpose (NCHW) + // Dimensions of transposed tensor. const auto tensor_dim = tensor->getDimensions(); - int num_groups = group; - if (num_groups == 0) num_groups = tensor_dim.d[0]; // depthwise convolution - VLOG(2) << "groups count: " << num_groups; + // group == 0 signifies that this is a depthwise convolution, so set + // num_groups to size of input's channel dim. For a non-depthwise conv, + // num_groups will be 1. + const int num_groups = (group == 0) ? tensor_dim.d[0] : group; - if (params->converter->precision_mode() == FP16MODE) { - weights_rsck = - ConvertFP32ToFP16(params->weight_store, inputs.at(1).weights()); + if (params->converter->precision_mode() == TrtPrecisionMode::FP16) { + weights_rsck = ConvertFP32ToFP16(params->weight_store, weights_rsck); } - + // For conv, TF weights are RSCK, and TRT expects KCRS. + // For backprop, TF weights are RSKC, and TRT expects CKRS. + // Therefore, this reorder will work for both cases. TRT_ShapedWeights weights = params->weight_store->GetTempWeights(weights_rsck); ReorderRSCKToKCRS(weights_rsck, &weights, num_groups); TRT_ShapedWeights biases(weights.type_); - const int noutput = weights.shape_.d[0] * num_groups; + const int output_axis = is_conv2d_backprop_input ? 1 : 0; + const int noutput = weights.shape_.d[output_axis] * num_groups; nvinfer1::DimsHW kernel_size; kernel_size.h() = weights.shape_.d[2]; kernel_size.w() = weights.shape_.d[3]; - VLOG(2) << "RSCK: " << weights.DebugString(); - VLOG(2) << "kernel size: " << kernel_size.h() << ", " << kernel_size.w(); - - // TODO(jie): stride. (NHWC/NCHW) - const auto tf_stride = attrs.get>("strides"); - VLOG(2) << "h_INDEX" << h_index << ", w_index " << w_index; - VLOG(2) << "stride: " << tf_stride[0] << tf_stride[1] << tf_stride[2] - << tf_stride[3]; - const nvinfer1::DimsHW stride(tf_stride[h_index], tf_stride[w_index]); + // Add padding. std::vector> padding; - // TODO(jie): padding. if (attrs.get("padding") == "SAME") { - // This is NCHW tensor with no batch dimension. - // 1 -> h - // 2 -> w - padding = CreateSamePadding( - stride, kernel_size, - {static_cast(tensor_dim.d[1]), static_cast(tensor_dim.d[2])}); + nvinfer1::DimsHW effective_kernel_size = kernel_size; + effective_kernel_size.h() += (kernel_size.h() - 1) * (dilation.h() - 1); + effective_kernel_size.w() += (kernel_size.w() - 1) * (dilation.w() - 1); + std::vector input_dims; + if (is_conv2d_backprop_input) { + // For backprop, calculate padding based on "input_sizes" input, which + // actually corresponds to output size. ("input_sizes" makes sense in the + // context of Conv2DBackpropInput). + // We use h_index and w_index instead of 1 and 2 because we havent + // transposed backprop_output_size along with the input. + auto output_size_weights = static_cast( + const_cast(backprop_output_size.weights().GetValues())); + input_dims = {output_size_weights[h_index], output_size_weights[w_index]}; + } else { + // Use 1 and 2 because tensor_dim has the dimensions of the transposed + // input. + input_dims = {static_cast(tensor_dim.d[1]), + static_cast(tensor_dim.d[2])}; + } + padding = CreateSamePadding(stride, effective_kernel_size, input_dims); } else { padding = {{0, 0}, {0, 0}}; } - if (padding[0].first != padding[0].second || padding[1].first != padding[1].second) { - // TODO(jie): handle asymmetric padding - VLOG(2) << "Padding!!!: " << padding[0].first << padding[0].second - << padding[1].first << padding[1].second; - VLOG(2) << "TENSOR before: " << DebugString(tensor->getDimensions()); + // Handle asymmetric padding. auto pad_layer = params->converter->network()->addPadding( *const_cast(tensor), nvinfer1::DimsHW(padding[0].first, padding[1].first), @@ -1628,24 +1763,38 @@ tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, int group) { const_cast(tensor), pad_layer->getOutput(0)); padding = {{0, 0}, {0, 0}}; tensor = pad_layer->getOutput(0); - VLOG(2) << "TENSOR after: " << DebugString(tensor->getDimensions()); } - nvinfer1::IConvolutionLayer* layer = - params->converter->network()->addConvolution( - *const_cast(tensor), noutput, kernel_size, - weights.GetTrtWeights(), biases.GetTrtWeights()); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + // Add convolution. + nvinfer1::ILayer* conv_layer = nullptr; + if (is_conv2d_backprop_input) { + nvinfer1::IDeconvolutionLayer* layer = + params->converter->network()->addDeconvolution( + *const_cast(tensor), noutput, kernel_size, + weights.GetTrtWeights(), biases.GetTrtWeights()); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + layer->setStride(stride); + layer->setPadding({padding[0].first, padding[1].first}); + layer->setName(node_def.name().c_str()); + layer->setNbGroups(num_groups); + conv_layer = layer; + } else { + nvinfer1::IConvolutionLayer* layer = + params->converter->network()->addConvolution( + *const_cast(tensor), noutput, kernel_size, + weights.GetTrtWeights(), biases.GetTrtWeights()); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + layer->setStride(stride); + layer->setPadding({padding[0].first, padding[1].first}); + layer->setName(node_def.name().c_str()); + layer->setNbGroups(num_groups); + layer->setDilation(dilation); + conv_layer = layer; + } + const nvinfer1::ITensor* output_tensor = conv_layer->getOutput(0); - layer->setStride(stride); - layer->setPadding({padding[0].first, padding[1].first}); - layer->setName(node_def.name().c_str()); - layer->setNbGroups(num_groups); - const nvinfer1::ITensor* output_tensor = layer->getOutput(0); - VLOG(2) << "TENSOR out: " << DebugString(output_tensor->getDimensions()); - VLOG(2) << "data_format: " << data_format; - if (data_format == "NHWC") { - // TODO(jie): transpose it back! + // Restore transpose. + if (need_transpose) { TF_RETURN_IF_ERROR(params->converter->TransposeTensor( const_cast(output_tensor), {0, 2, 3, 1}, &output_tensor)); @@ -1655,18 +1804,6 @@ tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, int group) { return tensorflow::Status::OK(); } -tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, - ConvolutionType type) { - switch (type) { - case ConvolutionType::DEFAULT: - return ConvertConv2DHelper(params, 1); - case ConvolutionType::DEPTHWISE_CONV: - return ConvertConv2DHelper(params, 0); - } - return tensorflow::errors::Unimplemented("Unsupported convolution type, at ", - params->node_def.name()); -} - Status BinaryTensorOpTensor(OpConverterParams* params, const TRT_TensorOrWeights& operand_l, const TRT_TensorOrWeights& operand_r) { @@ -1694,6 +1831,13 @@ Status BinaryTensorOpTensor(OpConverterParams* params, "Unsupported binary op broadcast scheme for op ", node_def.name(), ": ", status.error_message()); } + TFAttrs attrs(node_def); + nvinfer1::DataType dtype = attrs.get("T"); + if (dtype == nvinfer1::DataType::kINT32) { + return errors::Unimplemented("Binary op ", node_def.op(), + " does not support INT32, at ", + node_def.name()); + } if (params->validation_only) return Status::OK(); const nvinfer1::ITensor* tensor_l = nullptr; @@ -1710,8 +1854,6 @@ Status BinaryTensorOpTensor(OpConverterParams* params, } // Check type consistency. - TFAttrs attrs(node_def); - nvinfer1::DataType dtype = attrs.get("T"); TFTRT_CHECK_EQ_TYPE(tensor_l->getType(), dtype) << DebugString(tensor_l->getType()) << " vs " << DebugString(dtype); TFTRT_CHECK_EQ_TYPE(tensor_r->getType(), dtype) @@ -1771,12 +1913,8 @@ tensorflow::Status ConvertPlugin(OpConverterParams* params) { tensorflow::Status ConvertTranspose(OpConverterParams* params) { const auto& inputs = params->inputs; - if (inputs.size() != 2 || !inputs.at(0).is_tensor() || - !inputs.at(1).is_weights()) { - return tensorflow::errors::InvalidArgument( - "Input expects tensor and weights, at ", params->node_def.name()); - } - + TF_RETURN_IF_ERROR( + CheckInputsWeights(*params, {{"x", false}, {"perm", true}})); // Get the permutation from weights. TRT_ShapedWeights weights = inputs.at(1).weights(); const int* weights_ptr = @@ -1809,11 +1947,8 @@ tensorflow::Status ConvertTranspose(OpConverterParams* params) { tensorflow::Status ConvertReshape(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; - if (inputs.size() != 2 || !inputs.at(1).is_weights()) { - return tensorflow::errors::InvalidArgument( - "Input expects weights for shape, at ", node_def.name()); - } - + TF_RETURN_IF_ERROR( + CheckInputsWeights(*params, {{"tensor", false}, {"shape", true}})); TRT_TensorOrWeights input_tensor = inputs.at(0); TRT_ShapedWeights weights = inputs.at(1).weights(); if (weights.count() == 0) { @@ -1909,18 +2044,8 @@ tensorflow::Status ConvertReshape(OpConverterParams* params) { tensorflow::Status ConvertExpandDims(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; - if (inputs.size() != 2) { - return tensorflow::errors::InvalidArgument( - "Two inputs expected for ExpandDims, at ", node_def.name()); - } - if (inputs.at(0).is_weights()) { - return tensorflow::errors::Unimplemented( - "ExpandDims expects tensor for input, at ", node_def.name()); - } - if (!inputs.at(1).is_weights()) { - return tensorflow::errors::InvalidArgument( - "ExpandDims expects weights for axis, at ", node_def.name()); - } + TF_RETURN_IF_ERROR( + CheckInputsWeights(*params, {{"input", false}, {"axis", true}})); // Get input shape as vector. TRT_TensorOrWeights input_tensor = inputs.at(0); const nvinfer1::Dims dims = input_tensor.GetTrtDims(); @@ -1970,14 +2095,7 @@ tensorflow::Status ConvertExpandDims(OpConverterParams* params) { tensorflow::Status ConvertSqueeze(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; - if (inputs.size() != 1) { - return tensorflow::errors::InvalidArgument( - "One input expected for Squeeze, at ", node_def.name()); - } - if (inputs.at(0).is_weights()) { - return tensorflow::errors::Unimplemented( - "Squeeze expects tensor for input, at ", node_def.name()); - } + TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"input", false}})); // Get input shape. TRT_TensorOrWeights input_tensor = inputs.at(0); const nvinfer1::Dims dims = input_tensor.GetTrtDims(); @@ -1988,7 +2106,7 @@ tensorflow::Status ConvertSqueeze(OpConverterParams* params) { // Mark axes to remove by setting them to 0. TFAttrs attrs(node_def); auto squeeze_dims = attrs.get>("squeeze_dims"); - if (squeeze_dims.size() == 0) { + if (squeeze_dims.empty()) { return tensorflow::errors::Unimplemented( "Squeeze is only implemented for explicit dims, at ", node_def.name()); } @@ -2079,20 +2197,9 @@ tensorflow::Status GetStridedSliceBound(const std::vector& input_dims, tensorflow::Status ConvertStridedSlice(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; - if (inputs.size() != 4) { - return tensorflow::errors::InvalidArgument( - "StridedSlice expects 4 inputs, at ", node_def.name()); - } - if (!inputs.at(1).is_weights() || !inputs.at(2).is_weights() || - !inputs.at(3).is_weights()) { - return tensorflow::errors::InvalidArgument( - "StridedSlice expects weights for begin, end, and strides, at ", - node_def.name()); - } - if (!inputs.at(0).is_tensor()) { - return tensorflow::errors::Unimplemented( - "StridedSlice is only implemented for tensors, at ", node_def.name()); - } + TF_RETURN_IF_ERROR(CheckInputsWeights( + *params, + {{"input", false}, {"begin", true}, {"end", true}, {"strides", true}})); // Get input dims. nvinfer1::Dims dims = inputs.at(0).GetTrtDims(); std::vector input_dims(dims.d, dims.d + dims.nbDims); @@ -2167,7 +2274,7 @@ tensorflow::Status ConvertStridedSlice(OpConverterParams* params) { pad_dims.push_back(i); } } - if (pad_dims.size() == 0) { + if (pad_dims.empty()) { // No dimensions are changed. We could create a padding layer anyway with // values of 0. if (params->validation_only) return Status::OK(); @@ -2273,21 +2380,21 @@ tensorflow::Status ConvertStridedSlice(OpConverterParams* params) { } tensorflow::Status ConvertConv2D(OpConverterParams* params) { - return ConvertConv2DHelper(params, ConvolutionType::DEFAULT); + return ConvertConv2DHelper(params, 1, /*is_conv2d_backprop_input=*/false); } tensorflow::Status ConvertConv2DDepthwise(OpConverterParams* params) { - return ConvertConv2DHelper(params, ConvolutionType::DEPTHWISE_CONV); + return ConvertConv2DHelper(params, 0, /*is_conv2d_backprop_input=*/false); +} + +tensorflow::Status ConvertConv2DBackpropInput(OpConverterParams* params) { + return ConvertConv2DHelper(params, 1, /*is_conv2d_backprop_input=*/true); } tensorflow::Status ConvertPool(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; - if (inputs.at(0).is_weights()) { - return tensorflow::errors::Unimplemented( - node_def.op(), " is only implemented for tensors, not weights, at ", - node_def.name()); - } + TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"input", false}})); nvinfer1::PoolingType type; if (node_def.op() == "MaxPool") { type = nvinfer1::PoolingType::kMAX; @@ -2374,7 +2481,9 @@ tensorflow::Status ConvertPool(OpConverterParams* params) { return tensorflow::Status::OK(); } -tensorflow::Status ConvertActivation(OpConverterParams* params) { +// TODO(tmorris): Use ActivationType::kLEAKY_RELU in TRT 5.1+ once perf +// improves. +tensorflow::Status ConvertLeakyRelu(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; if (inputs.size() != 1) { @@ -2386,6 +2495,47 @@ tensorflow::Status ConvertActivation(OpConverterParams* params) { node_def.op(), " is only implemented for tensors, at ", node_def.name()); } + TFAttrs attrs(node_def); + const float alpha = attrs.get("alpha"); + if (alpha < 0.0f || alpha > 1.0f) { + return tensorflow::errors::Unimplemented( + "Alpha value for LeakyRelu must be between 0 and 1, at ", + node_def.name()); + } + if (params->validation_only) return tensorflow::Status::OK(); + + // Input Tensor + const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); + // Create const for alpha. + const nvinfer1::ITensor* const_alpha_tensor = nullptr; + TF_RETURN_IF_ERROR(CreateBroadcastableScalarConstant( + params, alpha, tensor->getDimensions(), &const_alpha_tensor)); + // alpha * x + nvinfer1::IElementWiseLayer* mul_layer = + params->converter->network()->addElementWise( + *const_cast(tensor), + *const_cast(const_alpha_tensor), + nvinfer1::ElementWiseOperation::kPROD); + TFTRT_RETURN_ERROR_IF_NULLPTR(mul_layer, node_def.name()); + // max(x, alpha * x) + nvinfer1::IElementWiseLayer* max_layer = + params->converter->network()->addElementWise( + *const_cast(tensor), + *const_cast(mul_layer->getOutput(0)), + nvinfer1::ElementWiseOperation::kMAX); + TFTRT_RETURN_ERROR_IF_NULLPTR(max_layer, node_def.name()); + nvinfer1::ITensor* output_tensor = max_layer->getOutput(0); + params->converter->MarkQuantizationRangesAsInferrable( + output_tensor, const_cast(mul_layer->getOutput(0))); + + params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return Status::OK(); +} + +tensorflow::Status ConvertActivation(OpConverterParams* params) { + const auto& inputs = params->inputs; + const auto& node_def = params->node_def; + TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"input", false}})); static const std::unordered_map ops{ {"Relu", nvinfer1::ActivationType::kRELU}, {"Sigmoid", nvinfer1::ActivationType::kSIGMOID}, @@ -2419,19 +2569,19 @@ tensorflow::Status ConvertActivation(OpConverterParams* params) { Status ConvertQuantize(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; - if ((inputs.size() == 0) || - (node_def.op() == "FakeQuantWithMinMaxArgs" && inputs.size() != 1) || - (node_def.op() == "FakeQuantWithMinMaxVars" && inputs.size() != 3) || - (node_def.op() == "QuantizeAndDequantizeV2" && inputs.size() != 3) || - (node_def.op() == "QuantizeAndDequantizeV3" && inputs.size() != 4)) { - return errors::InvalidArgument("Invalid number of inputs for ", - node_def.op(), ", at ", node_def.name()); - } - if (inputs.at(0).is_weights()) { - // TensorRT will automatically quantize weights, so we will ignore ranges - // for weights. - params->outputs->push_back(inputs.at(0)); - return Status::OK(); + if (node_def.op() == "FakeQuantWithMinMaxArgs") { + TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"input", false}})); + } else if (node_def.op() == "FakeQuantWithMinMaxVars") { + TF_RETURN_IF_ERROR(CheckInputsWeights( + *params, {{"input", false}, {"min", true}, {"max", true}})); + } else if (node_def.op() == "QuantizeAndDequantizeV2") { + TF_RETURN_IF_ERROR(CheckInputsWeights( + *params, {{"input", false}, {"input_min", true}, {"input_max", true}})); + } else if (node_def.op() == "QuantizeAndDequantizeV3") { + TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"input", false}, + {"input_min", true}, + {"input_max", true}, + {"num_bits", true}})); } float min_range = 0.0f; float max_range = 0.0f; @@ -2448,11 +2598,6 @@ Status ConvertQuantize(OpConverterParams* params) { node_def.op() == "QuantizeAndDequantizeV2" || node_def.op() == "QuantizeAndDequantizeV3") { // Get ranges via inputs. - if (!inputs.at(1).is_weights() || !inputs.at(2).is_weights()) { - return errors::InvalidArgument("Min and max inputs for ", node_def.op(), - " must be weights not tensors, at ", - node_def.name()); - } auto get_weights_value = [&inputs](int index) { auto raw_weights = static_cast( const_cast(inputs.at(index).weights().GetValues())); @@ -2483,20 +2628,11 @@ Status ConvertQuantize(OpConverterParams* params) { return Status::OK(); } -// TODO(pdavoodi): we should update relu6 implementation once TensorRT supports -// Relu6 natively. +// TODO(tmorris): Use ActivationType::kCLIP in TRT 5.1+ once perf improves. tensorflow::Status ConvertRelu6(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; - if (inputs.size() != 1) { - return tensorflow::errors::InvalidArgument( - "Invalid number of inputs for Relu6, at ", node_def.name()); - } - if (inputs.at(0).is_weights()) { - return tensorflow::errors::Unimplemented( - "Relu6 is only implemented for tensors, not weights, at ", - node_def.name()); - } + TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"input", false}})); if (params->validation_only) return Status::OK(); // *************************************************************************** // TensorRT does not implement Relu6 natively. This function converts Relu6 op @@ -2520,35 +2656,18 @@ tensorflow::Status ConvertRelu6(OpConverterParams* params) { params->converter->ProvideQuantizationRange(relu_layer->getOutput(0), 0.0f, 6.0f); - // Create a constant layer to store the floating point weight i.e. 6.0f This - // tensor will be broadcasted uniformly during elementwise `min` operation. - // The constant has to have the same rank as the input in order for TRT to - // broadcast - nvinfer1::Dims dims; - dims.nbDims = relu_layer->getOutput(0)->getDimensions().nbDims; - for (int i = 0; i < dims.nbDims; i++) { - dims.d[i] = 1; - } - TRT_ShapedWeights weights = params->weight_store->GetTempWeights( - tensorflow::DataType::DT_FLOAT, dims); - auto weights_ptr = - static_cast(const_cast(weights.GetValues())); - weights_ptr[0] = 6.0f; - nvinfer1::IConstantLayer* const6_layer = - params->converter->network()->addConstant(dims, weights.GetTrtWeights()); - TFTRT_RETURN_ERROR_IF_NULLPTR(const6_layer, node_def.name()); - params->converter->ProvideQuantizationRange(const6_layer->getOutput(0), 0.0f, - 6.0f); + // Create a constant layer to store the floating point weight i.e. 6.0f + const nvinfer1::ITensor* const6_tensor = nullptr; + TF_RETURN_IF_ERROR(CreateBroadcastableScalarConstant( + params, 6.0f, relu_layer->getOutput(0)->getDimensions(), &const6_tensor)); // ElementWise Min Operation // Min op is a nop for INT8 execution path, as the input tensor // to this layer will only have values in range [0.f, 6.0f]. - const nvinfer1::ITensor* tensor_l = relu_layer->getOutput(0); - const nvinfer1::ITensor* tensor_r = const6_layer->getOutput(0); nvinfer1::IElementWiseLayer* relu6_layer = params->converter->network()->addElementWise( - *const_cast(tensor_l), - *const_cast(tensor_r), + *const_cast(relu_layer->getOutput(0)), + *const_cast(const6_tensor), nvinfer1::ElementWiseOperation::kMIN); TFTRT_RETURN_ERROR_IF_NULLPTR(relu6_layer, node_def.name()); nvinfer1::ITensor* output_tensor = relu6_layer->getOutput(0); @@ -2561,17 +2680,20 @@ tensorflow::Status ConvertRelu6(OpConverterParams* params) { tensorflow::Status ConvertBiasAdd(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; - if (inputs.size() != 2 || !inputs.at(0).is_tensor() || - !inputs.at(1).is_weights()) { - return errors::InvalidArgument("Input expects tensor and weights, at ", - node_def.name()); + TF_RETURN_IF_ERROR( + CheckInputsWeights(*params, {{"value", false}, {"bias", true}})); + TFAttrs attrs(node_def); + tensorflow::DataType tf_dtype = attrs.get("T"); + if (tf_dtype != DataType::DT_FLOAT && tf_dtype != DataType::DT_HALF) { + return errors::Unimplemented("Data type is not supported, for node ", + node_def.name(), " got ", + DataTypeString(tf_dtype)); } if (params->validation_only) return Status::OK(); nvinfer1::ITensor* tensor = const_cast(inputs.at(0).tensor()); const nvinfer1::Dims original_dims = tensor->getDimensions(); - TFAttrs attrs(node_def); const string data_format = attrs.get("data_format"); const int channel_index = (data_format == "NHWC" ? original_dims.nbDims - 1 : 0); @@ -2617,7 +2739,7 @@ tensorflow::Status ConvertBiasAdd(OpConverterParams* params) { } TRT_ShapedWeights weights = inputs.at(1).weights(); - if (params->converter->precision_mode() == FP16MODE) { + if (params->converter->precision_mode() == TrtPrecisionMode::FP16) { weights = ConvertFP32ToFP16(params->weight_store, weights); } nvinfer1::ScaleMode mode = nvinfer1::ScaleMode::kCHANNEL; @@ -2661,43 +2783,69 @@ tensorflow::Status ConvertBiasAdd(OpConverterParams* params) { return Status::OK(); } -Status GetTensorDimsWithProtoShape(const Tensor& tensor, - int tensor_proto_array_len, - nvinfer1::Dims* dims) { +void GetTensorDimsWithProtoShape(const Tensor& tensor, nvinfer1::Dims* dims) { if (tensor.dims() > 0) { *dims = GetTrtDimsForTensor(tensor); - if (TrtDimsNumElements(*dims) != tensor_proto_array_len && - tensor_proto_array_len != 1) { - return errors::InvalidArgument( - "Broadcast on weights only supports kCHANNEL and kUNIFORM"); - } } else { dims->nbDims = 1; // No dimension provided. Flatten it. - dims->d[0] = tensor_proto_array_len; + dims->d[0] = tensor.NumElements(); dims->type[0] = nvinfer1::DimensionType::kSPATIAL; for (int i = 1; i < nvinfer1::Dims::MAX_DIMS; ++i) { dims->d[i] = 0; } } - return Status::OK(); } -template -Status TfTensorToTrtWeights(const DataType dtype, const Tensor& tensor, - const CType* tensor_proto_array, - int tensor_proto_array_len, TrtWeightStore* store, +Status TfTensorToTrtWeights(const Tensor& tensor, TrtWeightStore* weight_store, TRT_ShapedWeights* weights) { + const DataType dtype = tensor.dtype(); + + // We always convert the integer constants to INT32, since TRT INT8 is for + // quantized inference. + // + // TODO(aaroey): FP16 will remain in half format and is not converted to + // FP32, but the converter currently uses all float weights as FP32. Fix + // this. + const DataType converted_dtype = + (dtype == DT_INT16 || dtype == DT_INT8 || dtype == DT_UINT8 ? DT_INT32 + : dtype); + + // Verify that the dtype is supported by TensorRT. Otherwise, return an error. + nvinfer1::DataType trt_dtype; + TF_RETURN_IF_ERROR(ConvertDType(converted_dtype, &trt_dtype)); + + if (tensor.NumElements() == 0) { + // Return empty weights having converted dtype. + *weights = TRT_ShapedWeights(converted_dtype); + return Status::OK(); + } + nvinfer1::Dims weight_dims; - TF_RETURN_IF_ERROR(GetTensorDimsWithProtoShape(tensor, tensor_proto_array_len, - &weight_dims)); - *weights = store->GetTempWeights(dtype, weight_dims); - void* dst = const_cast(weights->GetValues()); - if (tensor_proto_array_len == 1) { - std::fill_n((CType*)dst, TrtDimsNumElements(weight_dims), - *tensor_proto_array); + GetTensorDimsWithProtoShape(tensor, &weight_dims); + *weights = weight_store->GetTempWeights(converted_dtype, weight_dims); + + // Copy the tensor directly if the tensor does not require cast to the + // supported type. + if (converted_dtype == dtype) { + char* dst = static_cast(const_cast(weights->GetValues())); + memcpy(dst, tensor.tensor_data().data(), tensor.TotalBytes()); + return Status::OK(); + } + + // Copy tensor elements after casting them to the converted DataType. + int32* dst = static_cast(const_cast(weights->GetValues())); + if (dtype == DT_INT16) { + const int16* src = tensor.flat().data(); + std::copy(src, src + tensor.NumElements(), dst); + } else if (dtype == DT_INT8) { + const int8* src = tensor.flat().data(); + std::copy(src, src + tensor.NumElements(), dst); } else { - memcpy(dst, tensor_proto_array, weights->size_bytes()); + // dtype can only be DT_UINT8 at this point. + TFTRT_CHECK_EQ_TYPE(dtype, DT_UINT8); + const uint8* src = tensor.flat().data(); + std::copy(src, src + tensor.NumElements(), dst); } return Status::OK(); } @@ -2715,15 +2863,6 @@ tensorflow::Status ConvertConst(OpConverterParams* params) { "Constant node is expected to have empty input list: ", node_def.name()); } - TFAttrs attrs(node_def); - const DataType dtype = attrs.get("dtype"); - // We always convert the integer constants to kINT32, since TRT kINT8 is for - // quantized inference. - const DataType converted_dtype = - (dtype == DT_INT16 || dtype == DT_INT8 || dtype == DT_UINT8 ? DT_INT32 - : dtype); - nvinfer1::DataType trt_dtype; - TF_RETURN_IF_ERROR(ConvertDType(converted_dtype, &trt_dtype)); // Create shaped weights as output const auto& tensor_proto = node_def.attr().at("value").tensor(); @@ -2733,78 +2872,18 @@ tensorflow::Status ConvertConst(OpConverterParams* params) { node_def.name()); } - TRT_ShapedWeights weights(converted_dtype); - if (tensor.NumElements() == 0) { - // Do nothing. - } else if (!tensor_proto.float_val().empty()) { - TF_RETURN_IF_ERROR(TfTensorToTrtWeights( - converted_dtype, tensor, tensor_proto.float_val().begin(), - tensor_proto.float_val_size(), params->weight_store, &weights)); - } else if (!tensor_proto.int_val().empty()) { - TF_RETURN_IF_ERROR(TfTensorToTrtWeights( - converted_dtype, tensor, tensor_proto.int_val().begin(), - tensor_proto.int_val_size(), params->weight_store, &weights)); - } else if (!tensor_proto.half_val().empty()) { - // TODO(aaroey): implement fp16 conversion. - return errors::Unimplemented("fp16 constant is not supported yet."); - } else if (!tensor_proto.tensor_content().empty()) { - // TODO(aaroey): fp16 will remain in half format and is not converted to - // fp32, but the converter currently uses all float weights as fp32. Fix - // this. - const auto& content = tensor_proto.tensor_content(); - if (content.size() > 0) { - const int dtype_size = tensorflow::DataTypeSize(dtype); - if (content.size() % dtype_size != 0) { - return errors::FailedPrecondition("Tensor content size ", - content.size(), - " is not a multiple of ", dtype_size); - } - nvinfer1::Dims weights_dim; - TF_RETURN_IF_ERROR(GetTensorDimsWithProtoShape( - tensor, content.size() / dtype_size, &weights_dim)); - const int64_t size_bytes = TrtDimsNumElements(weights_dim) * dtype_size; - if (content.size() != size_bytes) { - return errors::FailedPrecondition( - "Tensor size and TensorProto content size mismatch: ", size_bytes, - " vs ", content.size()); - } else if (tensor.NumElements() != content.size() / dtype_size) { - return errors::FailedPrecondition( - "Tensor elements count and TensorProto content size mismatch: ", - tensor.NumElements(), " vs ", content.size() / dtype_size); - } - weights = - params->weight_store->GetTempWeights(converted_dtype, weights_dim); - if (dtype_size == tensorflow::DataTypeSize(converted_dtype)) { - port::CopyToArray(content, static_cast( - const_cast(weights.GetValues()))); - } else { - // Copy out the weights as original data type. - std::vector temp_weights(content.size()); - port::CopyToArray(content, - reinterpret_cast(temp_weights.data())); - int32* dst = - static_cast(const_cast(weights.GetValues())); - // Copy to the weight store as converted data type. - if (dtype == DT_INT16) { - int16* data = reinterpret_cast(temp_weights.data()); - std::copy(data, data + tensor.NumElements(), dst); - } else if (dtype == DT_INT8) { - int8* data = reinterpret_cast(temp_weights.data()); - std::copy(data, data + tensor.NumElements(), dst); - } else if (dtype == DT_UINT8) { - uint8* data = reinterpret_cast(temp_weights.data()); - std::copy(data, data + tensor.NumElements(), dst); - } else { - return errors::FailedPrecondition( - "Unexpected data type: ", DataTypeString(dtype), - " at: ", node_def.name()); - } - } - } - } else { - return errors::Unimplemented("Not supported constant type, at ", - node_def.name()); + TFAttrs attrs(node_def); + const DataType dtype = attrs.get("dtype"); + if (dtype != tensor.dtype()) { + return errors::InvalidArgument("DataType mismatch between attr (", + DataTypeString(dtype), ") and tensor (", + DataTypeString(tensor.dtype()), ")"); } + + TRT_ShapedWeights weights; + TF_RETURN_IF_ERROR( + TfTensorToTrtWeights(tensor, params->weight_store, &weights)); + if (params->outputs != nullptr) { params->outputs->push_back(TRT_TensorOrWeights(weights)); } @@ -2822,9 +2901,13 @@ tensorflow::Status ConvertIdentity(OpConverterParams* params) { Status ConvertBinary(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 + // TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"x", false}, {"y", + // false}})); if (inputs.size() != 2) { - return errors::InvalidArgument("Binary ops require two inputs, at ", - node_def.name()); + return tensorflow::errors::InvalidArgument( + node_def.op(), " got ", inputs.size(), " inputs but expected 2, at ", + node_def.name()); } // Constant folding should have been done by TensorFlow @@ -2874,11 +2957,7 @@ tensorflow::Status ConvertUnary(OpConverterParams* params) { {"Abs", nvinfer1::UnaryOperation::kABS}, {"Reciprocal", nvinfer1::UnaryOperation::kRECIP}, }; - - if (inputs.size() != 1) { - return tensorflow::errors::FailedPrecondition( - "Unary ops require single tensor input, at ", node_def.name()); - } + TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"x", false}})); // TODO(jie): check type const nvinfer1::ITensor* tensor = nullptr; @@ -2893,7 +2972,7 @@ tensorflow::Status ConvertUnary(OpConverterParams* params) { // x -> [Sqrt] -> sqrt(x) -> [Recip] -> 1/sqrt(x) // ^ // need range here - if (params->converter->precision_mode() == INT8MODE && + if (params->converter->precision_mode() == TrtPrecisionMode::INT8 && !params->converter->use_calibration()) { return errors::Unimplemented( "Intermediate quantization range cannot be determined without" @@ -2927,14 +3006,7 @@ tensorflow::Status ConvertUnary(OpConverterParams* params) { tensorflow::Status ConvertSquare(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; - if (inputs.size() != 1) { - return tensorflow::errors::InvalidArgument("Square expects one input, at ", - node_def.name()); - } - if (inputs.at(0).is_weights()) { - return tensorflow::errors::Unimplemented( - "Square is only implemented for tensors, at ", node_def.name()); - } + TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"x", false}})); if (params->validation_only) return Status::OK(); // Constant 2 with same rank as input @@ -2947,18 +3019,15 @@ tensorflow::Status ConvertSquare(OpConverterParams* params) { auto weights_ptr = static_cast(const_cast(weights.GetValues())); weights_ptr[0] = 2.f; - nvinfer1::IConstantLayer* const2_layer = - params->converter->network()->addConstant(dims, weights.GetTrtWeights()); - TFTRT_RETURN_ERROR_IF_NULLPTR(const2_layer, node_def.name()); + nvinfer1::ITensor* const2_tensor = + params->converter->CreateConstantLayer(weights, dims); + TFTRT_RETURN_ERROR_IF_NULLPTR(const2_tensor, node_def.name()); // ElementWise Pow Operation - const nvinfer1::ITensor* tensor_l = inputs.at(0).tensor(); - const nvinfer1::ITensor* tensor_r = const2_layer->getOutput(0); nvinfer1::IElementWiseLayer* layer = params->converter->network()->addElementWise( - *const_cast(tensor_l), - *const_cast(tensor_r), - nvinfer1::ElementWiseOperation::kPOW); + *const_cast(inputs.at(0).tensor()), + *const2_tensor, nvinfer1::ElementWiseOperation::kPOW); TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); nvinfer1::ITensor* output_tensor = layer->getOutput(0); @@ -2969,11 +3038,8 @@ tensorflow::Status ConvertSquare(OpConverterParams* params) { tensorflow::Status ConvertReduce(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; - if (inputs.size() != 2 || !inputs.at(0).is_tensor() || - !inputs.at(1).is_weights()) { - return tensorflow::errors::InvalidArgument( - "Input expects tensor and weights, at", node_def.name()); - } + TF_RETURN_IF_ERROR( + CheckInputsWeights(*params, {{"input", false}, {"axis", true}})); const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); TRT_ShapedWeights index_list = inputs.at(1).weights(); @@ -3034,12 +3100,8 @@ tensorflow::Status ConvertReduce(OpConverterParams* params) { tensorflow::Status ConvertPad(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; - // TODO(aaroey): make a routine for this check and reuse it. - if (inputs.size() != 2 || !inputs.at(0).is_tensor() || - !inputs.at(1).is_weights()) { - return tensorflow::errors::InvalidArgument( - "Input expects tensor and weights, at", node_def.name()); - } + TF_RETURN_IF_ERROR( + CheckInputsWeights(*params, {{"tensor", false}, {"paddings", true}})); // Implement tensor binaryOp weight [channel wise] for now; const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); @@ -3076,7 +3138,7 @@ tensorflow::Status ConvertPad(OpConverterParams* params) { } // No padding at all, we should exit - if (pad_index.size() == 0) { + if (pad_index.empty()) { params->outputs->push_back(inputs.at(0)); return tensorflow::Status::OK(); } @@ -3220,6 +3282,11 @@ tensorflow::Status ConvertConcat(OpConverterParams* params) { tensorflow::Status ConvertFusedBatchNorm(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; + TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"x", false}, + {"scale", true}, + {"offset", true}, + {"mean", true}, + {"variance", true}})); TFAttrs attrs(node_def); float epsilon = attrs.get("epsilon"); auto data_format = attrs.get("data_format"); @@ -3240,21 +3307,6 @@ tensorflow::Status ConvertFusedBatchNorm(OpConverterParams* params) { node_def.op(), " only supports is_training=false, at ", node_def.name()); } - if (inputs.at(0).is_weights()) { - return tensorflow::errors::Unimplemented( - node_def.op(), - " is only implemented for tensor inputs, not weights, at ", - node_def.name()); - } - for (int i = 1; i < 5; i++) { - if (inputs.at(i).is_tensor()) { - return tensorflow::errors::Unimplemented( - node_def.op(), - " must have constant inputs for scale, offset, mean and variance, " - "at ", - node_def.name()); - } - } nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); // Check parameter types @@ -3276,7 +3328,7 @@ tensorflow::Status ConvertFusedBatchNorm(OpConverterParams* params) { TRT_ShapedWeights dummy_power_weights(parameter_type); size_t nweight = 0; for (int i = 1; i < 5; i++) { - nweight = std::max(nweight, (size_t)inputs.at(i).weights().count()); + nweight = std::max(nweight, inputs.at(i).weights().count()); } TRT_ShapedWeights* ptr_shape_weights = nullptr; for (int i = 1; i < 5; i++) { @@ -3411,14 +3463,9 @@ tensorflow::Status ConvertMatMulHelper(OpConverterParams* params, tensorflow::Status ConvertMatMul(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; - if (inputs.size() != 2 || !inputs.at(0).is_tensor() || - !inputs.at(1).is_weights()) { - return errors::InvalidArgument("Input expects tensor and weights, at ", - node_def.name()); - } + TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"a", false}, {"b", true}})); TFAttrs attrs(node_def); - // TODO(jie): INT32 should be converted? tensorflow::DataType tf_dtype = attrs.get("T"); if (tf_dtype != DataType::DT_FLOAT && tf_dtype != DataType::DT_HALF) { return errors::Unimplemented("Data type is not supported, for node ", @@ -3442,9 +3489,16 @@ tensorflow::Status ConvertMatMul(OpConverterParams* params) { tensorflow::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 + // 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()); + } TFAttrs attrs(node_def); - // TODO(jie): INT32 should be converted? tensorflow::DataType tf_dtype = attrs.get("T"); if (tf_dtype != tensorflow::DataType::DT_FLOAT && tf_dtype != tensorflow::DataType::DT_HALF) { @@ -3514,6 +3568,7 @@ tensorflow::Status ConvertBatchMatMul(OpConverterParams* params) { tensorflow::Status ConvertSoftmax(OpConverterParams* params) { const auto& inputs = params->inputs; const auto& node_def = params->node_def; + TF_RETURN_IF_ERROR(CheckInputsWeights(*params, {{"logits", false}})); const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); int nbDims = tensor->getDimensions().nbDims; @@ -3522,6 +3577,8 @@ tensorflow::Status ConvertSoftmax(OpConverterParams* params) { "TensorRT Softmax cannot apply on batch dimension, at" + node_def.name()); } + if (params->validation_only) return Status::OK(); + nvinfer1::ISoftMaxLayer* layer = params->converter->network()->addSoftMax( *const_cast(tensor)); TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); @@ -3537,31 +3594,36 @@ tensorflow::Status ConvertSoftmax(OpConverterParams* params) { tensorflow::Status ConvertTopK(OpConverterParams* params) { const auto& inputs = params->inputs; + if (inputs.size() != 2 || !inputs.at(0).is_tensor() || + !inputs.at(1).is_weights()) { + return errors::InvalidArgument("Input expects tensor and weights, at ", + params->node_def.name()); + } + const auto& node_def = params->node_def; + TF_RETURN_IF_ERROR( + CheckInputsWeights(*params, {{"input", false}, {"k", true}})); const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); - - int nbDims = tensor->getDimensions().nbDims; - if (nbDims == 0) { - return tensorflow::errors::InvalidArgument( - "TensorRT TopK cannot apply on batch dimension, at" + node_def.name()); + const int num_dims = tensor->getDimensions().nbDims; + if (num_dims == 0) { + return errors::InvalidArgument( + "TensorRT TopK cannot apply on batch dimension, at", node_def.name()); } TRT_ShapedWeights k_w = inputs.at(1).weights(); - int k = *(static_cast(const_cast(k_w.GetValues()))); - - nvinfer1::TopKOperation op; - uint32_t reducedAxes = 0; - if (node_def.op() == "TopKV2") { - op = nvinfer1::TopKOperation::kMAX; - reducedAxes |= 1 << (nbDims - 1); - } else { - return tensorflow::errors::Unimplemented( - "Operation: " + node_def.op() + - " not implemented, at: " + node_def.name()); + if (k_w.count() != 1) { + return errors::InvalidArgument("k value of TopK should be a scalar, at", + node_def.name()); } + // Note that ITopKLayer always have sorted outputs, so we don't need to handle + // the 'sorted' attribute of the node. + if (params->validation_only) return Status::OK(); + const nvinfer1::TopKOperation op = nvinfer1::TopKOperation::kMAX; + const int k = *(static_cast(const_cast(k_w.GetValues()))); + const uint32_t reduce_axes = 1 << (num_dims - 1); nvinfer1::ITopKLayer* layer = params->converter->network()->addTopK( - *const_cast(tensor), op, k, reducedAxes); + *const_cast(tensor), op, k, reduce_axes); TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); nvinfer1::ITensor* output_value_tensor = layer->getOutput(0); @@ -3578,8 +3640,10 @@ static void RegisterValidatableOpConverters( (*registration)["ConcatV2"] = ConvertConcat; (*registration)["Const"] = ConvertConst; (*registration)["Conv2D"] = ConvertConv2D; + (*registration)["Conv2DBackpropInput"] = ConvertConv2DBackpropInput; (*registration)["DepthwiseConv2dNative"] = ConvertConv2DDepthwise; (*registration)["ExpandDims"] = ConvertExpandDims; + (*registration)["LeakyRelu"] = ConvertLeakyRelu; (*registration)["MatMul"] = ConvertMatMul; (*registration)["Pad"] = ConvertPad; (*registration)["Relu6"] = ConvertRelu6; @@ -3588,6 +3652,7 @@ static void RegisterValidatableOpConverters( (*registration)["Squeeze"] = ConvertSqueeze; (*registration)["StridedSlice"] = ConvertStridedSlice; (*registration)["Transpose"] = ConvertTranspose; + (*registration)["TopKV2"] = ConvertTopK; for (auto quantization_op_type : {"QuantizeAndDequantizeV2", "QuantizeAndDequantizeV3", @@ -3634,14 +3699,13 @@ void Converter::RegisterOpConverters() { op_registry_["Mean"] = ConvertReduce; op_registry_["Softmax"] = ConvertSoftmax; op_registry_["BatchMatMul"] = ConvertBatchMatMul; - op_registry_["TopKV2"] = ConvertTopK; plugin_converter_ = ConvertPlugin; } tensorflow::Status ConvertGraphDefToEngine( - const tensorflow::GraphDef& gdef, int precision_mode, int max_batch_size, - size_t max_workspace_size_bytes, + 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, @@ -3656,9 +3720,9 @@ tensorflow::Status ConvertGraphDefToEngine( builder->setMaxBatchSize(max_batch_size); builder->setMaxWorkspaceSize(max_workspace_size_bytes); builder->setGpuAllocator(allocator); - if (precision_mode == FP16MODE) { + if (precision_mode == TrtPrecisionMode::FP16) { builder->setHalf2Mode(true); - } else if (precision_mode == INT8MODE) { + } else if (precision_mode == TrtPrecisionMode::INT8) { builder->setInt8Mode(true); if (use_calibration) { builder->setInt8Calibrator(calibrator); @@ -3678,15 +3742,14 @@ tensorflow::Status ConvertGraphDefToEngine( // Build the network VLOG(1) << "Starting engine conversion "; Converter converter(trt_network.get(), precision_mode, use_calibration); - std::vector> output_tensors; + std::vector output_tensors; // Graph nodes are already topologically sorted during construction for (const auto& node_def : gdef.node()) { string node_name = node_def.name(); VLOG(2) << "Converting op name=" << node_name << ", op=" << node_def.op(); - if (tensorflow::str_util::StartsWith(node_name, kInputPHName) && - (node_def.op() == "Placeholder")) { + if (IsEngineInput(node_name) && (node_def.op() == "Placeholder")) { int32 slot_number = -1; - if (!tensorflow::strings::safe_strto32( + if (!tensorflow::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); @@ -3712,18 +3775,23 @@ tensorflow::Status ConvertGraphDefToEngine( // engines offline, by calling sess.run() and cache/serialize the engines. TF_RETURN_IF_ERROR( converter.AddInputTensor(node_name, trt_dtype, trt_dims, batch_size)); - } else if (tensorflow::str_util::StartsWith(node_name, kOutputPHName) && - (node_def.op() == "Identity")) { + } else if (IsEngineOutput(node_name) && (node_def.op() == "Identity")) { int32 slot_number = -1; - if (!tensorflow::strings::safe_strto32( + if (!tensorflow::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); } + // Get output type that TensorFlow expects + TFAttrs attrs(node_def); + tensorflow::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) { output_tensors.resize(slot_number + 1); } - output_tensors.at(slot_number) = {node_def.input(0), node_name}; + output_tensors.at(slot_number) = {node_def.input(0), node_name, + trt_dtype}; } else { VLOG(2) << "Converting node: " << node_def.name() << " , " << node_def.op(); @@ -3749,8 +3817,7 @@ tensorflow::Status ConvertGraphDefToEngine( tensorflow::Status ConvertSegmentToGraphDef( const tensorflow::Graph* graph, const tensorflow::grappler::GraphProperties& graph_properties, - const std::set& subgraph_node_names, - const std::vector& subgraph_node_ids, // In topological order + const std::vector& subgraph_nodes, // In topological order std::vector* connections, tensorflow::GraphDef* segment_def, string* common_scope) { std::set marker_nodes; @@ -3813,8 +3880,10 @@ tensorflow::Status ConvertSegmentToGraphDef( marker_nodes.insert(node_name); auto seg_node = segment_def->add_node(); tensorflow::NodeDefBuilder builder(node_name, "Identity"); - auto status = builder.Input(connection.inside_node_name, 0, dtype) - .Finalize(seg_node); + auto status = + builder + .Input(connection.inside_node_name, connection.inside_port, dtype) + .Finalize(seg_node); VLOG(1) << "Constructing output " << node_name << " for the edge " << connection.inside_node_name << ":" << connection.inside_port << " -> " << connection.outside_node_name << ":" @@ -3824,13 +3893,12 @@ tensorflow::Status ConvertSegmentToGraphDef( std::unordered_map old_to_new_id_map; // Copy internal nodes to new graphdef - string local_scope = graph->FindNodeId(*subgraph_node_ids.begin())->name(); - for (const auto node_id : subgraph_node_ids) { - const auto node = graph->FindNodeId(node_id); + string local_scope = subgraph_nodes.front()->name(); + for (const Node* node : subgraph_nodes) { local_scope = GetCommonNameScope(local_scope, node->name()); - old_to_new_id_map[node_id] = segment_def->node_size(); + old_to_new_id_map[node->id()] = segment_def->node_size(); auto snode = segment_def->add_node(); - snode->CopyFrom(node->def()); + *snode = node->def(); VLOG(2) << "Copying " << snode->name() << " to subgraph"; } // Update the inputs of the new input nodes to point to placeholder nodes. @@ -3846,6 +3914,11 @@ tensorflow::Status ConvertSegmentToGraphDef( << placeholder_name; snode->set_input(connection.inside_port, placeholder_name); } + std::set subgraph_node_names; + for (const Node* node : subgraph_nodes) { + subgraph_node_names.insert(node->name()); + } + // Remove control inputs that are not inside the segment. for (int i = 0; i < segment_def->node_size(); ++i) { auto snode = segment_def->mutable_node(i); @@ -3856,7 +3929,7 @@ tensorflow::Status ConvertSegmentToGraphDef( TensorId input = ParseTensorName(snode->input(input_idx)); if (!subgraph_node_names.count( string(input.first.data(), input.first.size())) && - !str_util::StartsWith(input.first, kInputPHName)) { + !IsEngineInput(input.first)) { if (input.second == Graph::kControlSlot) { VLOG(1) << "... removing control inputs " << input.first << " from subgraph."; diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.h b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h similarity index 88% rename from tensorflow/contrib/tensorrt/convert/convert_nodes.h rename to tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h index 54e19b73957bccdae2b23bd3556de9ad00b864e5..d1e30eb848bd6ab62719ca6da561d14b05d8537d 100644 --- a/tensorflow/contrib/tensorrt/convert/convert_nodes.h +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_TENSORRT_CONVERT_CONVERT_NODES_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_CONVERT_CONVERT_NODES_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_NODES_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_NODES_H_ #include #include @@ -22,11 +22,11 @@ limitations under the License. #include #include -#include "tensorflow/contrib/tensorrt/convert/utils.h" -#include "tensorflow/contrib/tensorrt/log/trt_logger.h" -#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h" -#include "tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h" -#include "tensorflow/contrib/tensorrt/resources/trt_resources.h" +#include "tensorflow/compiler/tf2tensorrt/convert/utils.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_resources.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/grappler/costs/graph_properties.h" @@ -92,7 +92,7 @@ struct EngineInfo { EngineInfo() : engine_type(EngineType::TRTStatic), max_workspace_size_bytes(0), - precision_mode(FP32MODE), + precision_mode(TrtPrecisionMode::FP32), use_calibration(true) {} string engine_name; @@ -109,7 +109,7 @@ struct EngineInfo { int64 max_workspace_size_bytes; int maximum_cached_engines; std::vector cached_engine_batches; - int precision_mode; + TrtPrecisionMode precision_mode; bool use_calibration; }; @@ -128,8 +128,7 @@ struct EngineInfo { tensorflow::Status ConvertSegmentToGraphDef( const tensorflow::Graph* graph, const tensorflow::grappler::GraphProperties& graph_properties, - const std::set& subgraph_node_names, - const std::vector& subgraph_node_ids, + const std::vector& subgraph_nodes, std::vector* connections, tensorflow::GraphDef* segment_def, string* common_scope); @@ -142,8 +141,8 @@ tensorflow::Status ConvertSegmentToGraphDef( // is successful. This is different than successfully building the engine: // building can still fail afterwards. tensorflow::Status ConvertGraphDefToEngine( - const tensorflow::GraphDef& gdef, int precision_mode, int max_batch_size, - size_t max_workspace_size_bytes, + 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, @@ -159,7 +158,10 @@ class OutputEdgeValidator { bool operator()(const tensorflow::Edge* out_edge) const; }; +string DebugString(const nvinfer1::DimensionType type); +string DebugString(const nvinfer1::DataType trt_dtype); string DebugString(const nvinfer1::Dims& dims); +string DebugString(const nvinfer1::Permutation& permutation, int len); string DebugString(const nvinfer1::ITensor& tensor); int64_t TrtDimsNumElements(const nvinfer1::Dims& dims); @@ -176,6 +178,8 @@ class TRT_ShapedWeights { nvinfer1::Weights GetTrtWeights() const; + // Returns the raw pointer to the underlying buffer which holds the weights + // value. void* GetValues() const { return const_cast(tensor_.tensor_data().data()); } @@ -195,6 +199,10 @@ class TRT_ShapedWeights { // underlying buffer. TRT_ShapedWeights(DataType type, nvinfer1::Dims dims, Tensor tensor); + // All weights should be stored inside TrtWeightStore to make sure lifetime of + // all the underlying tensors are available until the engine is built. For + // this reason, tensor_ should never be reassigned to a different value that + // is not already present in the TrtWeightStore. Tensor tensor_; friend class TrtWeightStore; @@ -394,8 +402,21 @@ class TrtNodeValidator { // Class to convert TF nodes to TRT network. class Converter { public: - Converter(nvinfer1::INetworkDefinition* trt_network, int precision_mode, - bool use_calibration); + // Used for Converter::RenameAndMarkOutputTensors() + struct EngineOutputInfo { + // The TRT tensor name which produces the output. + string source_tensor_name; + // The TensorFlow node name which is receiving the output from the TRT + // engine. This should always be the Identity node created in + // ConvertSegmentToGraphDef. + string dest_node_name; + // Output type. TensorRT requires this to be explicitly set for engine + // outputs. + nvinfer1::DataType trt_dtype; + }; + + Converter(nvinfer1::INetworkDefinition* trt_network, + TrtPrecisionMode precision_mode, bool use_calibration); ////////////////////////////////////////////////////////////////////////////// // Methods used by the TRT engine builder to build a TRT network from a TF @@ -409,13 +430,10 @@ class Converter { Status AddInputTensor(const string& name, nvinfer1::DataType dtype, const nvinfer1::Dims& dims, int batch_size); - // Mark the tensors with names specified by output_tensors[i].first as output - // of the TRT network, and set their names in the TRT network as - // output_tensors[i].second. The tensor names (output_tensors[i].first) are - // standard TF tensor names, i.e. node names followed by output slot number - // (or just the node name if the tensor is the first output of the node). + // Mark the tensors with names specified by source_tensor_name as output of + // the TRT network, and set their names in the TRT network as dest_node_name. Status RenameAndMarkOutputTensors( - const std::vector>& output_tensors); + const std::vector& output_tensors); ////////////////////////////////////////////////////////////////////////////// // Methods used by op converters to convert individual TF node and add layers @@ -426,7 +444,7 @@ class Converter { nvinfer1::INetworkDefinition* network() { return trt_network_; } // What precision are we targeting? - int precision_mode() const { return precision_mode_; } + TrtPrecisionMode precision_mode() const { return precision_mode_; } // Calibration will be or was previously performed on this network? bool use_calibration() const { return use_calibration_; } @@ -469,6 +487,11 @@ class Converter { nvinfer1::Dims* operand_l_new_dims, nvinfer1::Dims* operand_r_new_dims) const; + // Creates an IConstantLayer using 'weights' whose dimensions are specified by + // 'dims', and returns the output ITensor. + nvinfer1::ITensor* CreateConstantLayer(const TRT_ShapedWeights& weights, + const nvinfer1::Dims& dims); + private: // Verify the provided batch_size is consistent with batch_size_ and update it // if necessary. @@ -523,7 +546,7 @@ class Converter { std::vector> quantization_infer_; - const int precision_mode_; + const TrtPrecisionMode precision_mode_; const bool use_calibration_; @@ -544,4 +567,4 @@ class Converter { #endif // GOOGLE_TENSORRT #endif // GOOGLE_CUDA -#endif // TENSORFLOW_CONTRIB_TENSORRT_CONVERT_CONVERT_NODES_H_ +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_NODES_H_ diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc similarity index 73% rename from tensorflow/contrib/tensorrt/convert/convert_nodes_test.cc rename to tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc index a2ddfbffa5b0d8c421bcfe054097a9e42b79fe8f..bb1341ada3766ea322029ec2904e4ae2c6f5544d 100644 --- a/tensorflow/contrib/tensorrt/convert/convert_nodes_test.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/convert/convert_nodes.h" +#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h" #include #include @@ -21,11 +21,16 @@ limitations under the License. #include #include +#include "absl/strings/match.h" +#include "absl/strings/numbers.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/framework/scope.h" #include "tensorflow/cc/ops/standard_ops.h" -#include "tensorflow/contrib/tensorrt/log/trt_logger.h" -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h" #include "tensorflow/core/framework/node_def.pb.h" // NOLINT #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor.pb.h" // NOLINT @@ -35,7 +40,9 @@ limitations under the License. #include "tensorflow/core/grappler/costs/graph_properties.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/protobuf.h" #include "tensorflow/core/platform/test.h" #include "tensorflow/core/protobuf/config.pb.h" // NOLINT #include "tensorflow/core/public/session.h" @@ -50,7 +57,7 @@ namespace tensorflow { namespace tensorrt { namespace convert { -using ::tensorflow::strings::StrCat; +using absl::StrCat; using ::testing::ElementsAre; using ::testing::ElementsAreArray; @@ -152,7 +159,7 @@ void ExpectTrtDimsEqualsArray(const std::vector& lhs, } template -void ExpectArrayNear(const std::vector& lhs, const std::vector& rhs) { +void ExpectArrayNear(const std::vector& lhs, absl::Span rhs) { ASSERT_EQ(lhs.size(), rhs.size()); for (int i = 0; i < lhs.size(); i++) { EXPECT_FLOAT_EQ(lhs[i], rhs[i]); @@ -163,7 +170,7 @@ void ExpectArrayNear(const std::vector& lhs, const std::vector& rhs) { // EXPECT_FLOAT_EQ. template <> void ExpectArrayNear(const std::vector& lhs, - const std::vector& rhs) { + absl::Span rhs) { ASSERT_EQ(lhs.size(), rhs.size()); for (int i = 0; i < lhs.size(); i++) { EXPECT_FLOAT_EQ(Eigen::half_impl::half_to_float(lhs[i]), @@ -364,9 +371,6 @@ TEST(TRT_TensorOrWeights_Test, Basic) { EXPECT_EQ(false, ptr->is_tensor()); EXPECT_EQ(true, ptr->is_weights()); EXPECT_TRUE(TrtShapedWeightsEquals(weights, ptr->weights())); - - nvinfer1::Dims dims; - dims.nbDims = 0; ExpectTrtDimsEqualsArray({}, ptr->GetTrtDims()); } } @@ -481,8 +485,7 @@ class ConverterTest : public ::testing::Test { ConverterTest() { builder_.reset(nvinfer1::createInferBuilder(logger_)); network_.reset(builder_->createNetwork()); - converter_.reset(new Converter(network_.get(), - /*precision_mode=*/FP32MODE, + converter_.reset(new Converter(network_.get(), TrtPrecisionMode::FP32, /*use_calibration=*/false)); weight_store_ = &converter_->weight_store_; } @@ -784,7 +787,7 @@ TEST_F(ConverterTest, MaybeApplyQuantizationRanges) { // input -> infer1 -> infer2 -> infer3 FakeITensor input, infer_1, infer_2, infer_3; FakeITensor not_infer; - Converter int8_converter(/*trt_network=*/nullptr, INT8MODE, + Converter int8_converter(/*trt_network=*/nullptr, TrtPrecisionMode::INT8, /*use_calibration=*/true); int8_converter.ProvideQuantizationRange(&input, -5.0f, 5.0f); int8_converter.ProvideQuantizationRange(¬_infer, -100.0f, 100.0f); @@ -915,6 +918,97 @@ TEST_F(ConverterTest, GetTrtBroadcastShape) { "(tensor #dims 4 vs broadcast #dims 5)"); } +TEST_F(ConverterTest, CreateConstantLayer) { + for (auto dtype : {DT_FLOAT, DT_INT32}) { + TRT_ShapedWeights weights = + weight_store_->GetTempWeights(dtype, GetTestDims({2, 3, 5})); + nvinfer1::ITensor* tensor = + converter_->CreateConstantLayer(weights, GetTestDims({3, 10})); + ASSERT_NE(nullptr, tensor); + EXPECT_EQ(TfDataTypeToTrt(dtype), tensor->getType()) + << "Expected " << DebugString(TfDataTypeToTrt(dtype)) << " vs. actual " + << DebugString(tensor->getType()); + ExpectTrtDimsEqualsArray({3, 10}, tensor->getDimensions()); + } +} + +class ConvertGraphDefToEngineTest : public ::testing::Test { + public: + Status RunConvertGraphDefToEngine(Scope* s) { + GraphDef gdef; + TF_EXPECT_OK(s->ToGraphDef(&gdef)); + std::vector input_shapes; + int batch_size = -1; + for (const NodeDef& node : gdef.node()) { + absl::string_view node_name(node.name()); + if (str_util::ConsumePrefix(&node_name, kInputPHName)) { + int port = -1; + EXPECT_TRUE(absl::SimpleAtoi(node_name, &port)) << node.name(); + if (input_shapes.size() < port + 1) input_shapes.resize(port + 1); + input_shapes[port] = + PartialTensorShape(node.attr().at("shape").shape()); + if (batch_size == -1) { + batch_size = input_shapes[port].dim_size(0); + } else { + EXPECT_EQ(batch_size, input_shapes[port].dim_size(0)); + } + } + } + // TODO(laigd): execute the engine and get outputs. + return ConvertGraphDefToEngine( + gdef, TrtPrecisionMode::FP32, /*max_batch_size=*/1, + /*max_workspace_size_bytes=*/64 << 20, input_shapes, &logger_, + /*allocator=*/nullptr, /*calibrator=*/nullptr, &engine_, + /*use_calibration=*/false, /*convert_successfully=*/nullptr); + } + + protected: + TrtUniquePtrType engine_; + + private: + Logger logger_; +}; + +TEST_F(ConvertGraphDefToEngineTest, IdentityGraph) { + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName(StrCat(kInputPHName, 0)), DT_FLOAT, + ops::Placeholder::Shape({1, 1})); + auto output = ops::Identity(s.WithOpName("identity1"), input); + output = ops::Identity(s.WithOpName("identity2"), output); + output = ops::Identity(s.WithOpName(StrCat(kOutputPHName, 0)), output); + // If the converter marks the input tensor as output tensor, the conversion + // below will fail with: + // > TensorRTOutputPH_0 cannot be both input and output + // > Network must have at least one output + TF_EXPECT_OK(RunConvertGraphDefToEngine(&s)); +} + +// Input/output data format for OpConverterTest::BuildAndRun(). +struct InputOutputData { + void* Buffer() const { + return const_cast(tensor.tensor_data().data()); + } + + size_t TotalBytes() const { return tensor.TotalBytes(); } + + const char* name; + Tensor tensor; +}; + +template +Tensor ConstructTensor(int data_size, const T& value = T()) { + std::vector values(data_size, value); + return test::AsTensor(values); +} + +using DataVec = std::vector; + +template +inline absl::Span GetSpanForData(const InputOutputData& data) { + const auto& tensor_map = data.tensor.flat(); + return absl::Span(tensor_map.data(), tensor_map.size()); +} + // Class to test various op converters, using both a TrtNodeValidator and // Converter. class OpConverterTest : public ::testing::Test { @@ -940,11 +1034,11 @@ class OpConverterTest : public ::testing::Test { builder_.reset(nvinfer1::createInferBuilder(logger_)); network_.reset(builder_->createNetwork()); builder_->setMaxBatchSize(1); + builder_->setMaxWorkspaceSize(1 << 26); // Reset the validator and converter. validator_.reset(new TrtNodeValidator); - converter_.reset(new Converter(network_.get(), - /*precision_mode=*/FP32MODE, + converter_.reset(new Converter(network_.get(), TrtPrecisionMode::FP32, /*use_calibration=*/false)); // Reset other related artifacts. @@ -953,14 +1047,14 @@ class OpConverterTest : public ::testing::Test { } // TODO(laigd): test fp16 and int8 support. - template - void BuildAndRun( - const std::vector>>& - input_data, - const char* output_name, std::vector* output_data) { + void BuildAndRun(const DataVec& input_data, DataVec* output_data) { // Mark the output tensor as TRT engine output. - TF_EXPECT_OK(converter_->RenameAndMarkOutputTensors( - {{string(output_name), string(output_name)}})); + std::vector output_info; + for (const auto& data : *output_data) { + output_info.push_back( + {data.name, data.name, TfDataTypeToTrt(data.tensor.dtype())}); + } + TF_EXPECT_OK(converter_->RenameAndMarkOutputTensors(output_info)); // Build the TRT engine. ASSERT_EQ(nullptr, engine_.get()); @@ -968,31 +1062,44 @@ class OpConverterTest : public ::testing::Test { CHECK_NOTNULL(engine_.get()); // Execute the TRT engine. - ASSERT_LE(input_data.size() + 1, 3); - void* buffers[3]; - for (const auto name_and_data : input_data) { - const int input_size = name_and_data.second.size() * sizeof(T); - const int input_index = engine_->getBindingIndex(name_and_data.first); - ASSERT_EQ(0, cudaMalloc(&buffers[input_index], input_size)); - ASSERT_EQ( - 0, cudaMemcpyAsync(buffers[input_index], name_and_data.second.data(), - input_size, cudaMemcpyHostToDevice, stream_)); + const int num_bindings = input_data.size() + output_data->size(); + std::vector buffers(num_bindings); + + for (const auto& data : input_data) { + const int input_index = engine_->getBindingIndex(data.name); + ASSERT_EQ(0, cudaMalloc(&buffers[input_index], data.TotalBytes())); + ASSERT_EQ(0, cudaMemcpyAsync(buffers[input_index], data.Buffer(), + data.TotalBytes(), cudaMemcpyHostToDevice, + stream_)); + } + struct SizeAndIndex { + SizeAndIndex(int in_size, int in_index) + : size(in_size), index(in_index) {} + int size; + int index; + }; + std::vector output_infos; + for (const auto& data : *output_data) { + const int output_index = engine_->getBindingIndex(data.name); + output_infos.emplace_back(data.TotalBytes(), output_index); + ASSERT_EQ(0, cudaMalloc(&buffers[output_index], data.TotalBytes())); } - const int output_size = output_data->size() * sizeof(T); - const int output_index = engine_->getBindingIndex(output_name); - ASSERT_EQ(0, cudaMalloc(&buffers[output_index], output_size)); - - ASSERT_EQ(engine_->getNbBindings(), input_data.size() + 1); - + ASSERT_EQ(engine_->getNbBindings(), num_bindings); TrtUniquePtrType execution_context( engine_->createExecutionContext()); - execution_context->enqueue(/*batchSize=*/1, buffers, stream_, nullptr); - ASSERT_EQ(0, cudaMemcpyAsync(output_data->data(), buffers[output_index], - output_size, cudaMemcpyDeviceToHost, stream_)); + execution_context->enqueue(/*batchSize=*/1, buffers.data(), stream_, + nullptr); + + for (int i = 0; i < output_infos.size(); ++i) { + const auto& output_info = output_infos[i]; + ASSERT_EQ(0, cudaMemcpyAsync(output_data->at(i).Buffer(), + buffers[output_info.index], output_info.size, + cudaMemcpyDeviceToHost, stream_)); + } cudaStreamSynchronize(stream_); - for (int i = 0; i < input_data.size() + 1; ++i) { + for (int i = 0; i < num_bindings; ++i) { ASSERT_EQ(0, cudaFree(buffers[i])); } } @@ -1111,6 +1218,30 @@ class OpConverterTest : public ::testing::Test { std::unordered_map validator_inputs_; }; +template +void CopyTensorElements(const Tensor& tensor, protobuf::RepeatedField* out) { + out->Clear(); + if (tensor.NumElements() == 0) return; + + // TensorProto does not need to have all the elements present and can truncate + // trailing elements with the same value for compressed representation. Such + // elements are derived based on the tensor shape. + const auto flat = tensor.flat(); + int64 last_index = 0; + for (int64 i = 0; i < tensor.NumElements(); ++i) { + if (flat(i) != flat(last_index)) { + last_index = i; + } + } + + int num_out_elements = last_index + 1; + out->Reserve(num_out_elements); + out->AddNAlreadyReserved(num_out_elements); + const T* src = flat.data(); + T* dst = out->mutable_data(); + std::copy(src, src + num_out_elements, dst); +} + template void TestConvertConst(OpConverterTest* test) { NodeDef node_def; @@ -1123,11 +1254,23 @@ void TestConvertConst(OpConverterTest* test) { const std::vector& expected_value) { test->Reset(); - auto& attr = *node_def.mutable_attr(); + TensorProto* tensor_attr = + (*node_def.mutable_attr())["value"].mutable_tensor(); + tensor_attr->Clear(); + if (as_tensor_content) { - tensor.AsProtoTensorContent(attr["value"].mutable_tensor()); + tensor.AsProtoTensorContent(tensor_attr); } else { - tensor.AsProtoField(attr["value"].mutable_tensor()); + tensor.shape().AsProto(tensor_attr->mutable_tensor_shape()); + tensor_attr->set_dtype(tensor.dtype()); + + if (tensor.dtype() == DT_FLOAT) { + CopyTensorElements(tensor, tensor_attr->mutable_float_val()); + } else if (tensor.dtype() == DT_INT32) { + CopyTensorElements(tensor, tensor_attr->mutable_int_val()); + } else { + tensor.AsProtoField(tensor_attr); + } } test->RunValidationAndConversion(node_def); TRT_TensorOrWeights output; @@ -1140,8 +1283,7 @@ void TestConvertConst(OpConverterTest* test) { { // By default empty tensor will pick DT_FLOAT as data type and we fix it // here. - attr["value"].mutable_tensor()->set_dtype(dtype); - Tensor t; // Empty tensor. + Tensor t(dtype); // Empty tensor. reset_and_test(t, false, {}, {}); } { @@ -1160,6 +1302,22 @@ void TestConvertConst(OpConverterTest* test) { 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})); + 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})); + 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}); + } } TEST_F(OpConverterTest, ConvertConst) { @@ -1189,7 +1347,7 @@ TEST_F(OpConverterTest, ConvertTranspose) { NodeDef node_def = MakeNodeDef("my_transpose", "Transpose", {}); RunValidationAndConversion( node_def, error::INVALID_ARGUMENT, - "Input expects tensor and weights, at my_transpose"); + "Transpose got 0 inputs but expected 2, at my_transpose"); } // Get the NodeDef for Transpose. @@ -1205,8 +1363,8 @@ TEST_F(OpConverterTest, ConvertTranspose) { AddTestTensor("input", {1, 2, 3}); AddTestTensor("weights", {3}); RunValidationAndConversion( - node_def, error::INVALID_ARGUMENT, - "Input expects tensor and weights, at my_transpose"); + node_def, error::UNIMPLEMENTED, + "The input \"perm\" for Transpose must be a constant, at my_transpose"); } { // Transpose at batch dimension, should fail. @@ -1236,10 +1394,12 @@ TEST_F(OpConverterTest, ConvertTranspose) { EXPECT_TRUE(output.is_tensor()); ExpectTrtDimsEqualsArray({3, 1, 2}, output.tensor()->getDimensions()); - std::vector output_data(6); - BuildAndRun({{"input", {1, 2, 3, 4, 5, 6}}}, "my_transpose", - &output_data); - EXPECT_THAT(output_data, ElementsAre(1, 4, 2, 5, 3, 6)); + const DataVec input_data{ + {"input", test::AsTensor({1, 2, 3, 4, 5, 6})}}; + DataVec output_data{{"my_transpose", ConstructTensor(6)}}; + BuildAndRun(input_data, &output_data); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(1, 4, 2, 5, 3, 6)); } } @@ -1249,7 +1409,7 @@ TEST_F(OpConverterTest, ConvertReshape) { NodeDef node_def = MakeNodeDef("my_reshape", "Reshape", {}); RunValidationAndConversion( node_def, error::INVALID_ARGUMENT, - "Input expects weights for shape, at my_reshape"); + "Reshape got 0 inputs but expected 2, at my_reshape"); } // Get the NodeDef for Reshape. @@ -1265,8 +1425,8 @@ TEST_F(OpConverterTest, ConvertReshape) { AddTestTensor("input", {1, 2, 3}); AddTestTensor("weights", {3}); RunValidationAndConversion( - node_def, error::INVALID_ARGUMENT, - "Input expects weights for shape, at my_reshape"); + node_def, error::UNIMPLEMENTED, + "The input \"shape\" for Reshape must be a constant, at my_reshape"); } { // Reshape to scalar, should fail. @@ -1279,11 +1439,6 @@ TEST_F(OpConverterTest, ConvertReshape) { } struct TestParams { - TestParams(int input_batch_size, const std::vector& input_tensor_dims, - const std::vector& input_shape) - : batch_size(input_batch_size), - tensor_dims(input_tensor_dims), - shape(input_shape) {} int batch_size; std::vector tensor_dims; std::vector shape; @@ -1326,10 +1481,12 @@ TEST_F(OpConverterTest, ConvertReshape) { EXPECT_TRUE(output.is_tensor()); ExpectTrtDimsEqualsArray({1, 3, 2}, output.tensor()->getDimensions()); - std::vector output_data(6); - BuildAndRun({{"input", {1, 2, 3, 4, 5, 6}}}, "my_reshape", - &output_data); - EXPECT_THAT(output_data, ElementsAre(1, 2, 3, 4, 5, 6)); + const DataVec input_data{ + {"input", test::AsTensor({1, 2, 3, 4, 5, 6})}}; + DataVec output_data{{"my_reshape", ConstructTensor(6)}}; + BuildAndRun(input_data, &output_data); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(1, 2, 3, 4, 5, 6)); } } @@ -1339,7 +1496,7 @@ TEST_F(OpConverterTest, ConvertMatMul) { NodeDef node_def = MakeNodeDef("my_matmul", "MatMul", {}); RunValidationAndConversion( node_def, error::INVALID_ARGUMENT, - "Input expects tensor and weights, at my_matmul"); + "MatMul got 0 inputs but expected 2, at my_matmul"); } // Get the NodeDef for MatMul. @@ -1389,12 +1546,13 @@ TEST_F(OpConverterTest, ConvertMatMul) { EXPECT_TRUE(output.is_tensor()); ExpectTrtDimsEqualsArray({2}, output.tensor()->getDimensions()); - std::vector output_data(2); - BuildAndRun({{"input", {0, 1}}}, "my_matmul", &output_data); + const DataVec input_data{{"input", test::AsTensor({0, 1})}}; + DataVec output_data{{"my_matmul", ConstructTensor(2)}}; + BuildAndRun(input_data, &output_data); if (transpose_b) { - EXPECT_THAT(output_data, ElementsAre(1, 3)); + EXPECT_THAT(GetSpanForData(output_data[0]), ElementsAre(1, 3)); } else { - EXPECT_THAT(output_data, ElementsAre(2, 3)); + EXPECT_THAT(GetSpanForData(output_data[0]), ElementsAre(2, 3)); } } } @@ -1448,23 +1606,28 @@ void TestConvertBiasAdd(OpConverterTest* test) { const int num_input = TrtDimsNumElements(GetTestDims(dims_array)); ASSERT_EQ(trt_input_rank > 1 ? 6 : (data_format == "NHWC" ? 3 : 2), num_input); - std::vector output_data(num_input); - test->BuildAndRun( - {{"input", std::vector(num_input, CType(0))}}, "my_biasadd", - &output_data); + + const DataVec input_data{ + {"input", ConstructTensor(num_input, CType(0))}}; + DataVec output_data{{"my_biasadd", ConstructTensor(num_input)}}; + test->BuildAndRun(input_data, &output_data); if (trt_input_rank == 1) { if (data_format == "NHWC") { - EXPECT_THAT(output_data, ElementsAre(CType(1), CType(2), CType(3))); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(CType(1), CType(2), CType(3))); } else { - EXPECT_THAT(output_data, ElementsAre(CType(1), CType(2))); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(CType(1), CType(2))); } } else { if (data_format == "NHWC") { - EXPECT_THAT(output_data, ElementsAre(CType(1), CType(2), CType(3), - CType(1), CType(2), CType(3))); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(CType(1), CType(2), CType(3), CType(1), + CType(2), CType(3))); } else { - EXPECT_THAT(output_data, ElementsAre(CType(1), CType(1), CType(1), - CType(2), CType(2), CType(2))); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(CType(1), CType(1), CType(1), CType(2), + CType(2), CType(2))); } } } @@ -1477,7 +1640,7 @@ TEST_F(OpConverterTest, ConvertBiasAdd) { NodeDef node_def = MakeNodeDef("my_biasadd", "BiasAdd", {}); RunValidationAndConversion( node_def, error::INVALID_ARGUMENT, - "Input expects tensor and weights, at my_biasadd"); + "BiasAdd got 0 inputs but expected 2, at my_biasadd"); } // OK. Note that kINT32 is not supported by IScaleLayer, so we don't test @@ -1542,21 +1705,25 @@ void TestBinaryTensorOpWeightNoBroadcast(OpConverterTest* test) { EXPECT_TRUE(output.is_tensor()); ExpectTrtDimsEqualsArray({1, 1, 2}, output.tensor()->getDimensions()); - std::vector output_data(2); - test->BuildAndRun( - {{"input", - /*input_data=*/swap_inputs ? operand2 : operand1}}, - "my_binary", &output_data); + const DataVec input_data{ + {"input", test::AsTensor(swap_inputs ? operand2 : operand1)}}; + DataVec output_data{{"my_binary", ConstructTensor(2)}}; + test->BuildAndRun(input_data, &output_data); if (node_def.op() == "Add") { - EXPECT_THAT(output_data, ElementsAre(CType(5), CType(10.5))); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(CType(5), CType(10.5))); } else if (node_def.op() == "Sub") { - EXPECT_THAT(output_data, ElementsAre(CType(1), CType(4.5))); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(CType(1), CType(4.5))); } else if (node_def.op() == "Mul") { - EXPECT_THAT(output_data, ElementsAre(CType(6), CType(22.5))); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(CType(6), CType(22.5))); } else if (node_def.op() == "Div") { - EXPECT_THAT(output_data, ElementsAre(CType(1.5), CType(2.5))); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(CType(1.5), CType(2.5))); } else if (node_def.op() == "RealDiv") { - EXPECT_THAT(output_data, ElementsAre(CType(1.5), CType(2.5))); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(CType(1.5), CType(2.5))); } else { ASSERT_TRUE(false); } @@ -1591,13 +1758,14 @@ void TestBinaryTensorOpWeightWithChannelWiseBroadcast(OpConverterTest* test) { EXPECT_TRUE(output.is_tensor()); ExpectTrtDimsEqualsArray({2, 1, 2}, output.tensor()->getDimensions()); - std::vector output_data(4); - test->BuildAndRun({{"input", input}}, "my_binary", &output_data); + const DataVec input_data{{"input", test::AsTensor(input)}}; + DataVec output_data{{"my_binary", ConstructTensor(4)}}; + test->BuildAndRun(input_data, &output_data); if (weights_dims.size() == 1) { - EXPECT_THAT(output_data, + EXPECT_THAT(GetSpanForData(output_data[0]), ElementsAre(CType(11), CType(22), CType(13), CType(24))); } else { - EXPECT_THAT(output_data, + EXPECT_THAT(GetSpanForData(output_data[0]), ElementsAre(CType(11), CType(12), CType(23), CType(24))); } } @@ -1625,9 +1793,10 @@ void TestBinaryTensorOpWeightWithUniformlyBroadcast(OpConverterTest* test) { EXPECT_TRUE(output.is_tensor()); ExpectTrtDimsEqualsArray({2, 1, 2}, output.tensor()->getDimensions()); - std::vector output_data(4); - test->BuildAndRun({{"input", input}}, "my_binary", &output_data); - EXPECT_THAT(output_data, + const DataVec input_data{{"input", test::AsTensor(input)}}; + DataVec output_data{{"my_binary", ConstructTensor(4)}}; + test->BuildAndRun(input_data, &output_data); + EXPECT_THAT(GetSpanForData(output_data[0]), ElementsAre(CType(11), CType(12), CType(13), CType(14))); } @@ -1675,17 +1844,19 @@ void TestBinaryTensorOpWeightFallback(OpConverterTest* test, // Check the result of running the engine. const int expected_num_outputs = TrtDimsNumElements(GetTestDims(expected_output_dims)); - std::vector output_data(expected_num_outputs); - test->BuildAndRun( - {{"input", - /*input_data=*/std::vector(num_inputs, CType(2))}}, - "my_binary", &output_data); + const DataVec input_data{ + {"input", ConstructTensor(num_inputs, CType(2))}}; + DataVec output_data{ + {"my_binary", ConstructTensor(expected_num_outputs)}}; + test->BuildAndRun(input_data, &output_data); if (node_def.op() == "Add") { - EXPECT_THAT(output_data, ElementsAreArray(std::vector( - expected_num_outputs, CType(3)))); + EXPECT_THAT( + GetSpanForData(output_data[0]), + ElementsAreArray(std::vector(expected_num_outputs, CType(3)))); } else if (node_def.op() == "Minimum") { - EXPECT_THAT(output_data, ElementsAreArray(std::vector( - expected_num_outputs, CType(1)))); + EXPECT_THAT( + GetSpanForData(output_data[0]), + ElementsAreArray(std::vector(expected_num_outputs, CType(1)))); } else { ASSERT_TRUE(false); } @@ -1712,32 +1883,33 @@ void TestBinaryTensorOpTensor(OpConverterTest* test) { EXPECT_TRUE(output.is_tensor()); ExpectTrtDimsEqualsArray({2, 2}, output.tensor()->getDimensions()); - std::vector output_data(4); + const DataVec input_data{ + {"input1", test::AsTensor({CType(3), CType(6)})}, + {"input2", test::AsTensor({CType(2), CType(3)})}}; + DataVec output_data{{"my_binary", ConstructTensor(4)}}; // After broadcasting first input becomes {3, 6, 3, 6} and second input // becomes {2, 3, 2, 3}. - test->BuildAndRun( - {{"input1", {CType(3), CType(6)}}, {"input2", {CType(2), CType(3)}}}, - "my_binary", &output_data); + test->BuildAndRun(input_data, &output_data); if (node_def.op() == "Add") { - EXPECT_THAT(output_data, + EXPECT_THAT(GetSpanForData(output_data[0]), ElementsAre(CType(5), CType(8), CType(6), CType(9))); } else if (node_def.op() == "Sub") { - EXPECT_THAT(output_data, + EXPECT_THAT(GetSpanForData(output_data[0]), ElementsAre(CType(1), CType(4), CType(0), CType(3))); } else if (node_def.op() == "Mul") { - EXPECT_THAT(output_data, + EXPECT_THAT(GetSpanForData(output_data[0]), ElementsAre(CType(6), CType(12), CType(9), CType(18))); } else if (node_def.op() == "Div") { - EXPECT_THAT(output_data, + EXPECT_THAT(GetSpanForData(output_data[0]), ElementsAre(CType(1.5), CType(3), CType(1), CType(2))); } else if (node_def.op() == "RealDiv") { - EXPECT_THAT(output_data, + EXPECT_THAT(GetSpanForData(output_data[0]), ElementsAre(CType(1.5), CType(3), CType(1), CType(2))); } else if (node_def.op() == "Minimum") { - EXPECT_THAT(output_data, + EXPECT_THAT(GetSpanForData(output_data[0]), ElementsAre(CType(2), CType(2), CType(3), CType(3))); } else if (node_def.op() == "Maximum") { - EXPECT_THAT(output_data, + EXPECT_THAT(GetSpanForData(output_data[0]), ElementsAre(CType(3), CType(6), CType(3), CType(6))); } else { ASSERT_TRUE(false); @@ -1751,7 +1923,9 @@ TEST_F(OpConverterTest, ConvertBinary) { NodeDef node_def = MakeNodeDef("my_add", "Add", {num_inputs, "input"}); AddTestTensor("input", {1}, /*batch_size=*/1, nvinfer1::DataType::kFLOAT); RunValidationAndConversion(node_def, error::INVALID_ARGUMENT, - "Binary ops require two inputs, at my_add"); + StrCat("Add got ", std::to_string(num_inputs), + " inputs but expected 2, at my_add") + .c_str()); } { // Both inputs are weights. @@ -1821,14 +1995,18 @@ TEST_F(OpConverterTest, ConvertBinary) { } TEST_F(OpConverterTest, ConvertQuantize) { - for (const string& op : - {"FakeQuantWithMinMaxArgs", "FakeQuantWithMinMaxVars", - "QuantizeAndDequantizeV2", "QuantizeAndDequantizeV3"}) { + const std::pair op_with_num_inputs[4] = { + {"FakeQuantWithMinMaxArgs", 1}, + {"FakeQuantWithMinMaxVars", 3}, + {"QuantizeAndDequantizeV2", 3}, + {"QuantizeAndDequantizeV3", 4}}; + for (const auto& pair : op_with_num_inputs) { // Input list is empty, should fail. - NodeDef node_def = MakeNodeDef("my_quantize", op, {}); + NodeDef node_def = MakeNodeDef("my_quantize", pair.first, {}); RunValidationAndConversion( node_def, error::INVALID_ARGUMENT, - StrCat("Invalid number of inputs for ", op, ", at my_quantize") + StrCat(pair.first, " got 0 inputs but expected ", + std::to_string(pair.second), ", at my_quantize") .c_str()); } { @@ -1915,9 +2093,9 @@ TEST_F(OpConverterTest, ConvertQuantize) { AddTestTensor("weights_min", {1}); AddTestTensor("weights_max", {1}); RunValidationAndConversion( - node_def, error::INVALID_ARGUMENT, - "Min and max inputs for QuantizeAndDequantizeV2 must be weights not " - "tensors, at my_quantize"); + node_def, error::UNIMPLEMENTED, + "The input \"input_min\" for QuantizeAndDequantizeV2 must be a constant" + ", at my_quantize"); } { // QuantizeAndDequantizeV3 ranges set via inputs, ok. @@ -1950,7 +2128,7 @@ TEST_F(OpConverterTest, ConvertRelu6) { NodeDef node_def = MakeNodeDef("my_relu6", "Relu6", {}); RunValidationAndConversion( node_def, error::INVALID_ARGUMENT, - "Invalid number of inputs for Relu6, at my_relu6"); + "Relu6 got 0 inputs but expected 1, at my_relu6"); } // Get the NodeDef for Relu6. @@ -1964,7 +2142,7 @@ TEST_F(OpConverterTest, ConvertRelu6) { AddTestWeights("input", {1}, {1.0f}); RunValidationAndConversion( node_def, error::UNIMPLEMENTED, - "Relu6 is only implemented for tensors, not weights, at my_relu6"); + "The input \"input\" for Relu6 must be a tensor, at my_relu6"); } { // Clip tensor values and set quantization ranges, ok. @@ -1977,10 +2155,12 @@ TEST_F(OpConverterTest, ConvertRelu6) { auto ranges = quantization_ranges(); EXPECT_EQ(ranges[output.tensor()], 6.0f); - std::vector output_data(6); - BuildAndRun({{"input", {-100, -1, 0, 3, 5, 9}}}, "my_relu6", - &output_data); - EXPECT_THAT(output_data, ElementsAre(0, 0, 0, 3, 5, 6)); + 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)); } } @@ -2002,24 +2182,26 @@ void TestConvertSquare(OpConverterTest* test) { ExpectTrtDimsEqualsArray({1, 20}, output.tensor()->getDimensions()); const int num_inputs = 20; - std::vector input_data(num_inputs); - std::vector expected_output_data(num_inputs); + std::vector inputs(num_inputs); + std::vector expected_outputs(num_inputs); for (int i = 0; i < 20; i++) { const CType value = CType(i - 9); - input_data[i] = value; - expected_output_data[i] = value * value; + inputs[i] = value; + expected_outputs[i] = value * value; } - std::vector output_data(num_inputs); - test->BuildAndRun({{"input", input_data}}, "my_square", &output_data); - ExpectArrayNear(expected_output_data, output_data); + const DataVec input_data{{"input", test::AsTensor(inputs)}}; + DataVec output_data{{"my_square", ConstructTensor(num_inputs)}}; + test->BuildAndRun(input_data, &output_data); + ExpectArrayNear(expected_outputs, GetSpanForData(output_data[0])); } TEST_F(OpConverterTest, ConvertSquare) { { // Input list is empty, should fail. NodeDef node_def = MakeNodeDef("my_square", "Square", {}); - RunValidationAndConversion(node_def, error::INVALID_ARGUMENT, - "Square expects one input, at my_square"); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Square got 0 inputs but expected 1, at my_square"); } { // Input is weights, should fail. @@ -2031,7 +2213,7 @@ TEST_F(OpConverterTest, ConvertSquare) { AddTestWeights("input", {1, 2, 3}, {1, 2, 3, 4, -5, 6}); RunValidationAndConversion( node_def, error::UNIMPLEMENTED, - "Square is only implemented for tensors, at my_square"); + "The input \"x\" for Square must be a tensor, at my_square"); } // OK. Note that kINT32 is not supported by IElementWiseLayer, so we don't @@ -2047,7 +2229,7 @@ TEST_F(OpConverterTest, ConvertActivation) { // Input list is empty, should fail. NodeDef node_def = MakeNodeDef("my_act", "Relu", {}); RunValidationAndConversion(node_def, error::INVALID_ARGUMENT, - "Relu expects one input, at my_act"); + "Relu got 0 inputs but expected 1, at my_act"); } { // Input is weights, should fail. @@ -2059,7 +2241,7 @@ TEST_F(OpConverterTest, ConvertActivation) { AddTestWeights("input", {1, 2, 3}, {-3, -2, -1, 0, 1, 2}); RunValidationAndConversion( node_def, error::UNIMPLEMENTED, - "Relu is only implemented for tensors, at my_act"); + "The input \"input\" for Relu must be a tensor, at my_act"); } // Get nodedef for activation layer. @@ -2103,12 +2285,14 @@ TEST_F(OpConverterTest, ConvertActivation) { EXPECT_TRUE(output.is_tensor()); ExpectTrtDimsEqualsArray({1, 2, 3}, output.tensor()->getDimensions()); - const std::vector input_data = {-100, -2, -1, 0, 1, 100}; - std::vector output_data(6); - BuildAndRun({{"input", input_data}}, "my_act", &output_data); - for (int i = 0; i < input_data.size(); i++) { - const float expected_output = get_act_output(op_name, input_data[i]); - EXPECT_FLOAT_EQ(output_data[i], expected_output); + const std::vector input = {-100, -2, -1, 0, 1, 100}; + const DataVec input_data{{"input", test::AsTensor(input)}}; + DataVec output_data{{"my_act", ConstructTensor(6)}}; + BuildAndRun(input_data, &output_data); + for (int i = 0; i < input.size(); i++) { + const float expected_output = get_act_output(op_name, input[i]); + EXPECT_FLOAT_EQ(GetSpanForData(output_data[0])[i], + expected_output); } } } @@ -2119,7 +2303,7 @@ TEST_F(OpConverterTest, ConvertExpandDims) { NodeDef node_def = MakeNodeDef("my_expanddims", "ExpandDims", {}); RunValidationAndConversion( node_def, error::INVALID_ARGUMENT, - "Two inputs expected for ExpandDims, at my_expanddims"); + "ExpandDims got 0 inputs but expected 2, at my_expanddims"); } // Get the NodeDef for ExpandDims. @@ -2134,18 +2318,18 @@ TEST_F(OpConverterTest, ConvertExpandDims) { Reset(); AddTestWeights("input", {1, 2, 3}, {1, 2, 3, 4, 5, 6}); AddTestWeights("weights", {1}, {1}); - RunValidationAndConversion( - node_def, error::UNIMPLEMENTED, - "ExpandDims expects tensor for input, at my_expanddims"); + RunValidationAndConversion(node_def, error::UNIMPLEMENTED, + "The input \"input\" for ExpandDims must be a " + "tensor, at my_expanddims"); } { // Axis is a tensor, should fail. Reset(); AddTestTensor("input", {1, 2, 3}); AddTestTensor("weights", {3}); - RunValidationAndConversion( - node_def, error::INVALID_ARGUMENT, - "ExpandDims expects weights for axis, at my_expanddims"); + RunValidationAndConversion(node_def, error::UNIMPLEMENTED, + "The input \"axis\" for ExpandDims must be a " + "constant, at my_expanddims"); } { // Add dim at batch dimension, should fail. @@ -2192,11 +2376,6 @@ TEST_F(OpConverterTest, ConvertExpandDims) { } struct TestParams { - TestParams(const std::vector& input_dims, int axis, - const std::vector& expected_output_dims) - : input_dims(input_dims), - axis(axis), - expected_output_dims(expected_output_dims) {} std::vector input_dims; int axis; std::vector expected_output_dims; @@ -2221,10 +2400,12 @@ TEST_F(OpConverterTest, ConvertExpandDims) { ExpectTrtDimsEqualsArray(ok_params[i].expected_output_dims, output.tensor()->getDimensions()); - std::vector output_data(6); - BuildAndRun({{"input", {1, 2, 3, 4, 5, 6}}}, "my_expanddims", - &output_data); - EXPECT_THAT(output_data, ElementsAre(1, 2, 3, 4, 5, 6)); + const DataVec input_data{ + {"input", test::AsTensor({1, 2, 3, 4, 5, 6})}}; + DataVec output_data{{"my_expanddims", ConstructTensor(6)}}; + BuildAndRun(input_data, &output_data); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(1, 2, 3, 4, 5, 6)); } } @@ -2232,8 +2413,9 @@ TEST_F(OpConverterTest, ConvertSqueeze) { { // Input list is empty, should fail. NodeDef node_def = MakeNodeDef("my_squeeze", "Squeeze", {}); - RunValidationAndConversion(node_def, error::INVALID_ARGUMENT, - "One input expected for Squeeze, at my_squeeze"); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Squeeze got 0 inputs but expected 1, at my_squeeze"); } { // No attrs, should fail. @@ -2253,7 +2435,7 @@ TEST_F(OpConverterTest, ConvertSqueeze) { Scope s = Scope::NewRootScope(); auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); ops::Squeeze::Attrs squeeze_attrs; - squeeze_attrs.axis_ = gtl::ArraySlice(axis); + squeeze_attrs.axis_ = gtl::ArraySlice(axis); // non-absl ok auto squeeze = ops::Squeeze(s.WithOpName("my_squeeze"), input, squeeze_attrs); return squeeze.operation.node()->def(); @@ -2266,7 +2448,7 @@ TEST_F(OpConverterTest, ConvertSqueeze) { AddTestWeights("input", {1, 2, 3}, {1, 2, 3, 4, 5, 6}); RunValidationAndConversion( node_def, error::UNIMPLEMENTED, - "Squeeze expects tensor for input, at my_squeeze"); + "The input \"input\" for Squeeze must be a tensor, at my_squeeze"); } { // Squeeze batch dim, should fail. @@ -2306,11 +2488,6 @@ TEST_F(OpConverterTest, ConvertSqueeze) { } struct TestParams { - TestParams(const std::vector& input_dims, const std::vector& axis, - const std::vector& expected_output_dims) - : input_dims(input_dims), - axis(axis), - expected_output_dims(expected_output_dims) {} std::vector input_dims; std::vector axis; std::vector expected_output_dims; @@ -2341,10 +2518,12 @@ TEST_F(OpConverterTest, ConvertSqueeze) { ExpectTrtDimsEqualsArray(ok_params[i].expected_output_dims, output.tensor()->getDimensions()); - std::vector output_data(6); - BuildAndRun({{"input", {1, 2, 3, 4, 5, 6}}}, "my_squeeze", - &output_data); - EXPECT_THAT(output_data, ElementsAre(1, 2, 3, 4, 5, 6)); + const DataVec input_data{ + {"input", test::AsTensor({1, 2, 3, 4, 5, 6})}}; + DataVec output_data{{"my_squeeze", ConstructTensor(6)}}; + BuildAndRun(input_data, &output_data); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(1, 2, 3, 4, 5, 6)); } } @@ -2354,7 +2533,7 @@ TEST_F(OpConverterTest, ConvertStridedSlice) { NodeDef node_def = MakeNodeDef("my_strided_slice", "StridedSlice", {}); RunValidationAndConversion( node_def, error::INVALID_ARGUMENT, - "StridedSlice expects 4 inputs, at my_strided_slice"); + "StridedSlice got 0 inputs but expected 4, at my_strided_slice"); } // Get nodedef for StridedSlice layer. @@ -2378,14 +2557,16 @@ TEST_F(OpConverterTest, ConvertStridedSlice) { }; { + // Input is weights, should fail. + Reset(); NodeDef node_def = get_strided_slice_nodedef(); AddTestWeights("input", {1, 2, 3}, {1, 2, 3, 4, 5, 6}); AddTestWeights("begin", {4}, {0, 0, 0, 0}); AddTestWeights("end", {4}, {1, 1, 2, 3}); AddTestWeights("strides", {4}, {1, 1, 1, 1}); - RunValidationAndConversion( - node_def, error::UNIMPLEMENTED, - "StridedSlice is only implemented for tensors, at my_strided_slice"); + RunValidationAndConversion(node_def, error::UNIMPLEMENTED, + "The input \"input\" for StridedSlice must be a " + "tensor, at my_strided_slice"); } { // Begin, end, strides are tensors, should fail. @@ -2396,8 +2577,8 @@ TEST_F(OpConverterTest, ConvertStridedSlice) { AddTestTensor("end", {4}); AddTestTensor("strides", {4}); RunValidationAndConversion( - node_def, error::INVALID_ARGUMENT, - "StridedSlice expects weights for begin, end, and strides, at " + node_def, error::UNIMPLEMENTED, + "The input \"begin\" for StridedSlice must be a constant, at " "my_strided_slice"); } { @@ -2479,29 +2660,6 @@ TEST_F(OpConverterTest, ConvertStridedSlice) { } struct TestParams { - TestParams(const std::vector& input_dims, - const std::vector& expected_output_dims, - const std::vector& begin, const std::vector& end, - const std::vector& begin_mask, - const std::vector& end_mask, - const std::vector& expected_output) - : input_dims(input_dims), - expected_output_dims(expected_output_dims), - begin(begin), - end(end), - expected_output(expected_output) { - // Masks are provided in terms of vectors for readability. Convert them to - // binary here. - this->begin_mask = 0; - for (int i = 0; i < begin_mask.size(); i++) { - if (begin_mask[i]) this->begin_mask |= (1 << i); - } - this->end_mask = 0; - for (int i = 0; i < end_mask.size(); i++) { - if (end_mask[i]) this->end_mask |= (1 << i); - } - } - std::vector input_dims; std::vector expected_output_dims; std::vector begin; @@ -2511,87 +2669,112 @@ TEST_F(OpConverterTest, ConvertStridedSlice) { std::vector expected_output; }; + auto get_mask = [](const std::vector& mask) { + int result = 0; + for (int i = 0; i < mask.size(); i++) { + if (mask[i]) result += (1 << i); + } + return result; + }; + // Ok. const int kStridedSliceOKCases = 18; TestParams ok_params[kStridedSliceOKCases] = { // 2D Crop. TestParams{/*input_dims=*/{1, 2, 3}, /*expected_output_dims=*/{1, 1, 2}, /*begin=*/{0, 0, 0, 0}, /*end=*/{0, 0, 1, 2}, - /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 1, 0, 0}, + /*begin_mask=*/get_mask({0, 0, 0, 0}), + /*end_mask=*/get_mask({1, 1, 0, 0}), /*expected_output=*/{1, 2}}, TestParams{/*input_dims=*/{1, 2, 3}, /*expected_output_dims=*/{1, 1, 2}, /*begin=*/{0, 0, 1, 1}, /*end=*/{0, 0, 0, 0}, - /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 1, 1, 1}, + /*begin_mask=*/get_mask({0, 0, 0, 0}), + /*end_mask=*/get_mask({1, 1, 1, 1}), /*expected_output=*/{5, 6}}, TestParams{/*input_dims=*/{1, 2, 3}, /*expected_output_dims=*/{1, 1, 2}, /*begin=*/{0, 0, 1, 1}, /*end=*/{0, 1, 2, 3}, - /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 1, 0, 0}, + /*begin_mask=*/get_mask({0, 0, 0, 0}), + /*end_mask=*/get_mask({1, 1, 0, 0}), /*expected_output=*/{5, 6}}, // 2D Crop, with transpose. TestParams{/*input_dims=*/{2, 3, 1}, /*expected_output_dims=*/{1, 2, 1}, /*begin=*/{0, 0, 0, 0}, /*end=*/{0, 1, 2, 1}, - /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 0, 0, 0}, + /*begin_mask=*/get_mask({0, 0, 0, 0}), + /*end_mask=*/get_mask({1, 0, 0, 0}), /*expected_output=*/{1, 2}}, TestParams{/*input_dims=*/{2, 3, 1}, /*expected_output_dims=*/{1, 2, 1}, /*begin=*/{0, 1, 1, 0}, /*end=*/{0, 2, 3, 1}, - /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 0, 0, 0}, + /*begin_mask=*/get_mask({0, 0, 0, 0}), + /*end_mask=*/get_mask({1, 0, 0, 0}), /*expected_output=*/{5, 6}}, TestParams{/*input_dims=*/{2, 1, 3}, /*expected_output_dims=*/{1, 1, 2}, /*begin=*/{0, 0, 0, 0}, /*end=*/{0, 1, 1, 2}, - /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 0, 0, 0}, + /*begin_mask=*/get_mask({0, 0, 0, 0}), + /*end_mask=*/get_mask({1, 0, 0, 0}), /*expected_output=*/{1, 2}}, TestParams{/*input_dims=*/{2, 1, 3}, /*expected_output_dims=*/{1, 1, 2}, /*begin=*/{0, 1, 0, 1}, /*end=*/{0, 2, 1, 3}, - /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 0, 0, 0}, + /*begin_mask=*/get_mask({0, 0, 0, 0}), + /*end_mask=*/get_mask({1, 0, 0, 0}), /*expected_output=*/{5, 6}}, // 2D Crop, with reshape. TestParams{/*input_dims=*/{2, 3}, /*expected_output_dims=*/{1, 2}, /*begin=*/{0, 0, 0}, /*end=*/{0, 1, 2}, - /*begin_mask=*/{0, 0, 0}, /*end_mask=*/{1, 0, 0}, + /*begin_mask=*/get_mask({0, 0, 0}), + /*end_mask=*/get_mask({1, 0, 0}), /*expected_output=*/{1, 2}}, TestParams{/*input_dims=*/{2, 3}, /*expected_output_dims=*/{1, 2}, /*begin=*/{0, 1, 1}, /*end=*/{0, 0, 0}, - /*begin_mask=*/{0, 0, 0}, /*end_mask=*/{1, 1, 1}, + /*begin_mask=*/get_mask({0, 0, 0}), + /*end_mask=*/get_mask({1, 1, 1}), /*expected_output=*/{5, 6}}, // 1D Crop. TestParams{/*input_dims=*/{1, 2, 3}, /*expected_output_dims=*/{1, 2, 2}, /*begin=*/{0, 0, 0, 0}, /*end=*/{0, 0, 0, 2}, - /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 1, 1, 0}, + /*begin_mask=*/get_mask({0, 0, 0, 0}), + /*end_mask=*/get_mask({1, 1, 1, 0}), /*expected_output=*/{1, 2, 4, 5}}, TestParams{/*input_dims=*/{1, 2, 3}, /*expected_output_dims=*/{1, 1, 3}, /*begin=*/{0, 0, 1, 0}, /*end=*/{0, 0, 0, 0}, - /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 1, 1, 1}, + /*begin_mask=*/get_mask({0, 0, 0, 0}), + /*end_mask=*/get_mask({1, 1, 1, 1}), /*expected_output=*/{4, 5, 6}}, // 1D Crop, with transpose. TestParams{/*input_dims=*/{2, 3, 1}, /*expected_output_dims=*/{1, 3, 1}, /*begin=*/{0, 0, 0, 0}, /*end=*/{0, 1, 0, 0}, - /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 0, 1, 1}, + /*begin_mask=*/get_mask({0, 0, 0, 0}), + /*end_mask=*/get_mask({1, 0, 1, 1}), /*expected_output=*/{1, 2, 3}}, TestParams{/*input_dims=*/{2, 3, 1}, /*expected_output_dims=*/{1, 3, 1}, /*begin=*/{0, 1, 0, 0}, /*end=*/{0, 0, 0, 0}, - /*begin_mask=*/{0, 0, 0, 0}, /*end_mask=*/{1, 1, 1, 1}, + /*begin_mask=*/get_mask({0, 0, 0, 0}), + /*end_mask=*/get_mask({1, 1, 1, 1}), /*expected_output=*/{4, 5, 6}}, // 1D Crop, with reshape. TestParams{/*input_dims=*/{6}, /*expected_output_dims=*/{3}, /*begin=*/{0, 0}, /*end=*/{0, 3}, - /*begin_mask=*/{0, 0}, /*end_mask=*/{1, 0}, + /*begin_mask=*/get_mask({0, 0}), /*end_mask=*/get_mask({1, 0}), /*expected_output=*/{1, 2, 3}}, TestParams{/*input_dims=*/{1, 6}, /*expected_output_dims=*/{1, 3}, /*begin=*/{0, 0, 2}, /*end=*/{0, 0, 5}, - /*begin_mask=*/{0, 0, 0}, /*end_mask=*/{1, 1, 0}, + /*begin_mask=*/get_mask({0, 0, 0}), + /*end_mask=*/get_mask({1, 1, 0}), /*expected_output=*/{3, 4, 5}}, TestParams{/*input_dims=*/{6, 1}, /*expected_output_dims=*/{3, 1}, /*begin=*/{0, 2, 0}, /*end=*/{0, 5, 0}, - /*begin_mask=*/{0, 0, 0}, /*end_mask=*/{1, 0, 1}, + /*begin_mask=*/get_mask({0, 0, 0}), + /*end_mask=*/get_mask({1, 0, 1}), /*expected_output=*/{3, 4, 5}}, // Negative axis. TestParams{/*input_dims=*/{6, 1}, /*expected_output_dims=*/{3, 1}, /*begin=*/{0, -6, 0}, /*end=*/{0, -3, 0}, - /*begin_mask=*/{0, 0, 0}, /*end_mask=*/{1, 0, 1}, + /*begin_mask=*/get_mask({0, 0, 0}), + /*end_mask=*/get_mask({1, 0, 1}), /*expected_output=*/{1, 2, 3}}, TestParams{/*input_dims=*/{6, 1}, /*expected_output_dims=*/{5, 1}, /*begin=*/{0, 0, 0}, /*end=*/{0, -1, 0}, - /*begin_mask=*/{0, 0, 0}, /*end_mask=*/{1, 0, 1}, + /*begin_mask=*/get_mask({0, 0, 0}), + /*end_mask=*/get_mask({1, 0, 1}), /*expected_output=*/{1, 2, 3, 4, 5}}, }; @@ -2612,10 +2795,337 @@ TEST_F(OpConverterTest, ConvertStridedSlice) { TRT_TensorOrWeights output; TF_EXPECT_OK(GetTensorOrWeights("my_strided_slice", &output)); - std::vector output_data(ok_params[i].expected_output.size()); - BuildAndRun({{"input", {1, 2, 3, 4, 5, 6}}}, "my_strided_slice", - &output_data); - EXPECT_THAT(output_data, ElementsAreArray(ok_params[i].expected_output)); + + const DataVec input_data{ + {"input", test::AsTensor({1, 2, 3, 4, 5, 6})}}; + DataVec output_data{ + {"my_strided_slice", + ConstructTensor(ok_params[i].expected_output.size())}}; + BuildAndRun(input_data, &output_data); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAreArray(ok_params[i].expected_output)); + } +} + +TEST_F(OpConverterTest, ConvertConv2D) { + { + // Input list is empty, should fail. + NodeDef node_def = MakeNodeDef("my_conv2d", "Conv2D", {}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Conv2D got 0 inputs but expected 2, at my_conv2d"); + } + + // Get nodedef for Conv2D layer. + auto get_conv2d_nodedef = + [](std::vector strides = {1, 1, 1, 1}, string padding = "SAME", + string data_format = "NCHW", std::vector dilations = {1, 1, 1, 1}, + bool is_conv2d_backprop_input = false) -> NodeDef { + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto filter = ops::Placeholder(s.WithOpName("weights"), DT_FLOAT); + if (is_conv2d_backprop_input) { + auto input_sizes = + ops::Placeholder(s.WithOpName("input_sizes"), DT_INT32); + ops::Conv2DBackpropInput::Attrs attrs = ops::Conv2DBackpropInput::Attrs() + .DataFormat(data_format) + .Dilations(dilations); + auto conv2d = + ops::Conv2DBackpropInput(s.WithOpName("my_conv2d"), input_sizes, + filter, input, strides, padding, attrs); + return conv2d.operation.node()->def(); + } else { + ops::Conv2D::Attrs attrs = + ops::Conv2D::Attrs().DataFormat(data_format).Dilations(dilations); + auto conv2d = ops::Conv2D(s.WithOpName("my_conv2d"), input, filter, + strides, padding, attrs); + return conv2d.operation.node()->def(); + } + }; + + { + // Input is weights, should fail. + Reset(); + NodeDef node_def = get_conv2d_nodedef(); + AddTestWeights("input", {1, 2, 3}, {1, 2, 3, 4, 5, 6}); + AddTestWeights("weights", {3, 3, 1, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9}); + RunValidationAndConversion( + node_def, error::UNIMPLEMENTED, + "The input \"input\" for Conv2D must be a tensor, at my_conv2d"); + } + { + // Filter is tensor, should fail. + Reset(); + NodeDef node_def = get_conv2d_nodedef(); + AddTestTensor("input", {1, 2, 3}); + AddTestTensor("weights", {3, 3, 1, 1}); + RunValidationAndConversion( + node_def, error::UNIMPLEMENTED, + "The input \"filter\" for Conv2D must be a constant, at my_conv2d"); + } + { + // Filter is not 4D, should fail. + Reset(); + NodeDef node_def = get_conv2d_nodedef(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("weights", {3, 3, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Conv2D expects kernel of dimension 4, at my_conv2d"); + } + { + // Dilations is not 4D, should fail. + Reset(); + NodeDef node_def = + get_conv2d_nodedef({1, 1, 1, 1}, "SAME", "NCHW", {1, 1, 1}); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("weights", {3, 3, 1, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Convolution dilations field must specify 4 dimensions, at my_conv2d"); + } + { + // Dilation value is not 1 for channel, should fail. + Reset(); + NodeDef node_def = + get_conv2d_nodedef({1, 1, 1, 1}, "SAME", "NCHW", {1, 2, 1, 1}); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("weights", {3, 3, 1, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9}); + RunValidationAndConversion(node_def, error::UNIMPLEMENTED, + "Dilation rate must be 1 for batch and channel " + "dimensions, at my_conv2d"); + } + { + // Dilation value is not 1 for channel (NHWC), should fail. + Reset(); + NodeDef node_def = + get_conv2d_nodedef({1, 1, 1, 1}, "SAME", "NHWC", {1, 1, 1, 2}); + AddTestTensor("input", {2, 3, 1}); + AddTestWeights("weights", {3, 3, 1, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9}); + RunValidationAndConversion(node_def, error::UNIMPLEMENTED, + "Dilation rate must be 1 for batch and channel " + "dimensions, at my_conv2d"); + } + { + // Dilation + Conv2DBackpropInput, should fail. + Reset(); + NodeDef node_def = + get_conv2d_nodedef({1, 1, 1, 1}, "SAME", "NHWC", {1, 1, 2, 1}, true); + AddTestTensor("input", {2, 3, 1}); + AddTestWeights("weights", {3, 3, 1, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9}); + AddTestWeights("input_sizes", {4}, {1, 2, 3, 1}); + RunValidationAndConversion(node_def, error::UNIMPLEMENTED, + "Dilation with Conv2DBackpropInput " + "(conv2d_transpose) is not supported, " + "at my_conv2d"); + } + { + // Strides is not 4D, should fail. + Reset(); + NodeDef node_def = + get_conv2d_nodedef({1, 1, 1}, "SAME", "NCHW", {1, 1, 1, 1}); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("weights", {3, 3, 1, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Convolution strides field must specify 4 dimensions, at my_conv2d"); + } + { + // Stride value is not 1 for channel, should fail. + Reset(); + NodeDef node_def = + get_conv2d_nodedef({1, 2, 1, 1}, "SAME", "NCHW", {1, 1, 1, 1}); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("weights", {3, 3, 1, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9}); + RunValidationAndConversion( + node_def, error::UNIMPLEMENTED, + "Stride must be 1 for batch and channel dimensions, at my_conv2d"); + } + + struct TestParams { + std::vector input_dims; + std::vector input; + std::vector filter_dims; + std::vector filter; + std::vector strides; + string padding; + string data_format; + std::vector dilations; + bool is_conv2d_backprop_input; + std::vector expected_output_dims; + std::vector expected_output; + }; + + // Ok. + const int kConv2DOKCases = 7; + TestParams ok_params[kConv2DOKCases] = { + // Basic + TestParams{/*input_dims=*/{1, 2, 3}, + /*input=*/{0, 1, 2, 3, 3, 4}, + /*filter_dims=*/{1, 2, 1, 1}, + /*filter=*/{-1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*padding=*/"VALID", + /*data_format=*/"NCHW", + /*dilations=*/{1, 1, 1, 1}, + /*is_conv2d_backprop_input=*/false, + /*expected_output_dims=*/{1, 2, 2}, + /*expected_output=*/{1, 1, 0, 1}}, + // SAME padding (Asymmetric) + TestParams{/*input_dims=*/{1, 2, 3}, + /*input=*/{0, 1, 2, 3, 3, 4}, + /*filter_dims=*/{1, 2, 1, 1}, + /*filter=*/{-1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*padding=*/"SAME", + /*data_format=*/"NCHW", + /*dilations=*/{1, 1, 1, 1}, + /*is_conv2d_backprop_input=*/false, + /*expected_output_dims=*/{1, 2, 3}, + /*expected_output=*/{1, 1, -2, 0, 1, -4}}, + // SAME padding (Symmetric) + TestParams{/*input_dims=*/{1, 2, 3}, + /*input=*/{0, 1, 2, 3, 3, 4}, + /*filter_dims=*/{1, 3, 1, 1}, + /*filter=*/{-1, 0, 1}, + /*strides=*/{1, 1, 1, 1}, + /*padding=*/"SAME", + /*data_format=*/"NCHW", + /*dilations=*/{1, 1, 1, 1}, + /*is_conv2d_backprop_input=*/false, + /*expected_output_dims=*/{1, 2, 3}, + /*expected_output=*/{1, 2, -1, 3, 1, -3}}, + // NHWC + TestParams{/*input_dims=*/{2, 3, 1}, + /*input=*/{0, 1, 2, 3, 3, 4}, + /*filter_dims=*/{1, 2, 1, 1}, + /*filter=*/{-1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*padding=*/"VALID", + /*data_format=*/"NHWC", + /*dilations=*/{1, 1, 1, 1}, + /*is_conv2d_backprop_input=*/false, + /*expected_output_dims=*/{2, 2, 1}, + /*expected_output=*/{1, 1, 0, 1}}, + // Dilated + TestParams{/*input_dims=*/{1, 2, 3}, + /*input=*/{0, 1, 2, 3, 3, 4}, + /*filter_dims=*/{1, 2, 1, 1}, + /*filter=*/{-1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*padding=*/"VALID", + /*data_format=*/"NCHW", + /*dilations=*/{1, 1, 1, 2}, + /*is_conv2d_backprop_input=*/false, + /*expected_output_dims=*/{1, 2, 1}, + /*expected_output=*/{2, 1}}, + // Strided + TestParams{/*input_dims=*/{1, 2, 4}, + /*input=*/{0, 1, 2, 2, 3, 4, 4, 7}, + /*filter_dims=*/{1, 2, 1, 1}, + /*filter=*/{-1, 1}, + /*strides=*/{1, 1, 1, 2}, + /*padding=*/"VALID", + /*data_format=*/"NCHW", + /*dilations=*/{1, 1, 1, 1}, + /*is_conv2d_backprop_input=*/false, + /*expected_output_dims=*/{1, 2, 2}, + /*expected_output=*/{1, 0, 1, 3}}, + // Transpose Strided + TestParams{/*input_dims=*/{1, 2, 2}, + /*input=*/{0, 1, 2, 3}, + /*filter_dims=*/{1, 2, 1, 1}, + /*filter=*/{-1, 1}, + /*strides=*/{1, 1, 1, 2}, + /*padding=*/"SAME", + /*data_format=*/"NCHW", + /*dilations=*/{1, 1, 1, 1}, + /*is_conv2d_backprop_input=*/true, + /*expected_output_dims=*/{1, 2, 4}, + /*expected_output=*/{0, 0, -1, 1, -2, 2, -3, 3}}, + }; + + for (int i = 0; i < kConv2DOKCases; i++) { + Reset(); + NodeDef node_def = get_conv2d_nodedef( + ok_params[i].strides, ok_params[i].padding, ok_params[i].data_format, + ok_params[i].dilations, ok_params[i].is_conv2d_backprop_input); + AddTestTensor("input", ok_params[i].input_dims); + AddTestWeights("weights", ok_params[i].filter_dims, + ok_params[i].filter); + if (ok_params[i].is_conv2d_backprop_input) { + AddTestWeights( + "input_sizes", + {static_cast(ok_params[i].expected_output.size())}, + ok_params[i].expected_output); + } + RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(GetTensorOrWeights("my_conv2d", &output)); + EXPECT_TRUE(output.is_tensor()); + ExpectTrtDimsEqualsArray(ok_params[i].expected_output_dims, + output.tensor()->getDimensions()); + + const DataVec input_data{ + {"input", test::AsTensor(ok_params[i].input)}}; + DataVec output_data{ + {"my_conv2d", + ConstructTensor(ok_params[i].expected_output.size())}}; + BuildAndRun(input_data, &output_data); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAreArray(ok_params[i].expected_output)); + } +} + +TEST_F(OpConverterTest, ConvertTopK) { + { + // Input list is empty, should fail. + NodeDef node_def = MakeNodeDef("my_topk", "TopKV2", {}); + RunValidationAndConversion(node_def, error::INVALID_ARGUMENT, + "Input expects tensor and weights, at my_topk"); + } + + for (const auto dtype : {DT_FLOAT, DT_INT32}) { + // Get the NodeDef for TopKV2. + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), dtype); + auto weights = ops::Placeholder(s.WithOpName("weights"), DT_INT32); + auto topk = ops::TopK(s.WithOpName("my_topk"), input, weights); + const NodeDef& node_def = topk.operation.node()->def(); + { + // K is a tensor, should fail. + Reset(); + AddTestTensor("input", {1, 2, 3}, /*batch_size=*/1, + /*trt_dtype=*/TfDataTypeToTrt(dtype)); + AddTestTensor("weights", {2}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Input expects tensor and weights, at my_topk"); + } + { + // Ok. + Reset(); + AddTestTensor("input", {1, 2, 5}); + AddTestWeights("weights", {1}, {2}); + RunValidationAndConversion(node_def); + TRT_TensorOrWeights outputs[2]; + TF_EXPECT_OK(GetTensorOrWeights("my_topk", &outputs[0])); + TF_EXPECT_OK(GetTensorOrWeights("my_topk:1", &outputs[1])); + for (auto& output : outputs) { + EXPECT_TRUE(output.is_tensor()); + ExpectTrtDimsEqualsArray({1, 2, 2}, output.tensor()->getDimensions()); + } + + const DataVec input_data{ + {"input", test::AsTensor({-9, 3, 5, 1, 6, -5, 7, 1, 0, -1})}}; + DataVec output_data{{"my_topk", ConstructTensor(4)}, + {"my_topk:1", ConstructTensor(4)}}; + BuildAndRun(input_data, &output_data); + EXPECT_THAT(GetSpanForData(output_data[0]), + ElementsAre(6, 5, 7, 1)); + EXPECT_THAT(GetSpanForData(output_data[1]), + ElementsAre(4, 2, 1, 2)); + } } } diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc similarity index 94% rename from tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc rename to tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc index d57f2300f8e6e6ce79c538133da6bc5cf5ead2f5..0eedfcacb4c11c8dc63fcfc13f044586b99b3c76 100644 --- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc @@ -12,9 +12,11 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h" -#include "tensorflow/contrib/tensorrt/convert/convert_graph.h" -#include "tensorflow/contrib/tensorrt/convert/utils.h" +#include "tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h" + +#include "absl/strings/str_cat.h" +#include "tensorflow/compiler/tf2tensorrt/convert/convert_graph.h" +#include "tensorflow/compiler/tf2tensorrt/convert/utils.h" #include "tensorflow/core/grappler/clusters/cluster.h" #include "tensorflow/core/grappler/grappler_item.h" #include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h" @@ -30,9 +32,9 @@ namespace tensorflow { namespace tensorrt { namespace convert { // TODO(sami): Remove VLOG messages once the code matures +using absl::StrAppend; +using absl::StrCat; using tensorflow::str_util::Uppercase; -using tensorflow::strings::StrAppend; -using tensorflow::strings::StrCat; tensorflow::Status TRTOptimizationPass::Init( const tensorflow::RewriterConfig_CustomGraphOptimizer* config) { @@ -64,7 +66,7 @@ tensorflow::Status TRTOptimizationPass::Init( max_workspace_size_bytes_ = params.at("max_workspace_size_bytes").i(); } if (params.count("precision_mode")) { - TF_RETURN_IF_ERROR(GetPrecisionMode( + TF_RETURN_IF_ERROR(TrtPrecisionModeFromName( Uppercase(params.at("precision_mode").s()), &precision_mode_)); } if (params.count("use_calibration")) { @@ -85,7 +87,7 @@ void TRTOptimizationPass::PrintDebugInfo( LOG(INFO) << offset << "type = " << cluster->type(); LOG(INFO) << offset << "num warmup steps = " << cluster->NumWarmupSteps(); const auto dev_names = cluster->GetDeviceNames(); - if (dev_names.size()) { + if (!dev_names.empty()) { LOG(INFO) << offset << " Device names:"; for (const auto s : dev_names) { LOG(INFO) << offset2 << s; @@ -101,7 +103,7 @@ void TRTOptimizationPass::PrintDebugInfo( } const auto dev_props = cluster->GetDevices(); - if (dev_props.size()) { + if (!dev_props.empty()) { LOG(INFO) << offset << "Device properties:"; for (auto k : dev_props) { LOG(INFO) << offset2 << k.first; @@ -129,7 +131,7 @@ void TRTOptimizationPass::PrintDebugInfo( } } LOG(INFO) << "item: " << item.id; - if (item.feed.size()) { + if (!item.feed.empty()) { LOG(INFO) << offset << "Feeds :"; for (const auto& f : item.feed) { const auto& shape = f.second.shape(); @@ -138,7 +140,7 @@ void TRTOptimizationPass::PrintDebugInfo( } else { LOG(INFO) << offset << "No Feeds"; } - if (item.fetch.size()) { + if (!item.fetch.empty()) { LOG(INFO) << offset << "Fetches :"; for (const auto& f : item.fetch) { LOG(INFO) << offset2 << f; @@ -147,7 +149,7 @@ void TRTOptimizationPass::PrintDebugInfo( LOG(INFO) << offset << "No Fetches"; } - if (item.init_ops.size()) { + if (!item.init_ops.empty()) { LOG(INFO) << offset << "init ops :"; for (const auto& f : item.init_ops) { LOG(INFO) << offset2 << f; @@ -158,7 +160,7 @@ void TRTOptimizationPass::PrintDebugInfo( LOG(INFO) << "Save Op = " << item.save_op; LOG(INFO) << "Restore Op = " << item.restore_op; LOG(INFO) << "save_restore_loc_tensor = " << item.save_restore_loc_tensor; - if (item.keep_ops.size()) { + if (!item.keep_ops.empty()) { LOG(INFO) << offset << "keep ops :"; for (const auto& f : item.keep_ops) { LOG(INFO) << offset2 << f; @@ -195,7 +197,7 @@ tensorflow::Status TRTOptimizationPass::Optimize( PrintDebugInfo(cluster, item); } int max_dim = -1; - if (item.feed.size()) { + if (!item.feed.empty()) { for (const auto& f : item.feed) { const auto& shape = f.second.shape(); if (shape.dims() > 0) { @@ -225,7 +227,7 @@ tensorflow::Status TRTOptimizationPass::Optimize( TF_RETURN_IF_ERROR(static_graph_properties.InferStatically(true)); tensorflow::tensorrt::convert::ConversionParams cp; - if (use_calibration_ && precision_mode_ != INT8MODE) { + if (use_calibration_ && precision_mode_ != TrtPrecisionMode::INT8) { VLOG(1) << "Calibration with FP32 or FP16 is not implemented. " << "Falling back to use_calibration = False." << "Note that the default value of use_calibration is True."; @@ -243,7 +245,7 @@ tensorflow::Status TRTOptimizationPass::Optimize( // If the last token is not an integer, it must be part of the name. // Otherwise it is port number. if (tokens.size() > 1 && - !strings::safe_strto32(tokens.back(), &dumm_port)) { + !strings::safe_strto32(tokens.back(), &dumm_port)) { // non-absl ok StrAppend(&s, ":", tokens.back()); } nodes_to_preserve.push_back(s); diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h similarity index 87% rename from tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h rename to tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h index 3e8dc0978e43e2e9ba07aaa09f74acfe8e59b9a7..b2aed2a37afb6c01863f5617bad0bafe004eec24 100644 --- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h +++ b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h @@ -13,11 +13,12 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_ #include +#include "tensorflow/compiler/tf2tensorrt/convert/utils.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h" #include "tensorflow/core/platform/logging.h" @@ -34,7 +35,7 @@ class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer { TRTOptimizationPass(const string& name = "TRTOptimizationPass") : name_(name), minimum_segment_size_(3), - precision_mode_(0), + precision_mode_(TrtPrecisionMode::FP32), maximum_batch_size_(-1), is_dynamic_op_(false), max_cached_batches_(1), @@ -62,7 +63,7 @@ class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer { private: const string name_; int minimum_segment_size_; - int precision_mode_; + TrtPrecisionMode precision_mode_; int maximum_batch_size_; bool is_dynamic_op_; std::vector batches_; @@ -77,4 +78,4 @@ class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer { #endif // GOOGLE_CUDA #endif // GOOGLE_TENSORRT -#endif // TENSORFLOW_CONTRIB_TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_ +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_ diff --git a/tensorflow/contrib/tensorrt/convert/utils.cc b/tensorflow/compiler/tf2tensorrt/convert/utils.cc similarity index 73% rename from tensorflow/contrib/tensorrt/convert/utils.cc rename to tensorflow/compiler/tf2tensorrt/convert/utils.cc index e7a1febb8c076891596741fe30721e7acca15a73..0ca3a5a4a58e6a3e29d3d515f496b8cb5e9f7eb0 100644 --- a/tensorflow/contrib/tensorrt/convert/utils.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/utils.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/convert/utils.h" +#include "tensorflow/compiler/tf2tensorrt/convert/utils.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" @@ -34,33 +34,32 @@ bool IsGoogleTensorRTEnabled() { #endif } -Status GetPrecisionModeName(const int precision_mode, string* name) { - switch (precision_mode) { - case FP32MODE: +Status TrtPrecisionModeToName(TrtPrecisionMode mode, string* name) { + switch (mode) { + case TrtPrecisionMode::FP32: *name = "FP32"; break; - case FP16MODE: + case TrtPrecisionMode::FP16: *name = "FP16"; break; - case INT8MODE: + case TrtPrecisionMode::INT8: *name = "INT8"; break; default: - return tensorflow::errors::OutOfRange("Unknown precision mode"); + return errors::OutOfRange("Unknown precision mode"); } return Status::OK(); } -Status GetPrecisionMode(const string& name, int* precision_mode) { +Status TrtPrecisionModeFromName(const string& name, TrtPrecisionMode* mode) { if (name == "FP32") { - *precision_mode = FP32MODE; + *mode = TrtPrecisionMode::FP32; } else if (name == "FP16") { - *precision_mode = FP16MODE; + *mode = TrtPrecisionMode::FP16; } else if (name == "INT8") { - *precision_mode = INT8MODE; + *mode = TrtPrecisionMode::INT8; } else { - return tensorflow::errors::InvalidArgument("Invalid precision mode name: ", - name); + return errors::InvalidArgument("Invalid precision mode name: ", name); } return Status::OK(); } diff --git a/tensorflow/contrib/tensorrt/convert/utils.h b/tensorflow/compiler/tf2tensorrt/convert/utils.h similarity index 72% rename from tensorflow/contrib/tensorrt/convert/utils.h rename to tensorflow/compiler/tf2tensorrt/convert/utils.h index 0592f31462af2b20f3a13fe5119e89c2ba42dd8a..0aa602dda2f3e98095bf72b5810a246c690d6741 100644 --- a/tensorflow/contrib/tensorrt/convert/utils.h +++ b/tensorflow/compiler/tf2tensorrt/convert/utils.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_TENSORRT_CONVERT_UTILS_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_CONVERT_UTILS_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_UTILS_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_UTILS_H_ #include @@ -35,16 +35,13 @@ using TrtUniquePtrType = std::unique_ptr>; bool IsGoogleTensorRTEnabled(); -// TODO(aaroey): use an enum instead. -const int FP32MODE = 0; -const int FP16MODE = 1; -const int INT8MODE = 2; +enum class TrtPrecisionMode { FP32, FP16, INT8 }; -Status GetPrecisionModeName(const int precision_mode, string* name); +Status TrtPrecisionModeToName(TrtPrecisionMode mode, string* name); -Status GetPrecisionMode(const string& name, int* precision_mode); +Status TrtPrecisionModeFromName(const string& name, TrtPrecisionMode* mode); } // namespace tensorrt } // namespace tensorflow -#endif // TENSORFLOW_CONTRIB_TENSORRT_CONVERT_UTILS_H_ +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_UTILS_H_ diff --git a/tensorflow/compiler/tf2tensorrt/kernels/get_serialized_resource_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/get_serialized_resource_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..81406b6e301ca350a3e52c97f5fcb575e88c3a90 --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/kernels/get_serialized_resource_op.cc @@ -0,0 +1,69 @@ +/* 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 + +#include "tensorflow/compiler/tf2tensorrt/utils/trt_resources.h" +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/resource_mgr.h" +#include "tensorflow/core/lib/core/refcount.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + +namespace tensorflow { +namespace tensorrt { + +class GetSerializedResourceOp : public OpKernel { + public: + explicit GetSerializedResourceOp(OpKernelConstruction* context) + : OpKernel(context) {} + + ~GetSerializedResourceOp() override {} + + void Compute(OpKernelContext* context) override { + // TODO(laigd): it will allocate the tensor on the device and copy the + // serialized string to that tensor, and later sess.run() will copy it back + // to host. We need to optimize this. + const string& container = context->input(0).scalar()(); + const string& resource_name = context->input(1).scalar()(); + + // Get the resource. + SerializableResourceBase* resource = nullptr; + OP_REQUIRES_OK(context, context->resource_manager()->Lookup( + container, resource_name, &resource)); + ::tensorflow::core::ScopedUnref sc(resource); + + // Serialize the resource as output. + string serialized_resource; + OP_REQUIRES_OK(context, resource->SerializeToString(&serialized_resource)); + + Tensor* output = nullptr; + OP_REQUIRES_OK(context, + context->allocate_output(0, TensorShape({}), &output)); + output->scalar()() = serialized_resource; + } +}; + +REGISTER_KERNEL_BUILDER(Name("GetSerializedResourceOp").Device(DEVICE_GPU), + GetSerializedResourceOp); + +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/compiler/tf2tensorrt/kernels/get_serialized_resource_op_test.cc b/tensorflow/compiler/tf2tensorrt/kernels/get_serialized_resource_op_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..ec038ebda073c8050321d5668b15a2c6faa72a4b --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/kernels/get_serialized_resource_op_test.cc @@ -0,0 +1,80 @@ +/* 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 +#include +#include +#include + +#include "tensorflow/compiler/tf2tensorrt/utils/trt_resources.h" +#include "tensorflow/core/framework/fake_input.h" +#include "tensorflow/core/framework/node_def_builder.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/kernels/ops_testutil.h" +#include "tensorflow/core/platform/test.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + +namespace tensorflow { +namespace tensorrt { + +class GetSerializedResourceOpTest : public OpsTestBase {}; + +TEST_F(GetSerializedResourceOpTest, Basic) { + // Create the GPU device. + std::unique_ptr device( + DeviceFactory::NewDevice("GPU", {}, "/job:worker/replica:0/task:0")); + + // Create the resource. + class MySerializableResource : public SerializableResourceBase { + public: + string DebugString() const override { return ""; } + Status SerializeToString(string* serialized) override { + *serialized = "my_serialized_str"; + return Status::OK(); + } + }; + const string container = "mycontainer"; + const string resource_name = "myresource"; + SerializableResourceBase* resource = new MySerializableResource(); + ResourceMgr* rm = device->resource_manager(); + EXPECT_TRUE(rm->Create(container, resource_name, resource).ok()); + + // Create the op. + SetDevice(DEVICE_GPU, std::move(device)); + TF_ASSERT_OK(NodeDefBuilder("op", "GetSerializedResourceOp") + .Input(FakeInput(DT_STRING)) + .Input(FakeInput(DT_STRING)) + .Finalize(node_def())); + TF_ASSERT_OK(InitOp()); + + // Execute the op. + AddInputFromArray(TensorShape({}), {container}); + AddInputFromArray(TensorShape({}), {resource_name}); + TF_ASSERT_OK(RunOpKernel()); + + // Verify the result. + // TODO(laigd): OpsTestBase::GetOutput() doesn't work. + Tensor* output = context_->mutable_output(0); + EXPECT_EQ("my_serialized_str", output->scalar()()); +} + +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc similarity index 57% rename from tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc rename to tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc index bad568644bb1f8d01d4cb0a7c853ec47d6f19e45..e3b31d736eb89b079410bba34b26d259ef2c2527 100644 --- a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc +++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc @@ -12,35 +12,45 @@ 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/contrib/tensorrt/kernels/trt_engine_op.h" - #include - -#include "tensorflow/contrib/tensorrt/convert/convert_nodes.h" -#include "tensorflow/contrib/tensorrt/convert/utils.h" -#include "tensorflow/contrib/tensorrt/log/trt_logger.h" -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.h" -#include "tensorflow/contrib/tensorrt/resources/trt_resource_manager.h" -#include "tensorflow/contrib/tensorrt/resources/trt_resources.h" -#include "tensorflow/contrib/tensorrt/test/utils.h" +#include +#include + +#include "absl/memory/memory.h" +#include "absl/strings/str_cat.h" +#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h" +#include "tensorflow/compiler/tf2tensorrt/convert/utils.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h" +#include "tensorflow/compiler/tf2tensorrt/utils/test_utils.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_resource_manager.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_resources.h" +#include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/graph_to_functiondef.h" +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/lib/core/refcount.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/stream_executor.h" +#include "tensorflow/core/platform/thread_annotations.h" #include "tensorflow/core/platform/types.h" #if GOOGLE_CUDA #if GOOGLE_TENSORRT #include "cuda/include/cuda_runtime_api.h" +#include "tensorrt/include/NvInfer.h" namespace tensorflow { namespace tensorrt { static Logger logger; +using absl::StrAppend; +using absl::StrCat; using ::nvinfer1::IRuntime; -using ::tensorflow::strings::StrAppend; -using ::tensorflow::strings::StrCat; // A helper class to call done() when destructed for asynchronous execution. // Helps simultaneous execution of native and TRT engines. @@ -53,6 +63,83 @@ class AsyncHelper : public tensorflow::core::RefCounted { AsyncOpKernel::DoneCallback done_; }; +// This OP can construct TRTEngine on the fly and if construction of engine +// fails, executes equivalent subgraph as a TensorFlow function. +class TRTEngineOp : public AsyncOpKernel { + public: + explicit TRTEngineOp(OpKernelConstruction* context); + + void ComputeAsync(OpKernelContext* context, + AsyncOpKernel::DoneCallback done) override; + + private: + // Execute calibration + void ExecuteCalibration(OpKernelContext* ctx, AsyncHelper* helper); + + // Construct a function handle for executing native funcdef graph + Status ConstructFunctionHandle(OpKernelContext* ctx); + + // Execute replaced native segment as function Op. + void ExecuteNativeSegment(OpKernelContext* ctx, AsyncHelper* helper); + + // Execute the tensorrt engine. Returns whether we need to retry by running + // the native segment. + bool ExecuteTrtEngine(OpKernelContext* ctx, EngineContext* engine_context); + + // Allocate necessary resources for calibration + Status AllocateCalibrationResources(OpKernelContext* ctx, + SerializableResourceBase** cr); + + // Get engine for the input shape + EngineContext* GetEngine(const std::vector& input_shapes, + OpKernelContext* ctx); + + // Return engine batch in cached_engne_batch_sizes_ which is closest to input + // batch. + bool GetCompatibleCachedEngine( + const std::vector& actual_input_shapes, + std::vector* engine_input_shapes); + + std::vector input_nodes_; + std::vector output_nodes_; + + // serialized protobuf segment or trt engine depending on static_engine_ flag. + string serialized_segment_; + + // Name of the function for TF native execution of the segment. + string funcdef_name_; + + // GraphDef representation of the segment. + GraphDef segment_graph_; + + // Engine Precision mode. + TrtPrecisionMode precision_mode_; + + // Whether engine is constructed during the conversion or needs to be + // constructed from protobuf segment. + bool static_engine_; + + // Whether to calibrate INT8 engine. + bool calibration_mode_; + + // Batches of the cached engines + std::vector cached_engine_batches_; + + // Maximum number of cached engines + int max_cached_engines_; + + int64 workspace_size_; + mutex engine_mutex_; + FunctionLibraryRuntime::Handle native_func_; + + // The finalized calibrator for inference. + std::unique_ptr calibrator_; + + // If true, create calibration graph for INT8 mode. Otherwise, we are using + // user-provided quantization ranges. + bool use_calibration_; +}; + #define TYPECASE(dt, X, Y) \ case dt: { \ return (void*)X->flat::Type>().data(); \ @@ -123,20 +210,20 @@ TRTEngineOp::TRTEngineOp(OpKernelConstruction* context) context->GetAttr("calibration_data", &calibration_data)); OP_REQUIRES_OK(context, context->GetAttr("segment_funcdef_name", &funcdef_name_)); - OP_REQUIRES_OK(context, GetPrecisionMode(precision_string, &precision_mode_)); + OP_REQUIRES_OK(context, + TrtPrecisionModeFromName(precision_string, &precision_mode_)); OP_REQUIRES_OK(context, context->GetAttr("use_calibration", &use_calibration_)); - calibration_mode_ = (use_calibration_ && precision_mode_ == INT8MODE && - calibration_data.size() == 0); - if (calibration_data.size()) { + calibration_mode_ = + (use_calibration_ && precision_mode_ == TrtPrecisionMode::INT8 && + calibration_data.empty()); + if (!calibration_data.empty()) { calibrator_.reset(new TRTInt8Calibrator(calibration_data)); calibration_data.resize(0); } native_func_ = tensorflow::kInvalidHandle; OP_REQUIRES_OK(context, context->GetAttr("max_cached_engines_count", &max_cached_engines_)); - OP_REQUIRES_OK(context, - context->GetAttr("fixed_input_size", &fixed_input_size_)); OP_REQUIRES_OK(context, context->GetAttr("cached_engine_batches", &cached_engine_batches_)); std::sort(cached_engine_batches_.begin(), cached_engine_batches_.end()); @@ -167,6 +254,7 @@ void TRTEngineOp::ExecuteNativeSegment(OpKernelContext* ctx, opts.rendezvous = ctx->rendezvous(); opts.cancellation_manager = ctx->cancellation_manager(); opts.runner = ctx->runner(); + inputs.reserve(ctx->num_inputs()); for (int i = 0; i < ctx->num_inputs(); i++) { inputs.push_back(ctx->input(i)); } @@ -175,11 +263,13 @@ void TRTEngineOp::ExecuteNativeSegment(OpKernelContext* ctx, lib->Run(opts, native_func_, inputs, outputs, [this, ctx, outputs, helper](const tensorflow::Status& s) { tensorflow::core::ScopedUnref sc(helper); - VLOG(1) << "Native Segment completed"; if (!s.ok()) { + LOG(ERROR) << "Failed to execute native segment " << this->name() + << ": " << s; ctx->SetStatus(s); return; } + VLOG(1) << "Native Segment completed"; for (size_t t = 0; t < outputs->size(); ++t) { ctx->set_output(t, outputs->at(t)); } @@ -194,18 +284,37 @@ void TRTEngineOp::ExecuteCalibration(OpKernelContext* ctx, VLOG(1) << "Executing TRT calibration: " << name(); helper->Ref(); tensorflow::core::ScopedUnref sc(helper); - // TODO(aaroey): remove the ResourceMgr singleton. - auto trt_rm = TRTResourceManager::instance(); - auto res_mgr = trt_rm->getManager("TRTCalibration"); + auto res_mgr = ctx->resource_manager(); TRTCalibrationResource* calib_res = nullptr; - auto status = res_mgr->LookupOrCreate( - funcdef_name_, "Calibrator", &calib_res, - {[ctx, this](TRTCalibrationResource** cr) -> tensorflow::Status { - return this->AllocateCalibrationResources(ctx, cr); - }}); - if (!status.ok()) { - ctx->SetStatus(status); - return; + 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); + // TODO(aaroey): here we also add the resource to the ResourceMgr singleton. + // This is needed before we migrate all uses of calib_graph_to_infer_graph() + // to the new calibration workflow. After that we'll remove this block. + { + auto deprecated_rm = + TRTResourceManager::instance()->getManager("TRTCalibration"); + TRTCalibrationResource* copied_resource = nullptr; + // Check whether the resource exists, and create it if not. + if (deprecated_rm->Lookup(funcdef_name_, "Calibrator", &copied_resource) + .ok()) { + // Do nothing if the resource exists. + copied_resource->Unref(); + } else { + copied_resource = calib_res; + // Increase the refcount by 1 then transfer the ownership of that refcount + // to the ResourceMgr singleton. + copied_resource->Ref(); + OP_REQUIRES_OK(ctx, deprecated_rm->Create(funcdef_name_, "Calibrator", + copied_resource)); + } } int num_inputs = ctx->num_inputs(); // Pass input data to calibrator @@ -219,7 +328,8 @@ void TRTEngineOp::ExecuteCalibration(OpKernelContext* ctx, return; } // Check the allocated buffer is sufficient for input - const auto device_tensor = dev_tensors_.at(i).AccessTensor(ctx); + const auto device_tensor = + calib_res->device_tensors_.at(i).AccessTensor(ctx); CHECK_EQ(t.TotalBytes(), device_tensor->TotalBytes()); input_data.emplace(StrCat(kInputPHName, i), data_address); } @@ -236,32 +346,34 @@ void TRTEngineOp::ExecuteCalibration(OpKernelContext* ctx, ExecuteNativeSegment(ctx, helper); } -int TRTEngineOp::GetEngineBatch(OpKernelContext* ctx) { - int num_batch = ctx->input(0).shape().dim_size(0); - int smallest_engine = 0; - for (const auto i : cached_engine_batches_) { - if (i >= num_batch) { - smallest_engine = i; - break; - } - } - // TODO(sami): Need an LRU here - if (smallest_engine == 0) { - if (max_cached_engines_ > cached_engine_batches_.size()) { - smallest_engine = num_batch; - cached_engine_batches_.push_back(num_batch); - VLOG(1) << "Running with batch size " << num_batch; - } else { - string msg = - StrCat("Engine buffer is full. buffer limit=", max_cached_engines_, - ", current entries="); - for (auto i : cached_engine_batches_) StrAppend(&msg, i, ","); - StrAppend(&msg, " requested batch=", num_batch); - LOG(WARNING) << msg; - return -1; +bool TRTEngineOp::GetCompatibleCachedEngine( + const std::vector& actual_input_shapes, + std::vector* engine_input_shapes) { + const int batch_size = actual_input_shapes[0].dim_size(0); + int smallest_batch_size = -1; + // Output shape will always be the same as the input but we will overwrite the + // batch size. + *engine_input_shapes = actual_input_shapes; + for (const int cached_batch_size : cached_engine_batches_) { + // Check if compatible: batch <= cached batch. + // + // TODO(laigd): here it only compare the first dim a.k.a the batch size, + // we'll need to to support non-batch dimensions as well. This will be done + // as part of the offline conversion implementation. + if (batch_size <= cached_batch_size) { + // First case: first compatible engine found + // Second case: smaller batch size engine found + if ((smallest_batch_size == -1) || + (cached_batch_size < smallest_batch_size)) { + smallest_batch_size = cached_batch_size; + // Overwrite batch size for output + for (int i = 0; i < engine_input_shapes->size(); i++) { + (*engine_input_shapes)[i].set_dim(0, smallest_batch_size); + } + } } } - return smallest_engine; + return (smallest_batch_size != -1); } void TRTEngineOp::ComputeAsync(OpKernelContext* ctx, @@ -272,25 +384,21 @@ void TRTEngineOp::ComputeAsync(OpKernelContext* ctx, ExecuteCalibration(ctx, helper); return; } - const int smallest_engine = GetEngineBatch(ctx); - if (smallest_engine < 0) { - LOG(WARNING) << "Failed to get engine batch, running native segment for " - << name(); - ExecuteNativeSegment(ctx, helper); - return; + // Get shapes of inputs to engine. + 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()); } - - const int num_batch = ctx->input(0).shape().dim_size(0); - auto& engine_ctx_pair = GetEngine(smallest_engine, ctx); - auto& trt_engine_ptr = engine_ctx_pair.first; - if (!trt_engine_ptr) { - LOG(WARNING) << "Engine retrieval for batch size " << num_batch + EngineContext* engine_context = GetEngine(input_shapes, ctx); + if (!engine_context->cuda_engine) { + LOG(WARNING) << "Engine retrieval for input shapes: " + << TensorShapeUtils::ShapeListString(input_shapes) << " failed. Running native segment for " << name(); ExecuteNativeSegment(ctx, helper); return; } - const bool retry = ExecuteTrtEngine(ctx, num_batch, trt_engine_ptr.get(), - engine_ctx_pair.second.get()); + const bool retry = ExecuteTrtEngine(ctx, engine_context); if (retry) { LOG(WARNING) << "Failed to execute engine, " << "retrying with native segment for " << name(); @@ -299,18 +407,19 @@ void TRTEngineOp::ComputeAsync(OpKernelContext* ctx, } } -bool TRTEngineOp::ExecuteTrtEngine( - OpKernelContext* ctx, const int num_batch, - nvinfer1::ICudaEngine* trt_engine_ptr, - nvinfer1::IExecutionContext* trt_execution_context_ptr) { +bool TRTEngineOp::ExecuteTrtEngine(OpKernelContext* ctx, + EngineContext* engine_context) { VLOG(1) << "Executing TRT engine: " << name(); + auto& cuda_engine = engine_context->cuda_engine; const bool kRetry = true; + // All inputs must have the same batch size, so just get it from the first + // input. + const int num_batch = ctx->input(0).shape().dim_size(0); const int num_binding = ctx->num_inputs() + ctx->num_outputs(); std::vector buffers(num_binding); for (int i = 0; i < ctx->num_inputs(); i++) { const string input_name = StrCat(kInputPHName, i); - const int binding_index = - trt_engine_ptr->getBindingIndex(input_name.c_str()); + const int binding_index = cuda_engine->getBindingIndex(input_name.c_str()); if (binding_index == -1) { LOG(ERROR) << "Input node not found, at " << input_name; return kRetry; @@ -323,10 +432,11 @@ bool TRTEngineOp::ExecuteTrtEngine( << " vs " << input_shape.dim_size(0); return kRetry; } - auto dtype = trt_engine_ptr->getBindingDataType(binding_index); + auto dtype = cuda_engine->getBindingDataType(binding_index); switch (dtype) { case nvinfer1::DataType::kFLOAT: - buffers[binding_index] = (void*)(input_tensor.flat().data()); + buffers[binding_index] = + const_cast(input_tensor.flat().data()); break; case nvinfer1::DataType::kHALF: LOG(ERROR) << "FP16 inputs are not supported yet!"; @@ -335,10 +445,11 @@ bool TRTEngineOp::ExecuteTrtEngine( LOG(ERROR) << "INT8 inputs are not supported yet!"; return kRetry; case nvinfer1::DataType::kINT32: - buffers[binding_index] = (void*)(input_tensor.flat().data()); + buffers[binding_index] = + const_cast(input_tensor.flat().data()); break; default: - LOG(ERROR) << "Unknown TRT data type: " << int(dtype); + LOG(ERROR) << "Unknown TRT data type: " << static_cast(dtype); return kRetry; } } @@ -346,13 +457,12 @@ bool TRTEngineOp::ExecuteTrtEngine( for (int i = 0; i < ctx->num_outputs(); i++) { // Create an output tensor const string output_name = StrCat(kOutputPHName, i); - const int binding_index = - trt_engine_ptr->getBindingIndex(output_name.c_str()); + const int binding_index = cuda_engine->getBindingIndex(output_name.c_str()); Tensor* output_tensor = nullptr; TensorShape output_shape; if (binding_index != -1) { - auto dims = trt_engine_ptr->getBindingDimensions(binding_index); + auto dims = cuda_engine->getBindingDimensions(binding_index); std::vector trt_shape(dims.nbDims + 1); trt_shape[0] = num_batch; for (int j = 0; j < dims.nbDims; j++) trt_shape[j + 1] = dims.d[j]; @@ -374,11 +484,11 @@ bool TRTEngineOp::ExecuteTrtEngine( // TODO(aaroey): ideally we should retry, fix this. return !kRetry; } - auto dtype = trt_engine_ptr->getBindingDataType(binding_index); + auto dtype = cuda_engine->getBindingDataType(binding_index); switch (dtype) { case nvinfer1::DataType::kFLOAT: buffers[binding_index] = - reinterpret_cast(output_tensor->flat().data()); + const_cast(output_tensor->flat().data()); break; case nvinfer1::DataType::kHALF: LOG(WARNING) << "half size is not supported yet!"; @@ -388,7 +498,7 @@ bool TRTEngineOp::ExecuteTrtEngine( return kRetry; case nvinfer1::DataType::kINT32: buffers[binding_index] = - reinterpret_cast(output_tensor->flat().data()); + const_cast(output_tensor->flat().data()); break; default: LOG(WARNING) << "Unknown TRT data type: " << static_cast(dtype); @@ -402,9 +512,12 @@ bool TRTEngineOp::ExecuteTrtEngine( ->implementation() ->GpuStreamMemberHack())); + // nvinfer1::IExecutionContext::enqueue is not thread safe and we need a mutex + // for it. + tensorflow::mutex_lock lock(engine_context->mu); // TODO(jie): trt enqueue does not return error - auto ret = trt_execution_context_ptr->enqueue(num_batch, &buffers[0], *stream, - nullptr); + auto ret = engine_context->execution_context->enqueue(num_batch, &buffers[0], + *stream, nullptr); if (!ret) { LOG(WARNING) << "Failed to enqueue batch for TRT engine: " << name(); return kRetry; @@ -414,50 +527,45 @@ bool TRTEngineOp::ExecuteTrtEngine( return !kRetry; } -TRTEngineOp::~TRTEngineOp() { - // We need to manually destroy the engine and execution context before - // the allocator is destructed. - for (auto& eng : engine_map_) { - eng.second.first.reset(); - eng.second.second.reset(); +EngineContext* TRTEngineOp::GetEngine( + const std::vector& input_shapes, OpKernelContext* ctx) { + static EngineContext empty_context; + tensorflow::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); + + // Get engine cache + TRTEngineCacheResource* cache_res = nullptr; + auto status = ctx->resource_manager()->LookupOrCreate( + "TRTEngineCache", funcdef_name_, &cache_res, + {[this, ctx](TRTEngineCacheResource** cr) -> tensorflow::Status { + *cr = new TRTEngineCacheResource(ctx, this->max_cached_engines_); + return Status::OK(); + }}); + if (!status.ok()) { + ctx->SetStatus(status); + return &empty_context; } - allocator_.reset(); -} - -nvinfer1::IGpuAllocator* TRTEngineOp::GetAllocator(OpKernelContext* ctx) { - if (allocator_) return allocator_.get(); - auto device = ctx->device(); - auto alloc = device->GetAllocator(tensorflow::AllocatorAttributes()); - if (!alloc) { - LOG(ERROR) << "Can't find device allocator for gpu device " - << device->name(); - return nullptr; + tensorflow::core::ScopedUnref sc(cache_res); + auto& cache = cache_res->cache_; + auto allocator = cache_res->allocator_.get(); + if (allocator == nullptr) { + return &empty_context; } - allocator_.reset(new TRTDeviceAllocator(alloc)); - return allocator_.get(); -} - -TRTEngineOp::EngineCtxPair& TRTEngineOp::GetEngine(int batch_size, - OpKernelContext* ctx) { - static EngineCtxPair null_pair = { - TrtUniquePtrType(nullptr), - TrtUniquePtrType(nullptr)}; - // TODO(sami): This method needs to be re-written to use resource manager and - // with LRU mechanism option. - tensorflow::mutex_lock lock(engine_mutex_); + // Handle the static engine case. For static engines, the cache will have a + // single element containing the only engine. if (static_engine_) { - if (engine_map_.size()) { - if (engine_map_.begin()->first >= batch_size) { - return engine_map_.begin()->second; + if (cache.size()) { + // Batch size of engine must be >= the input batch size + // TODO(tmorris): use match compatible function? + if (cache.begin()->first[0].dim_size(0) >= batch_size) { + return cache.begin()->second.get(); } - return null_pair; + return &empty_context; } + TrtUniquePtrType infer(nvinfer1::createInferRuntime(logger)); - auto allocator = GetAllocator(ctx); - if (allocator == nullptr) { - return null_pair; - } infer->setGpuAllocator(allocator); TrtUniquePtrType static_engine( infer->deserializeCudaEngine(serialized_segment_.c_str(), @@ -465,62 +573,87 @@ TRTEngineOp::EngineCtxPair& TRTEngineOp::GetEngine(int batch_size, PluginFactoryTensorRT::GetInstance())); auto raw_static_engine = static_engine.get(); const auto max_batch_size = raw_static_engine->getMaxBatchSize(); - engine_map_[max_batch_size] = { - std::move(static_engine), - TrtUniquePtrType( - raw_static_engine->createExecutionContext())}; + // Static engine will have max_batch_size for batch size so that all inputs + // will map to this single engine. + std::vector engine_input_shapes(input_shapes); + for (int i = 0; i < engine_input_shapes.size(); i++) { + // TODO(tmorris): will all inputs have batch size as first dimension?? + engine_input_shapes[i].set_dim(0, max_batch_size); + } + // TODO(laigd): here we assume engine_input_shapes matches the actual input + // shapes of the engine, we should verify that. + cache.emplace(engine_input_shapes, + absl::make_unique( + std::move(static_engine), + TrtUniquePtrType( + raw_static_engine->createExecutionContext()))); // Runtime is safe to delete after engine creation serialized_segment_.clear(); if (max_batch_size < batch_size) { - return null_pair; + return &empty_context; } - return engine_map_.at(max_batch_size); + return cache.at(engine_input_shapes).get(); } // static_engine_ // Handle the dynamic engine case. - auto engine_it = engine_map_.find(batch_size); - if (engine_it == engine_map_.end() && - engine_map_.size() < (size_t)max_cached_engines_) { - nvinfer1::IGpuAllocator* allocator = nullptr; - allocator = GetAllocator(ctx); - if (allocator == nullptr) { - return null_pair; - } - std::vector shapes; - for (int i = 0; i < ctx->num_inputs(); ++i) { - shapes.emplace_back(ctx->input(i).shape()); + // See if there is a compatible engine cached. The batch size should be <= the + // cached batch size. + 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 + // exact shape and possibly create a new engine if it is not in cache. + if (!matched_successfully) { + engine_input_shapes = input_shapes; + if (!cached_engine_batches_.empty()) { + // If user has explicitly defined cached_engine_batches, we should + // warn them that their input was non-compatible (batch size too high) + LOG(WARNING) << "No compatible cached engine was found for batch size: " + << batch_size << ". A new engine will be created."; + cached_engine_batches_.push_back(batch_size); } + } + + if (!cache.count(engine_input_shapes)) { TrtUniquePtrType engine; bool convert_successfully = false; LOG(INFO) << "Building a new TensorRT engine for " << name() - << " with batch size " << batch_size; + << " input shapes: " + << TensorShapeUtils::ShapeListString(engine_input_shapes); + + // Convert to partial shapes + std::vector partial_shapes(engine_input_shapes.begin(), + engine_input_shapes.end()); + // Up to this point, calibrator_ can never be empty, since otherwise it // means calibration_mode_ is true and this path won't get executed. auto status = convert::ConvertGraphDefToEngine( - segment_graph_, precision_mode_, batch_size, workspace_size_, shapes, - &logger, allocator, calibrator_.get(), &engine, use_calibration_, - &convert_successfully); + segment_graph_, precision_mode_, batch_size, workspace_size_, + partial_shapes, &logger, allocator, calibrator_.get(), &engine, + use_calibration_, &convert_successfully); if (!status.ok()) { if (convert_successfully) { // This means it fail to build the engine even when the network is built // successfully, probably due to internal issues. In this case we don't // retry in the future. - engine_map_[batch_size] = {nullptr, nullptr}; + cache.emplace(engine_input_shapes, absl::make_unique()); } LOG(WARNING) << "Engine creation for batch size " << batch_size << " failed " << status; - return null_pair; + return &empty_context; } VLOG(1) << "Conversion is done"; TrtUniquePtrType exec_context( engine->createExecutionContext()); - engine_map_[batch_size] = {std::move(engine), std::move(exec_context)}; + cache.emplace(engine_input_shapes, + absl::make_unique(std::move(engine), + std::move(exec_context))); } - return engine_map_.at(batch_size); + return cache.at(engine_input_shapes).get(); } tensorflow::Status TRTEngineOp::AllocateCalibrationResources( - OpKernelContext* ctx, TRTCalibrationResource** cr) { + OpKernelContext* ctx, SerializableResourceBase** cr) { auto cres = new TRTCalibrationResource(); *cr = cres; // Get the allocator. @@ -536,7 +669,7 @@ tensorflow::Status TRTEngineOp::AllocateCalibrationResources( const int batch_size = ctx->input(0).dim_size(0); const int num_inputs = ctx->num_inputs(); std::vector shapes; - dev_tensors_.resize(num_inputs); + cres->device_tensors_.resize(num_inputs); VLOG(1) << " Constructing calibrator"; for (int i = 0; i < num_inputs; i++) { // allocate workspace on device for inputs @@ -544,19 +677,19 @@ tensorflow::Status TRTEngineOp::AllocateCalibrationResources( shapes.emplace_back(t.shape()); Tensor* device_tensor; TF_RETURN_IF_ERROR(ctx->allocate_persistent( - t.dtype(), t.shape(), &dev_tensors_.at(i), &device_tensor)); + t.dtype(), t.shape(), &cres->device_tensors_.at(i), &device_tensor)); CHECK_EQ(t.TotalBytes(), device_tensor->TotalBytes()); void* device_address = GetTensorAddress(device_tensor); if (device_address == nullptr) { return tensorflow::errors::InvalidArgument( "Unsupported data type encountered in input ", i); } - device_buffers_.emplace( + cres->device_buffers_.emplace( StrCat(kInputPHName, i), std::pair(device_address, device_tensor->TotalBytes())); } cres->calibrator_.reset( - new TRTInt8Calibrator(device_buffers_, batch_size, name())); + new TRTInt8Calibrator(cres->device_buffers_, batch_size, name())); const string label(name()); auto segment_graph = &segment_graph_; const int platform_gpu_id = @@ -585,9 +718,10 @@ tensorflow::Status TRTEngineOp::AllocateCalibrationResources( // TODO(aaroey): maybe setting the max batch size using the python // calibration wrapper class. auto s = convert::ConvertGraphDefToEngine( - *segment_graph, INT8MODE, cres->calibrator_->getBatchSize(), - workspace_size_bytes, shapes, &cres->logger_, cres->allocator_.get(), - cres->calibrator_.get(), &cres->engine_, + *segment_graph, TrtPrecisionMode::INT8, + cres->calibrator_->getBatchSize(), workspace_size_bytes, shapes, + &cres->logger_, cres->allocator_.get(), cres->calibrator_.get(), + &cres->engine_, /*use_calibration=*/true, /*convert_successfully=*/nullptr); if (!s.ok()) { diff --git a/tensorflow/core/kernels/data/model_dataset_op.h b/tensorflow/compiler/tf2tensorrt/ops/get_serialized_resource_op.cc similarity index 51% rename from tensorflow/core/kernels/data/model_dataset_op.h rename to tensorflow/compiler/tf2tensorrt/ops/get_serialized_resource_op.cc index 38d3fcb234feeec49d2954622831f23262239488..59da73f5efc8eedc20c35cf35cb1eae6cda136c9 100644 --- a/tensorflow/core/kernels/data/model_dataset_op.h +++ b/tensorflow/compiler/tf2tensorrt/ops/get_serialized_resource_op.cc @@ -13,30 +13,28 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CORE_KERNELS_DATA_MODEL_DATASET_OP_H_ -#define TENSORFLOW_CORE_KERNELS_DATA_MODEL_DATASET_OP_H_ +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT -#include "tensorflow/core/framework/dataset.h" -#include "tensorflow/core/framework/model.h" +#include "tensorflow/core/framework/common_shape_fns.h" +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/shape_inference.h" +#include "tensorflow/core/framework/tensor_shape.h" namespace tensorflow { -namespace data { -class ModelDatasetOp : public UnaryDatasetOpKernel { - public: - using ModelFactory = std::function()>; +REGISTER_OP("GetSerializedResourceOp") + .Input("container: string") + .Input("resource_name: string") + .Output("serialized_resource: string") + .SetShapeFn(shape_inference::ScalarShape) + .SetIsStateful() + .Doc(R"doc( +Gets a resource from a container managed by the resource manager and returns +its serialized representation. +)doc"); - explicit ModelDatasetOp(OpKernelConstruction* ctx); - - void MakeDataset(OpKernelContext* ctx, DatasetBase* input, - DatasetBase** output) override; - - virtual ModelFactory CreateModelFactory() const; - - private: - class Dataset; -}; - -} // namespace data } // namespace tensorflow -#endif // TENSORFLOW_CORE_KERNELS_DATA_MODEL_DATASET_OP_H_ + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc similarity index 80% rename from tensorflow/contrib/tensorrt/ops/trt_engine_op.cc rename to tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc index 92405906eb76b043bc08b68e25e16ab40197dddf..b84d2fe0b8cef3475f2a7d0f5383d5e11cde099a 100644 --- a/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc +++ b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc @@ -28,16 +28,22 @@ namespace shape_inference { extern Status TRTEngineOpShapeInference(InferenceContext* c); } +// NOTE: please try NOT to add/modify/remove attributes or inputs/outputs to the +// list below, this will break backward compatibility! +// +// TODO(laigd): consider making this op stateful. The only problem is it uses TF +// function which has to be stateless, but we can use function library as the +// key to cache the instantiated functions for different executor subgraphs. REGISTER_OP("TRTEngineOp") .Attr("serialized_segment: string") .Attr("input_shapes: list(shape)") .Attr("output_shapes: list(shape)") .Attr("segment_funcdef_name: string") - .Attr("InT: list({int8,float16,float32})") - .Attr("OutT: list({int8,float16,float32})") + .Attr("InT: list({int8,float16,float32,int32})") + .Attr("OutT: list({int8,float16,float32,int32})") .Attr("static_engine: bool = true") .Attr("fixed_input_size: bool = true") - .Attr("cached_engine_batches: list(int) = []") + .Attr("cached_engine_batches: list(int) >= 0 = []") .Attr("max_cached_engines_count: int = 1") .Attr("workspace_size_bytes: int") .Attr("precision_mode: {'FP32', 'FP16', 'INT8'}") diff --git a/tensorflow/contrib/tensorrt/plugin/trt_plugin.cc b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc similarity index 96% rename from tensorflow/contrib/tensorrt/plugin/trt_plugin.cc rename to tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc index 062f86e8bb4dc753925e4e2baf0bc80a5312a94f..a4341c530fffca88c82813cc2ace2c0ae1df5345 100644 --- a/tensorflow/contrib/tensorrt/plugin/trt_plugin.cc +++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.cc @@ -13,10 +13,12 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h" + #include #include -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin_utils.h" + +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_utils.h" #if GOOGLE_CUDA #if GOOGLE_TENSORRT diff --git a/tensorflow/contrib/tensorrt/plugin/trt_plugin.h b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h similarity index 92% rename from tensorflow/contrib/tensorrt/plugin/trt_plugin.h rename to tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h index 754920b60ca7439513a91ad0354833a2482b29c1..f495d857037c79a1783f8eb232fb57c20e229169 100644 --- a/tensorflow/contrib/tensorrt/plugin/trt_plugin.h +++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_TENSORRT_PLUGIN_TRT_PLUGIN_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_PLUGIN_TRT_PLUGIN_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_PLUGIN_TRT_PLUGIN_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_PLUGIN_TRT_PLUGIN_H_ #include #include @@ -71,4 +71,4 @@ class PluginTensorRT : public nvinfer1::IPlugin { #endif // GOOGLE_TENSORRT #endif // GOOGLE_CUDA -#endif // TENSORFLOW_CONTRIB_TENSORRT_PLUGIN_TRT_PLUGIN_H_ +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_PLUGIN_TRT_PLUGIN_H_ diff --git a/tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.cc b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.cc similarity index 96% rename from tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.cc rename to tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.cc index cccc91226265ed139fb8db0b71c40b868f729562..871fb1210bd495dc3f5e8153bb6c3a361bf569f5 100644 --- a/tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.cc +++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h" #if GOOGLE_CUDA #if GOOGLE_TENSORRT diff --git a/tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.h b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h similarity index 91% rename from tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.h rename to tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h index bbae9fb65c22cf69d2e7954436fd04dd16f7f6c8..9aa99a40b80de92a4d9b9ad36e88e693b8aa42dc 100644 --- a/tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.h +++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h @@ -13,14 +13,14 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_TENSORRT_PLUGIN_TRT_PLUGIN_FACTORY_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_PLUGIN_TRT_PLUGIN_FACTORY_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_PLUGIN_TRT_PLUGIN_FACTORY_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_PLUGIN_TRT_PLUGIN_FACTORY_H_ #include #include -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin.h" -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin_utils.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_utils.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/mutex.h" @@ -99,4 +99,4 @@ class TrtPluginRegistrar { #endif // GOOGLE_TENSORRT #endif // GOOGLE_CUDA -#endif // TENSORFLOW_CONTRIB_TENSORRT_PLUGIN_TRT_PLUGIN_FACTORY_H_ +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_PLUGIN_TRT_PLUGIN_FACTORY_H_ diff --git a/tensorflow/contrib/tensorrt/plugin/trt_plugin_factory_test.cc b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory_test.cc similarity index 96% rename from tensorflow/contrib/tensorrt/plugin/trt_plugin_factory_test.cc rename to tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory_test.cc index 129bdcdbc2f8d9d5215f45f381bcadf35e4fa75e..7d9c465c22beed0e252cbc26d6c533a0789d4f49 100644 --- a/tensorflow/contrib/tensorrt/plugin/trt_plugin_factory_test.cc +++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory_test.cc @@ -13,9 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h" -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/test.h" diff --git a/tensorflow/contrib/tensorrt/plugin/trt_plugin_utils.cc b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_utils.cc similarity index 94% rename from tensorflow/contrib/tensorrt/plugin/trt_plugin_utils.cc rename to tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_utils.cc index a8f60886c03c174a612e7a135b6eb7bb7cb9997a..f3d6b4ff476139693a5251ddf58a3200d8af8efc 100644 --- a/tensorflow/contrib/tensorrt/plugin/trt_plugin_utils.cc +++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_utils.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin_utils.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_utils.h" #include #if GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/plugin/trt_plugin_utils.h b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_utils.h similarity index 82% rename from tensorflow/contrib/tensorrt/plugin/trt_plugin_utils.h rename to tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_utils.h index 274ce42fec9283c643004d45fba461879fc5f2dc..e5eff15c19694093c7a5ea933a41375e8e01c8b9 100644 --- a/tensorflow/contrib/tensorrt/plugin/trt_plugin_utils.h +++ b/tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_utils.h @@ -13,12 +13,12 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_TENSORRT_PLUGIN_TRT_PLUGIN_UTILS_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_PLUGIN_TRT_PLUGIN_UTILS_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_PLUGIN_TRT_PLUGIN_UTILS_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_PLUGIN_TRT_PLUGIN_UTILS_H_ #include -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h" #include "tensorflow/core/platform/types.h" #if GOOGLE_CUDA @@ -43,4 +43,4 @@ string ExtractOpName(const void* serial_data, size_t serial_length, #endif // GOOGLE_TENSORRT #endif // GOOGLE_CUDA -#endif // TENSORFLOW_CONTRIB_TENSORRT_PLUGIN_TRT_PLUGIN_UTILS_H_ +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_PLUGIN_TRT_PLUGIN_UTILS_H_ diff --git a/tensorflow/compiler/tf2tensorrt/python/ops/trt_ops.py b/tensorflow/compiler/tf2tensorrt/python/ops/trt_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..7503d4d984e00b0adfbae350f298360aa2eae01d --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/python/ops/trt_ops.py @@ -0,0 +1,62 @@ +# 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. +# ============================================================================= +"""Exposes the Python wrapper of TRTEngineOp.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import threading + +import platform +from tensorflow.python.framework import errors + +_trt_ops_so = None +_module_lock = threading.Lock() + + +def load_trt_ops(): + """Load TF-TRT op libraries so if it hasn't been loaded already.""" + global _trt_ops_so + + if platform.system() == "Windows": + raise RuntimeError("Windows platforms are not supported") + + with _module_lock: + if _trt_ops_so: + return + + # TODO(laigd): we should load TF-TRT kernels here as well after removing the + # swig binding. + try: + # TODO(lagid): It is not known why these unused imports were introduced. + # Investigate and get rid of these, if not required. + # pylint: disable=unused-import,g-import-not-at-top,unused-variable + from tensorflow.compiler.tf2tensorrt.ops.gen_trt_ops import trt_engine_op + from tensorflow.python.framework import load_library + from tensorflow.python.platform import resource_loader + # pylint: enable=unused-import,g-import-not-at-top,unused-variable + + _trt_ops_so = load_library.load_op_library( + resource_loader.get_path_to_datafile("_trt_ops.so")) + except errors.NotFoundError as e: + no_trt_message = ( + "**** Failed to initialize TensorRT. This is either because the " + "TensorRT installation path is not in LD_LIBRARY_PATH, or because " + "you do not have it installed. If not installed, please go to " + "https://developer.nvidia.com/tensorrt to download and install " + "TensorRT ****") + print(no_trt_message) + raise e diff --git a/tensorflow/contrib/tensorrt/segment/segment.cc b/tensorflow/compiler/tf2tensorrt/segment/segment.cc similarity index 97% rename from tensorflow/contrib/tensorrt/segment/segment.cc rename to tensorflow/compiler/tf2tensorrt/segment/segment.cc index 084a96e0fa5c97edc58adf2590ed94e5ef0e4d85..3794929b1df3fa999de6ab218dc2ddfb96e4ac81 100644 --- a/tensorflow/contrib/tensorrt/segment/segment.cc +++ b/tensorflow/compiler/tf2tensorrt/segment/segment.cc @@ -13,14 +13,15 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/segment/segment.h" +#include "tensorflow/compiler/tf2tensorrt/segment/segment.h" #include #include #include #include -#include "tensorflow/contrib/tensorrt/segment/union_find.h" +#include "absl/strings/str_cat.h" +#include "tensorflow/compiler/tf2tensorrt/segment/union_find.h" #include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/graph/graph_constructor.h" @@ -29,11 +30,14 @@ limitations under the License. #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/types.h" +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + namespace tensorflow { namespace tensorrt { namespace segment { -using ::tensorflow::strings::StrAppend; -using ::tensorflow::strings::StrCat; +using absl::StrAppend; +using absl::StrCat; // A simple graph representation to mirror tensorflow::Graph. This structure // helps saving memory since segmenter modifies the graph in place, preventing @@ -673,10 +677,11 @@ tensorflow::Status SegmentGraph( // --------------------------------- Step 3 --------------------------------- // Convert the segments into the expected return format for (const auto& itr : sg_map) { - const std::set& segment_nodes = - itr.second; + const string& segment_root = itr.first; + // Return format does not require set comparator. + std::set segment_nodes(itr.second.begin(), itr.second.end()); if (VLOG_IS_ON(1)) { - string s = "parent=" + itr.first + ":"; + string s = "parent=" + segment_root + ":"; for (auto node : segment_nodes) s += " " + node->name(); VLOG(1) << "Segment " << segments->size() << ": " << s; } @@ -689,12 +694,10 @@ tensorflow::Status SegmentGraph( } // TODO(sami): Make segmenter placement aware once trtscopes are in place - std::set segment_node_names; - for (auto node : itr.second) segment_node_names.insert(node->name()); - const auto& dev_itr = device_maps.find(itr.first); + const auto& dev_itr = device_maps.find(segment_root); if (dev_itr == device_maps.end() || dev_itr->second.empty()) { VLOG(1) << "No device assigned to segment " << segments->size(); - segments->emplace_back(std::make_pair(segment_node_names, string())); + segments->emplace_back(std::make_pair(segment_nodes, string())); } else if (dev_itr->second.size() > 1) { string s("Segment "); StrAppend(&s, segments->size(), " has multiple devices attached: "); @@ -703,10 +706,10 @@ tensorflow::Status SegmentGraph( } LOG(WARNING) << s << " choosing " << *(dev_itr->second.begin()); segments->emplace_back( - std::make_pair(segment_node_names, *(dev_itr->second.begin()))); + std::make_pair(segment_nodes, *(dev_itr->second.begin()))); } else { segments->emplace_back( - std::make_pair(segment_node_names, *(dev_itr->second.begin()))); + std::make_pair(segment_nodes, *(dev_itr->second.begin()))); } } if (VLOG_IS_ON(1)) { @@ -725,3 +728,6 @@ tensorflow::Status SegmentGraph( } // namespace segment } // namespace tensorrt } // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/segment/segment.h b/tensorflow/compiler/tf2tensorrt/segment/segment.h similarity index 81% rename from tensorflow/contrib/tensorrt/segment/segment.h rename to tensorflow/compiler/tf2tensorrt/segment/segment.h index b9693aad1b764515459db6833b05221ea5b3a2d1..9622ddd593990e93ba1b54e9dfd0052006e20ced 100644 --- a/tensorflow/contrib/tensorrt/segment/segment.h +++ b/tensorflow/compiler/tf2tensorrt/segment/segment.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_TENSORRT_SEGMENT_SEGMENT_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_SEGMENT_SEGMENT_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_SEGMENT_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_SEGMENT_H_ #include #include @@ -24,15 +24,17 @@ limitations under the License. #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" -namespace tensorflow { +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +namespace tensorflow { namespace tensorrt { namespace segment { -// Vector of segments, each entry contains a set of node names and a device name -// in the segment. -// TODO(aaroey): use node pointer instead of node name. -using SegmentNodesVector = std::vector, string>>; +// Vector of segments, each entry contains a set of node pointers and a device +// name in the segment. +using SegmentNodesVector = + std::vector, string>>; struct SegmentOptions { // Segment must contain at least this many nodes. @@ -60,4 +62,7 @@ tensorflow::Status SegmentGraph( } // namespace tensorrt } // namespace tensorflow -#endif // TENSORFLOW_CONTRIB_TENSORRT_SEGMENT_SEGMENT_H_ +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA + +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_SEGMENT_H_ diff --git a/tensorflow/contrib/tensorrt/segment/segment_test.cc b/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc similarity index 97% rename from tensorflow/contrib/tensorrt/segment/segment_test.cc rename to tensorflow/compiler/tf2tensorrt/segment/segment_test.cc index 4805ef9c61a7784a1c08cf5eaf504691bc9dbedc..e11ad2719740d908f93ef580a6b308469365f402 100644 --- a/tensorflow/contrib/tensorrt/segment/segment_test.cc +++ b/tensorflow/compiler/tf2tensorrt/segment/segment_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/segment/segment.h" +#include "tensorflow/compiler/tf2tensorrt/segment/segment.h" #include "tensorflow/cc/framework/scope.h" #include "tensorflow/cc/ops/standard_ops.h" @@ -26,6 +26,9 @@ limitations under the License. #include "tensorflow/core/platform/types.h" #include "tensorflow/core/public/session.h" +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + namespace tensorflow { namespace tensorrt { namespace segment { @@ -75,7 +78,10 @@ class SegmentTest : public ::testing::Test { const std::vector>& expected_segments) { EXPECT_EQ(expected_segments.size(), segments.size()); for (int i = 0; i < segments.size(); ++i) { - const auto& segment_node_names = segments[i].first; + std::set segment_node_names; + for (const Node* node : segments[i].first) { + segment_node_names.insert(node->name()); + } const auto& expected = expected_segments[i]; for (const auto& name : expected) { EXPECT_TRUE(segment_node_names.count(name)) @@ -262,3 +268,6 @@ TEST_F(SegmentTest, BigIfElse) { } // namespace segment } // namespace tensorrt } // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/segment/union_find.h b/tensorflow/compiler/tf2tensorrt/segment/union_find.h similarity index 92% rename from tensorflow/contrib/tensorrt/segment/union_find.h rename to tensorflow/compiler/tf2tensorrt/segment/union_find.h index 1c64ebbb0ae532a4776ab8963515d19fd3b23b4c..6458ae692fd7c922b5fc3bea2e55b613447dbde0 100644 --- a/tensorflow/contrib/tensorrt/segment/union_find.h +++ b/tensorflow/compiler/tf2tensorrt/segment/union_find.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_TENSORRT_SEGMENT_UNION_FIND_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_SEGMENT_UNION_FIND_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_UNION_FIND_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_UNION_FIND_H_ namespace tensorflow { namespace tensorrt { @@ -76,4 +76,4 @@ UnionFind* UnionFind::FindRoot() { } // namespace tensorrt } // namespace tensorflow -#endif // TENSORFLOW_CONTRIB_TENSORRT_SEGMENT_UNION_FIND_H_ +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_SEGMENT_UNION_FIND_H_ diff --git a/tensorflow/contrib/tensorrt/tensorrt_test.cc b/tensorflow/compiler/tf2tensorrt/tensorrt_test.cc similarity index 100% rename from tensorflow/contrib/tensorrt/tensorrt_test.cc rename to tensorflow/compiler/tf2tensorrt/tensorrt_test.cc diff --git a/tensorflow/contrib/tensorrt/test/utils.cc b/tensorflow/compiler/tf2tensorrt/utils/test_utils.cc similarity index 97% rename from tensorflow/contrib/tensorrt/test/utils.cc rename to tensorflow/compiler/tf2tensorrt/utils/test_utils.cc index 276308b3a0a6ce864969afb0179c6a3f00d6b70b..3bcca99afbff8b84d2dd628ae9211ee94e86af2a 100644 --- a/tensorflow/contrib/tensorrt/test/utils.cc +++ b/tensorflow/compiler/tf2tensorrt/utils/test_utils.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/test/utils.h" +#include "tensorflow/compiler/tf2tensorrt/utils/test_utils.h" #include #include diff --git a/tensorflow/contrib/tensorrt/test/utils.h b/tensorflow/compiler/tf2tensorrt/utils/test_utils.h similarity index 89% rename from tensorflow/contrib/tensorrt/test/utils.h rename to tensorflow/compiler/tf2tensorrt/utils/test_utils.h index 4bb4120206cfaae70107e55d1818e3af2f02717a..bcd628b62f0320f7ce9dfe6240316d876f1d5a20 100644 --- a/tensorflow/contrib/tensorrt/test/utils.h +++ b/tensorflow/compiler/tf2tensorrt/utils/test_utils.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_TENSORRT_TEST_UTILS_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_TEST_UTILS_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TEST_UTILS_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TEST_UTILS_H_ #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" @@ -41,4 +41,4 @@ string GetTestValue(const string& label); } // namespace tensorrt } // namespace tensorflow -#endif // TENSORFLOW_CONTRIB_TENSORRT_TEST_UTILS_H_ +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TEST_UTILS_H_ diff --git a/tensorflow/contrib/tensorrt/resources/trt_allocator.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc similarity index 98% rename from tensorflow/contrib/tensorrt/resources/trt_allocator.cc rename to tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc index 7a2e93414aed56525eaeac876cdac20404bcf6ab..1636cdc30c4df157ed124b160449af645f917252 100644 --- a/tensorflow/contrib/tensorrt/resources/trt_allocator.cc +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h" #include "tensorflow/core/platform/logging.h" diff --git a/tensorflow/contrib/tensorrt/resources/trt_allocator.h b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h similarity index 93% rename from tensorflow/contrib/tensorrt/resources/trt_allocator.h rename to tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h index f857a9de055ee7668f0bf9bc97e030354505081b..59ffb42bad348c78cde32035aff8c7081528b3a6 100644 --- a/tensorflow/contrib/tensorrt/resources/trt_allocator.h +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_ALLOCATOR_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_ALLOCATOR_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ALLOCATOR_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ALLOCATOR_H_ #include @@ -81,4 +81,4 @@ class TRTDeviceAllocator : public TRTBaseAllocator { #endif // GOOGLE_TENSORRT #endif // GOOGLE_CUDA -#endif // TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_ALLOCATOR_H_ +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_ALLOCATOR_H_ diff --git a/tensorflow/contrib/tensorrt/resources/trt_allocator_test.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator_test.cc similarity index 98% rename from tensorflow/contrib/tensorrt/resources/trt_allocator_test.cc rename to tensorflow/compiler/tf2tensorrt/utils/trt_allocator_test.cc index beb1284208e4c10ffe1d36ef411cf08f11dbcb78..e457c64928e5df84c7e2726ba3621420f013dbc9 100644 --- a/tensorflow/contrib/tensorrt/resources/trt_allocator_test.cc +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h" #include "tensorflow/core/platform/test.h" diff --git a/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc similarity index 97% rename from tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.cc rename to tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc index dab1dd9343be7d5b033a3e04bf0b49fbbf37e9e5..5213fced1ea9220422245172f5b4a3f584a2a566 100644 --- a/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.cc +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h" #include #include @@ -135,7 +135,7 @@ void TRTInt8Calibrator::setDone() { void TRTInt8Calibrator::writeCalibrationCache(const void* ptr, std::size_t length) { - calibration_table_ = string((const char*)ptr, length); + calibration_table_ = string(static_cast(ptr), length); VLOG(1) << "Got calibration data for " << engine_name_ << " @" << ptr << " length=" << length; } diff --git a/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h similarity index 93% rename from tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h rename to tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h index 65466c9741989fda5f82fc27d813d026f35fe386..10587e99624acfb97730bbbd9dfbcde020ffc669 100644 --- a/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_INT8_CALIBRATOR_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_INT8_CALIBRATOR_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_INT8_CALIBRATOR_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_INT8_CALIBRATOR_H_ #include #include @@ -96,4 +96,4 @@ struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator { #endif #endif -#endif // TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_INT8_CALIBRATOR_H_ +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_INT8_CALIBRATOR_H_ diff --git a/tensorflow/contrib/tensorrt/log/trt_logger.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc similarity index 90% rename from tensorflow/contrib/tensorrt/log/trt_logger.cc rename to tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc index dda0dc9e712eb726800abfb6084f4f708d04825b..6bc842ed5ca7e03018157060a332338cdc926f14 100644 --- a/tensorflow/contrib/tensorrt/log/trt_logger.cc +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/log/trt_logger.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h" #if GOOGLE_CUDA #if GOOGLE_TENSORRT @@ -26,6 +26,9 @@ namespace tensorrt { void Logger::log(Severity severity, const char* msg) { // Suppress info-level messages switch (severity) { +#if NV_TENSORRT_MAJOR > 5 || (NV_TENSORRT_MAJOR == 5 && NV_TENSORRT_MINOR >= 1) + case Severity::kVERBOSE: +#endif case Severity::kINFO: { // Mark TRT info messages as debug! VLOG(2) << name_ << " " << msg; break; diff --git a/tensorflow/contrib/tensorrt/log/trt_logger.h b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h similarity index 86% rename from tensorflow/contrib/tensorrt/log/trt_logger.h rename to tensorflow/compiler/tf2tensorrt/utils/trt_logger.h index 96ccacb791e40143c5c4d9d691bb353702f9a28b..22f4de970a80765b0e1e7e8816134d83aaec7c73 100644 --- a/tensorflow/contrib/tensorrt/log/trt_logger.h +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_logger.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_TENSORRT_LOG_TRT_LOGGER_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_LOG_TRT_LOGGER_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LOGGER_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LOGGER_H_ #include "tensorflow/core/platform/types.h" @@ -41,4 +41,4 @@ class Logger : public nvinfer1::ILogger { #endif // GOOGLE_TENSORRT #endif // GOOGLE_CUDA -#endif // TENSORFLOW_CONTRIB_TENSORRT_LOG_TRT_LOGGER_H_ +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LOGGER_H_ diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h new file mode 100644 index 0000000000000000000000000000000000000000..09c47b36b0ad8074e749342e7d08f139da7ea1f4 --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h @@ -0,0 +1,192 @@ +/* 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_TF2TENSORRT_UTILS_TRT_LRU_CACHE_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LRU_CACHE_H_ + +#include +#include + +#include "tensorflow/compiler/tf2tensorrt/convert/utils.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h" +#include "tensorflow/core/framework/resource_mgr.h" +#include "tensorflow/core/lib/core/errors.h" + +#if GOOGLE_CUDA && GOOGLE_TENSORRT +#include "tensorrt/include/NvInfer.h" +#endif // GOOGLE_CUDA && GOOGLE_TENSORRT + +namespace tensorflow { +namespace tensorrt { + +template +class LRUCache { + public: + typedef Value value_type; + typedef Key key_type; + typedef HashFunction hasher; + typedef typename std::unordered_map map_type; + typedef typename map_type::iterator iterator; + typedef typename map_type::const_iterator const_iterator; + + LRUCache() : capacity_(0) {} + explicit LRUCache(size_t capacity) : capacity_(capacity) {} + + size_t capacity() const { return capacity_; } + + void reserve(size_t capacity) { + capacity_ = capacity; + DiscardOld(); + } + + size_t size() const { return objects_.size(); } + + size_t count(const key_type& key) const { return objects_.count(key); } + + value_type& at(const key_type& key) { return Touch(key); } + + const_iterator begin() const { return objects_.begin(); } + const_iterator end() const { return objects_.end(); } + + iterator begin() { return objects_.begin(); } + iterator end() { return objects_.end(); } + + template + std::pair emplace(Args&&... args) { + DiscardOld(1); + std::pair result = + objects_.emplace(std::forward(args)...); + key_type key = result.first->first; + if (result.second) { + keys_.push_front(key); + } else { + TouchNoCheck(key); // The key must exist in this case. + } + return result; + } + + private: + std::unordered_map objects_; + std::list keys_; + size_t capacity_; + value_type not_found_value_; + + value_type& Touch(const key_type& key) { + // Check that the key exists, and let it return std::out_of_range error if + // not. + value_type& value = objects_.at(key); + TouchNoCheck(key); + return value; + } + + void TouchNoCheck(const key_type& key) { + auto rank = std::find(keys_.begin(), keys_.end(), key); + if (rank != keys_.begin()) { + keys_.erase(rank); + keys_.push_front(key); + } + } + + // Creates n free positions in cache + tensorflow::Status DiscardOld(size_t n = 0) { + if (n > capacity_) { + return tensorflow::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(); + } +}; + +// 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 { + return std::hash()(TensorShapeUtils::ShapeListString(key)); + } +}; + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + +struct EngineContext { + EngineContext() {} // Creates an empty context. + EngineContext( + TrtUniquePtrType&& input_cuda_engine, + TrtUniquePtrType&& input_execution_context) + : cuda_engine(std::move(input_cuda_engine)), + execution_context(std::move(input_execution_context)) {} + + mutex mu; + TrtUniquePtrType cuda_engine; + TrtUniquePtrType execution_context + GUARDED_BY(mu); +}; + +class TRTEngineCacheResource : public tensorflow::ResourceBase { + public: + TRTEngineCacheResource(OpKernelContext* ctx, size_t capacity) + : cache_(capacity) { + auto device = ctx->device(); + auto alloc = device->GetAllocator(tensorflow::AllocatorAttributes()); + if (!alloc) { + LOG(ERROR) << "Can't find device allocator for gpu device " + << device->name(); + allocator_ = nullptr; + } else { + allocator_.reset(new TRTDeviceAllocator(alloc)); + } + } + + string DebugString() const override { + std::stringstream oss; + using std::dec; + using std::endl; + using std::hex; + oss << "TRTEngineCacheResource: "; + oss << "TRTBaseAllocator = " << hex << allocator_.get() << dec << ", "; + oss << "LRUCache = " << hex << &cache_ << dec << endl; + oss << "Containing " << cache_.size() << " entries: " << endl; + for (const auto& item : cache_) { + oss << TensorShapeUtils::ShapeListString(item.first) << ": " << hex + << "ICudaEngine: " << item.second.get()->cuda_engine.get() << ", " + << "IExecutionContext: " << item.second.get()->execution_context.get() + << dec << endl; + } + return oss.str(); + } + + // Keep device allocator for TRT. + std::unique_ptr allocator_; + + // Declare cache after allocator so that it is destroyed before allocator is. + LRUCache, std::unique_ptr, + VectorTensorShapeHasher> + cache_; +}; + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA + +} // namespace tensorrt +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_LRU_CACHE_H_ diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache_test.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..0aa5eb8f7d4ad062c2d8622fa5aa55f823f80dd5 --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache_test.cc @@ -0,0 +1,57 @@ +/* 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/tf2tensorrt/utils/trt_lru_cache.h" + +#include "tensorflow/core/platform/test.h" + +namespace tensorflow { +namespace tensorrt { + +TEST(LRUCacheTest, Basic) { + LRUCache> cache; + cache.reserve(2); + // Insert 10 + cache.emplace(10, 100); + EXPECT_EQ(cache.size(), 1); + EXPECT_EQ(cache.count(10), 1); + EXPECT_EQ(cache.at(10), 100); + EXPECT_EQ(cache.count(100), 0); + // Insert 20 + cache.emplace(20, 200); + EXPECT_EQ(cache.size(), 2); + EXPECT_EQ(cache.count(10), 1); + EXPECT_EQ(cache.count(20), 1); + EXPECT_EQ(cache.at(10), 100); + EXPECT_EQ(cache.at(20), 200); + EXPECT_EQ(cache.count(100), 0); + EXPECT_EQ(cache.count(200), 0); + // Insert 30, Evicting 10 + cache.emplace(30, 300); + EXPECT_EQ(cache.count(10), 0); + EXPECT_EQ(cache.count(20), 1); + EXPECT_EQ(cache.count(30), 1); + // Touch 20 + cache.at(20); + // Insert 40, Evicting 30 + cache.emplace(40, 400); + EXPECT_EQ(cache.count(10), 0); + EXPECT_EQ(cache.count(20), 1); + EXPECT_EQ(cache.count(30), 0); + EXPECT_EQ(cache.count(40), 1); +} + +} // namespace tensorrt +} // namespace tensorflow diff --git a/tensorflow/contrib/tensorrt/resources/trt_resource_manager.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_resource_manager.cc similarity index 96% rename from tensorflow/contrib/tensorrt/resources/trt_resource_manager.cc rename to tensorflow/compiler/tf2tensorrt/utils/trt_resource_manager.cc index 9c3698e5d1cc5d6d8d31a8fcaf03d103f1e1915d..0a72a88bc740101bcbadb40bfe106a5b8d284bbf 100644 --- a/tensorflow/contrib/tensorrt/resources/trt_resource_manager.cc +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_resource_manager.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tensorrt/resources/trt_resource_manager.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_resource_manager.h" #include "tensorflow/core/platform/logging.h" namespace tensorflow { diff --git a/tensorflow/contrib/tensorrt/resources/trt_resource_manager.h b/tensorflow/compiler/tf2tensorrt/utils/trt_resource_manager.h similarity index 87% rename from tensorflow/contrib/tensorrt/resources/trt_resource_manager.h rename to tensorflow/compiler/tf2tensorrt/utils/trt_resource_manager.h index 19f39e6d3db1571573fb290dd2c30fd43ea604ef..03879ffff2fa724b05cb1919753e4aaa99e2e702 100644 --- a/tensorflow/contrib/tensorrt/resources/trt_resource_manager.h +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_resource_manager.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_RESOURCE_MANAGER_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_RESOURCE_MANAGER_H_ +#ifndef TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_RESOURCE_MANAGER_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_RESOURCE_MANAGER_H_ #include #include @@ -42,4 +42,4 @@ class TRTResourceManager { } // namespace tensorrt } // namespace tensorflow -#endif // TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_RESOURCE_MANAGER_H_ +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_RESOURCE_MANAGER_H_ diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_resources.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_resources.cc new file mode 100644 index 0000000000000000000000000000000000000000..2e553079b19a3e5d0739cc6ac79a84f3b6a1fc4e --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_resources.cc @@ -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. +==============================================================================*/ + +#include "tensorflow/compiler/tf2tensorrt/utils/trt_resources.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + +namespace tensorflow { +namespace tensorrt { + +TRTCalibrationResource::~TRTCalibrationResource() { + VLOG(0) << "Destroying Calibration Resource " << std::endl << DebugString(); + builder_.reset(); + engine_.reset(); + // We need to manually destroy the builder and engine before the allocator + // is destroyed. + allocator_.reset(); +} + +string TRTCalibrationResource::DebugString() const { + std::stringstream oss; + using std::dec; + using std::endl; + using std::hex; + oss << " Calibrator = " << hex << calibrator_.get() << dec << endl + << " Builder = " << hex << builder_.get() << dec << endl + << " Engine = " << hex << engine_.get() << dec << endl + << " Logger = " << hex << &logger_ << dec << endl + << " Allocator = " << hex << allocator_.get() << dec << endl + << " Thread = " << hex << thr_.get() << dec << endl; + return oss.str(); +} + +Status TRTCalibrationResource::SerializeToString(string* serialized) { + calibrator_->waitAndSetDone(); + thr_->join(); + *serialized = calibrator_->getCalibrationTableAsString(); + if (serialized->empty()) { + return tensorflow::errors::Unknown("Calibration table is empty."); + } + return Status::OK(); +} + +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_resources.h b/tensorflow/compiler/tf2tensorrt/utils/trt_resources.h new file mode 100644 index 0000000000000000000000000000000000000000..5e8d4b3b738df09b0c2ea82dcc06e9b23a708385 --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_resources.h @@ -0,0 +1,73 @@ +/* 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_TF2TENSORRT_UTILS_TRT_RESOURCES_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_RESOURCES_H_ + +#include +#include +#include +#include +#include +#include + +#include "tensorflow/compiler/tf2tensorrt/convert/utils.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/resource_mgr.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "tensorrt/include/NvInfer.h" + +namespace tensorflow { +namespace tensorrt { + +class SerializableResourceBase : public tensorflow::ResourceBase { + public: + virtual Status SerializeToString(string* serialized) = 0; +}; + +class TRTCalibrationResource : public SerializableResourceBase { + public: + ~TRTCalibrationResource() override; + + string DebugString() const override; + + Status SerializeToString(string* serialized) override; + + // Lookup table for temporary staging areas of input tensors for calibration. + std::unordered_map> device_buffers_; + + // Temporary staging areas for calibration inputs. + std::vector device_tensors_; + + std::unique_ptr calibrator_; + TrtUniquePtrType builder_; + TrtUniquePtrType engine_; + std::unique_ptr allocator_; + tensorflow::tensorrt::Logger logger_; + // TODO(sami): Use threadpool threads! + std::unique_ptr thr_; +}; + +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_UTILS_TRT_RESOURCES_H_ diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index d8123e956fac04912b4fed5bf75cc9cb55c5baf9..5a1a9435c19160cfb8130253a9fa756af423165c 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -1,6 +1,6 @@ licenses(["notice"]) # Apache 2.0 -load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test") +load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test", "tf_cuda_cc_test") package_group( name = "internal", @@ -204,6 +204,7 @@ cc_library( "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client", "//tensorflow/compiler/xla/client:client_library", @@ -224,6 +225,7 @@ cc_library( "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", + "@com_google_absl//absl/types:variant", ], alwayslink = 1, ) @@ -315,11 +317,13 @@ tf_cc_test( ":tf2xla_util", "//tensorflow/cc:cc_ops", "//tensorflow/cc:function_ops", + "//tensorflow/cc:functional_ops", "//tensorflow/cc:ops", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:math_ops_op_lib", + "//tensorflow/core:ops", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", @@ -444,6 +448,7 @@ cc_library( hdrs = [ "dump_graph.h", ], + visibility = ["//visibility:public"], deps = [ "//tensorflow/compiler/jit:flags", "//tensorflow/core:framework", @@ -669,8 +674,31 @@ cc_library( name = "side_effect_util", srcs = ["side_effect_util.cc"], hdrs = ["side_effect_util.h"], + visibility = [":friends"], deps = [ "//tensorflow/core:core_cpu", "@com_google_absl//absl/strings", ], ) + +tf_cuda_cc_test( + name = "fused_batchnorm_reserve_space_test", + size = "medium", + srcs = ["fused_batchnorm_reserve_space_test.cc"], + deps = [ + "//tensorflow/cc:cc_ops", + "//tensorflow/compiler/jit", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:tensorflow", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + "//tensorflow/core/kernels:ops_testutil", + "//tensorflow/core/kernels:ops_util", + "@com_google_absl//absl/algorithm:container", + ], +) diff --git a/tensorflow/compiler/tf2xla/functionalize_cond.cc b/tensorflow/compiler/tf2xla/functionalize_cond.cc index 7ae96e1d484900e28e8c23c3bb2232401144ad82..8aa162be47c9181e215de6a2eb660215135ff6eb 100644 --- a/tensorflow/compiler/tf2xla/functionalize_cond.cc +++ b/tensorflow/compiler/tf2xla/functionalize_cond.cc @@ -34,6 +34,8 @@ limitations under the License. #include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/graph/control_flow.h" #include "tensorflow/core/graph/node_builder.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/lib/strings/strcat.h" using xla::StatusOr; @@ -41,6 +43,43 @@ using xla::StatusOr; namespace tensorflow { namespace functionalize_cond { +bool AncestorNode::operator<(const AncestorNode& other) const { + return (output_tensor.node->id() < other.output_tensor.node->id()) || + (output_tensor.node->id() == other.output_tensor.node->id() && + output_tensor.index < other.output_tensor.index) || + (output_tensor.node->id() == other.output_tensor.node->id() && + output_tensor.index == other.output_tensor.index && + type < other.type); +} + +bool AncestorNode::operator==(const AncestorNode& other) const { + return output_tensor.node->id() == other.output_tensor.node->id() && + output_tensor.index == other.output_tensor.index && type == other.type; +} + +size_t AncestorNode::Hash::operator()(const AncestorNode& ancestor) const { + size_t h = std::hash()(ancestor.output_tensor.node->id()); + h = Hash64Combine(h, std::hash()(ancestor.output_tensor.index)); + return Hash64Combine(h, std::hash()(static_cast(ancestor.type))); +} + +typedef std::tuple + ClusterTuple; + +struct ClusterTupleLessThan { + bool operator()(const ClusterTuple& a, const ClusterTuple& b) const { + if (std::tie(std::get<0>(a), std::get<1>(a)) < + std::tie(std::get<0>(b), std::get<1>(b))) { + return true; + } else if (std::tie(std::get<0>(a), std::get<1>(a)) == + std::tie(std::get<0>(b), std::get<1>(b))) { + return StateMap::OutputTensorLess()(std::get<2>(a), std::get<2>(b)); + } else { + return false; + } + } +}; + // TODO(jpienaar): Move to OutputTensor. string DebugString(const OutputTensor& tensor) { return absl::StrCat(tensor.node->name(), ":", tensor.index); @@ -145,10 +184,10 @@ size_t StateMap::Hash::operator()(const StateMap::AncestorState& map) const { if (map.empty()) return 0; // Compute hash of the front element. auto it = map.begin(); - size_t h = hash()(*it); + size_t h = AncestorNode::Hash()(*it); for (++it; it != map.end(); ++it) { // Combine the has with the different elements in the map. - h = Hash64Combine(h, hash()(*it)); + h = Hash64Combine(h, AncestorNode::Hash()(*it)); } return h; } @@ -229,7 +268,17 @@ string StateMap::CondStateToString(StateMap::CondId id) const { } string StateMap::AncestorStateToString(const Node* node) const { - if (auto id = LookupAncestorId(node)) return NodesToString(*id); + if (auto id = LookupAncestorId(node)) { + return absl::StrCat( + "{", + absl::StrJoin(*id, ",", + [](string* output, const AncestorNode& ancestor) { + absl::StrAppend(output, + ancestor.output_tensor.node->name(), + ":", ancestor.output_tensor.index); + }), + "}"); + } return "{}"; } @@ -247,7 +296,9 @@ class Conditional { Status AddMerge(Node* m); // Constructs an If node from the merge nodes. - Status BuildAndReplace(Graph* graph, FunctionLibraryDefinition* library); + Status BuildAndReplace( + Graph* graph, FunctionLibraryDefinition* library, + std::unordered_map* merge_to_replacement); private: // Extracts the then/else bodies: creates new graphs with the nodes @@ -262,10 +313,15 @@ class Conditional { Status BuildIfNode(Graph* graph, FunctionLibraryDefinition* library); // Adds input edges to If node. - Status AddInputEdges(Graph* graph); + Status AddInputEdges( + Graph* graph, + const std::unordered_map& merge_to_replacement); // Adds output edges from If node. - Status AddOutputEdges(Graph* graph); + // Record new output tensor for all Merge nodes in 'merge_to_replacement'. + Status AddOutputEdges( + Graph* graph, + std::unordered_map* merge_to_replacement); // Adds switch node that is part of this conditional. Status AddSwitch(Node* s); @@ -705,9 +761,9 @@ Status Conditional::BuildIfNode(Graph* graph, } builder.Device(predicate_.node->assigned_device_name()); // Conditional should be the first input ... - builder.Input(NodeDefBuilder::NodeOut(predicate_.node->name(), - predicate_.index, - predicate_.node->output_type(0))); + builder.Input( + NodeDefBuilder::NodeOut(predicate_.node->name(), predicate_.index, + predicate_.node->output_type(predicate_.index))); // ... followed by the other inputs. builder.Input(inputs); @@ -720,12 +776,29 @@ Status Conditional::BuildIfNode(Graph* graph, return Status::OK(); } -Status Conditional::AddInputEdges(Graph* graph) { +Status Conditional::AddInputEdges( + Graph* graph, + const std::unordered_map& merge_to_replacement) { VLOG(2) << "AddInputEdges for " << if_node_->name(); int index = 0; // Add predicate input. - graph->AddEdge(const_cast(predicate_.node), predicate_.index, if_node_, - index++); + if (predicate_.node->IsMerge()) { + // If the predicate is a Merge node, we should not use Merge output as + // predicate. Instead, we should use the corresponding If output in + // 'merge_to_replacement'. Otherwise, this Conditional's If node is still + // connected to the predicate Merge node; and when we call + // DeleteReachableAndDeadNodes(), the predicate Merge node and this + // Conditional's If node will be removed. + auto iter = merge_to_replacement.find(predicate_.node); + if (iter == merge_to_replacement.end()) { + return errors::Internal("Cannot find replacement for Merge node ", + predicate_.node->name()); + } + graph->AddEdge(iter->second.node, iter->second.index, if_node_, index++); + } else { + graph->AddEdge(const_cast(predicate_.node), predicate_.index, + if_node_, index++); + } // Add function body inputs. for (auto& arg : cond_arg_nodes_) { if (arg.src_output == Graph::kControlSlot) { @@ -740,7 +813,9 @@ Status Conditional::AddInputEdges(Graph* graph) { return Status::OK(); } -Status Conditional::AddOutputEdges(Graph* graph) { +Status Conditional::AddOutputEdges( + Graph* graph, + std::unordered_map* merge_to_replacement) { VLOG(2) << "AddOutputEdges for " << if_node_->name(); int i = 0; for (Node* node : merges_) { @@ -764,6 +839,10 @@ Status Conditional::AddOutputEdges(Graph* graph) { graph->AddEdge(if_node_, i, dst, dst_input); } } + + // Record corresponding output tensor in 'merge_to_replacement'. + (*merge_to_replacement)[node] = OutputTensor{if_node_, i}; + ++i; } for (Node* n : external_control_outputs_) { @@ -773,8 +852,9 @@ Status Conditional::AddOutputEdges(Graph* graph) { return Status::OK(); } -Status Conditional::BuildAndReplace(Graph* graph, - FunctionLibraryDefinition* library) { +Status Conditional::BuildAndReplace( + Graph* graph, FunctionLibraryDefinition* library, + std::unordered_map* merge_to_replacement) { VLOG(1) << "Build If and replace merge nodes " << NodesToString(this->merges_); if (replaced_) return Status::OK(); @@ -793,8 +873,8 @@ Status Conditional::BuildAndReplace(Graph* graph, } TF_RETURN_IF_ERROR(BuildIfNode(graph, library)); - TF_RETURN_IF_ERROR(AddInputEdges(graph)); - TF_RETURN_IF_ERROR(AddOutputEdges(graph)); + TF_RETURN_IF_ERROR(AddInputEdges(graph, *merge_to_replacement)); + TF_RETURN_IF_ERROR(AddOutputEdges(graph, merge_to_replacement)); TF_RETURN_IF_ERROR(parent_->PropagateUpdatedState(if_node_)); // Check that the if_node doesn't feed into itself. @@ -936,6 +1016,10 @@ StatusOr FunctionalizeCond::JoinCondStatesMerge( VLOG(4) << "Joining (for merge) " << DebugString(src) << " and " << DebugString(dst); if (state_map_.IsEmpty(dst)) return src; + if (state_map_.IsEmpty(src)) { + return errors::Internal("Merge node ", merge->name(), + " has input that's not in any CondContext."); + } if (state_map_.IsDead(src)) return src; if (state_map_.IsDead(dst)) return dst; @@ -1170,8 +1254,17 @@ Status FunctionalizeCond::DetermineAncestorState(Node* dst) { if (other_id != id && other_id != nullptr) { state.insert(other_id->begin(), other_id->end()); } - if (IsSwitch(src) || IsMerge(src)) { - state.insert(src); + if (IsMerge(src)) { + state.insert({{src, 0}, AncestorNode::AncestorNodeType::kMerge}); + } else if (IsSwitch(src)) { + OutputTensor pred; + // For dead switch nodes, GetSwitchPredicate() will fail, and we use + // the switch node directly as ancestor. + if (GetSwitchPredicate(*src, &pred).ok()) { + state.insert({pred, AncestorNode::AncestorNodeType::kPred}); + } else { + state.insert({{src, 0}, AncestorNode::AncestorNodeType::kSwitch}); + } } return state_map_.GetAncestorId(state); }; @@ -1317,16 +1410,30 @@ Status FunctionalizeCond::FunctionalizeInternal() { // Sort the merge nodes from innermost outwards. SortMergeNodes(&merge_order); - // Cluster merge nodes by CondId and AncestorId in order of nesting. - using ClusterPair = std::pair; + // Cluster merge nodes by (CondId, AncestorId, predicate) in order of + // nesting. (CondId, AncestorId) is not enough, e.g. + // pred1 = array_ops.placeholder(dtypes.bool, name='pred1') + // pred2 = array_ops.placeholder(dtypes.bool, name='pred2') + // cond1 = control_flow_ops.cond(pred1, ...) + // cond2 = control_flow_ops.cond(pred2, ...) + // cond3 = control_flow_ops.cond(pred1, use cond1 and cond2) + // cond4 = control_flow_ops.cond(pred2, use cond1 and cond2) + // cond3 and cond4 have the same (CondId, AncestorId), but they should not + // be merged into one "If" node (because they have different predicates). std::deque> merge_clusters; - std::map merge_cluster_index; + std::map merge_cluster_index; for (Node* merge : merge_order) { auto cond_id = state_map_.LookupCondId(merge); if (state_map_.IsDead(cond_id)) continue; - ClusterPair key = - std::make_pair(cond_id, state_map_.LookupAncestorId(merge)); + auto predicate = merge_to_predicate_.find(merge); + if (predicate == merge_to_predicate_.end()) { + return errors::Internal("Cannot find predicate for Merge node ", + merge->name()); + } + + ClusterTuple key = std::make_tuple( + cond_id, state_map_.LookupAncestorId(merge), predicate->second); auto idx = merge_cluster_index.find(key); if (idx == merge_cluster_index.end()) { merge_cluster_index[key] = merge_clusters.size(); @@ -1345,7 +1452,8 @@ Status FunctionalizeCond::FunctionalizeInternal() { Conditional cond(merge_to_predicate_.at(cluster.front()), this, &state_map_); for (Node* merge : cluster) TF_RETURN_IF_ERROR(cond.AddMerge(merge)); - TF_RETURN_IF_ERROR(cond.BuildAndReplace(graph_, library_)); + TF_RETURN_IF_ERROR( + cond.BuildAndReplace(graph_, library_, &merge_to_replacement_)); if (VLOG_IS_ON(4)) DumpGraphWithCondState("after_extract"); } diff --git a/tensorflow/compiler/tf2xla/functionalize_cond.h b/tensorflow/compiler/tf2xla/functionalize_cond.h index 8525d7af61b4471e53a9ae16b081060bfd234c9c..d85800fb8ee65a354716bf6601c6bc40eca9a10d 100644 --- a/tensorflow/compiler/tf2xla/functionalize_cond.h +++ b/tensorflow/compiler/tf2xla/functionalize_cond.h @@ -43,6 +43,33 @@ enum class BranchType { kNeither = 3, }; +// When we keep track of which switch/merge node's feed into a node, we record +// 1) predicate for non-dead switch node, +// 2) the switch node itself for dead switch node, +// 3) the merge node itself for merge node. +// Case 1) is an optimization. With this optimization, if there are nodes from +// different switch nodes but those switch nodes have the same predicate, the +// nodes will still have same AncestorState, and they will be clustered into a +// single "If". +struct AncestorNode { + enum class AncestorNodeType { + kPred = 0, + kSwitch = 1, + kMerge = 2, + }; + + OutputTensor output_tensor; + AncestorNodeType type; + + // Compare two AncestorNodes by (node id, index, type). + bool operator<(const AncestorNode& other) const; + bool operator==(const AncestorNode& other) const; + + struct Hash { + size_t operator()(const AncestorNode&) const; + }; +}; + // StateMap is responsible for mapping from each graph Node to // * a CondState, where each CondState is a map from predicate to branch (i,e., // what predicates have to hold or not hold). @@ -68,7 +95,7 @@ class StateMap { using CondId = const CondState*; // Keep track of which switch/merge node's feed into a node's values. - using AncestorState = std::set; + using AncestorState = std::set; // Every unique ID is mapped to a AncestorState. using AncestorId = const AncestorState*; @@ -232,6 +259,9 @@ class FunctionalizeCond { // Mapping from merge nodes to predicate. std::unordered_map merge_to_predicate_; + // Mapping from merge nodes to corresponding If node outputs. + std::unordered_map merge_to_replacement_; + FunctionLibraryDefinition* library_; Graph* graph_; diff --git a/tensorflow/compiler/tf2xla/functionalize_cond_test.cc b/tensorflow/compiler/tf2xla/functionalize_cond_test.cc index b0aabd63bbda784b3b7103a438ce025eea0cd93b..05fa1ee92dc172bd11cec9f99e3884996e00791f 100644 --- a/tensorflow/compiler/tf2xla/functionalize_cond_test.cc +++ b/tensorflow/compiler/tf2xla/functionalize_cond_test.cc @@ -101,6 +101,17 @@ TEST_F(FunctionalizeCondTest, JoinCondStates) { TF_EXPECT_OK(t.status()); } +TEST_F(FunctionalizeCondTest, JoinCondStatesMergeWithInputNotInCondContext) { + Tensor val_tensor(DT_INT32, TensorShape()); + val_tensor.flat().setZero(); + Node* val = test::graph::Constant(graph_.get(), val_tensor, "val"); + Node* m = test::graph::Merge(graph_.get(), val, val); + + StateMap::CondState cond_state; + auto joined_or = JoinCondStatesMerge(m, /*src=*/nullptr, &cond_state); + EXPECT_FALSE(joined_or.ok()); +} + } // namespace } // namespace functionalize_cond } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/fused_batchnorm_reserve_space_test.cc b/tensorflow/compiler/tf2xla/fused_batchnorm_reserve_space_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..4535ece374ceb801e450af98a21d5a4c5e8f2a29 --- /dev/null +++ b/tensorflow/compiler/tf2xla/fused_batchnorm_reserve_space_test.cc @@ -0,0 +1,130 @@ +/* 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 +#include + +#include "absl/algorithm/container.h" +#include "tensorflow/cc/ops/array_ops.h" +#include "tensorflow/cc/ops/const_op.h" +#include "tensorflow/cc/ops/nn_ops.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/kernels/ops_testutil.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/public/session.h" + +namespace tensorflow { +namespace { +Status GetTestDevice(Session* session, string* test_device) { + std::vector devices; + TF_RETURN_IF_ERROR(session->ListDevices(&devices)); + + bool found_cpu = absl::c_any_of(devices, [&](const DeviceAttributes& device) { + return device.device_type() == "CPU"; + }); + + bool found_gpu = absl::c_any_of(devices, [&](const DeviceAttributes& device) { + return device.device_type() == "GPU"; + }); + + if (!found_gpu && !found_cpu) { + return errors::Internal("Expected at least one CPU or GPU!"); + } + + *test_device = found_gpu ? "GPU" : "CPU"; + VLOG(2) << "Using test device " << *test_device; + return Status::OK(); +} + +void FillZeros(Tensor* tensor) { + auto flat = tensor->flat(); + for (int i = 0; i < flat.size(); i++) { + flat.data()[i] = 0.0f; + } +} + +// This tests check that the implementation outputs from FusedBatchnorm +// training, reserve_space_{1|2}, are what we assume them to be in the TF/XLA +// lowering. +// +// If this test starts failing then it doesn't indicate that TF/cudnn have +// violated their contract, but it indicates that we need to update the TF/XLA +// lowering for FusedBatchnorm training to match the new implementation defined +// behavior. +TEST(FusedBatchnormReserveSpaceTest, Test) { + using ::tensorflow::ops::Const; + using ::tensorflow::ops::FusedBatchNorm; + + std::unique_ptr session( + tensorflow::NewSession(tensorflow::SessionOptions{})); + + string test_device; + TF_ASSERT_OK(GetTestDevice(session.get(), &test_device)); + + Scope root = tensorflow::Scope::NewRootScope(); + Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT); + + Tensor scale_data(DT_FLOAT, TensorShape({10})); + FillZeros(&scale_data); + Output scale = + Const(root.WithOpName("scale"), Input::Initializer(scale_data)); + + Tensor offset_data(DT_FLOAT, TensorShape({10})); + FillZeros(&offset_data); + Output offset = + Const(root.WithOpName("offset"), Input::Initializer(offset_data)); + + Tensor mean_data(DT_FLOAT, TensorShape({0})); + Output mean = Const(root.WithOpName("offset"), Input::Initializer(mean_data)); + + Tensor variance_data(DT_FLOAT, TensorShape({0})); + Output variance = + Const(root.WithOpName("variance"), Input::Initializer(variance_data)); + + string tf_device = absl::StrCat("/device:", test_device, ":0"); + string xla_device = absl::StrCat("/device:XLA_", test_device, ":0"); + + FusedBatchNorm fused_batch_norm_tf( + root.WithOpName("fused_batch_norm_tf").WithDevice(tf_device), input, + scale, offset, mean, variance, FusedBatchNorm::Attrs{}.IsTraining(true)); + FusedBatchNorm fused_batch_norm_xla( + root.WithOpName("fused_batch_norm_xla").WithDevice(xla_device), input, + scale, offset, mean, variance, FusedBatchNorm::Attrs{}.IsTraining(true)); + + tensorflow::GraphDef graph; + TF_ASSERT_OK(root.ToGraphDef(&graph)); + + TF_ASSERT_OK(session->Create(graph)); + + Tensor input_data(DT_FLOAT, TensorShape({10, 10, 10, 10})); + auto flat_input = input_data.flat(); + for (int i = 0; i < flat_input.size(); i++) { + flat_input.data()[i] = (i - 5) / 1000.0f; + } + + std::vector results; + TF_ASSERT_OK(session->Run({{"input", input_data}}, + {fused_batch_norm_tf.reserve_space_1.name(), + fused_batch_norm_xla.reserve_space_1.name(), + fused_batch_norm_tf.reserve_space_2.name(), + fused_batch_norm_xla.reserve_space_2.name()}, + {}, &results)); + + test::ExpectClose(results[0], results[1], /*atol=*/1e-4); + test::ExpectClose(results[2], results[3], /*atol=*/1e-4); +} +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/graph_compiler.cc b/tensorflow/compiler/tf2xla/graph_compiler.cc index efb75749722893100494e089c0beb96944e9f1d4..5e4699bbb6218089d2e76a36c7351bf7fbd23264 100644 --- a/tensorflow/compiler/tf2xla/graph_compiler.cc +++ b/tensorflow/compiler/tf2xla/graph_compiler.cc @@ -22,6 +22,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/dump_graph.h" #include "tensorflow/compiler/tf2xla/literal_util.h" #include "tensorflow/compiler/tf2xla/shape_util.h" +#include "tensorflow/compiler/tf2xla/side_effect_util.h" #include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_context.h" @@ -88,6 +89,9 @@ Status PrepareArguments(XlaOpKernelContext* ctx, Graph* graph, case XlaExpression::Kind::kResource: return errors::Unimplemented( "Resource as function argument is not yet implemented."); + case XlaExpression::Kind::kTensorList: + return errors::Unimplemented( + "TensorList as function argument is not yet implemented."); case XlaExpression::Kind::kInvalid: return errors::InvalidArgument("Invalid function argument"); } @@ -191,6 +195,9 @@ Status GraphCompiler::CompileFunctionalNode(Node* n, // into the functions. XlaOpKernelContext xla_op_context(op_context); + XlaContext& context = XlaContext::Get(op_context); + auto* b = context.builder(); + XlaCompiler* compiler = xla_op_context.compiler(); NameAttrList func; @@ -219,8 +226,12 @@ Status GraphCompiler::CompileFunctionalNode(Node* n, TF_RETURN_IF_ERROR( PrepareArguments(&xla_op_context, graph.get(), expressions, &arguments)); + bool add_token_input_output = + HasNodeAttr(n->def(), kXlaTokenInputNodesAttrName); + XlaCompiler::CompileOptions compile_options; compile_options.is_entry_computation = false; + compile_options.add_token_input_output = add_token_input_output; XlaCompiler::CompilationResult result; TF_RETURN_IF_ERROR( compiler->CompileFunction(compile_options, func, arguments, &result)); @@ -234,9 +245,19 @@ Status GraphCompiler::CompileFunctionalNode(Node* n, } handles.push_back(expressions[i]->handle()); } - - XlaContext& context = XlaContext::Get(op_context); - auto* b = context.builder(); + if (add_token_input_output) { + std::vector token_input_nodes; + TF_RETURN_IF_ERROR( + GetNodeAttr(n->def(), kXlaTokenInputNodesAttrName, &token_input_nodes)); + std::vector token_inputs; + for (const string& node_name : token_input_nodes) { + auto token_or = compiler->GetNodeToken(node_name); + TF_RETURN_IF_ERROR(token_or.status()); + token_inputs.push_back(token_or.ConsumeValueOrDie()); + } + xla::XlaOp token_input = xla::AfterAll(b, token_inputs); + handles.push_back(token_input); + } auto output_handle = xla::Call(b, *result.computation, handles); // The output handle of `Call` computation is a tuple type. Unzip it so @@ -251,6 +272,10 @@ Status GraphCompiler::CompileFunctionalNode(Node* n, ++computation_output; } } + if (add_token_input_output) { + TF_RETURN_IF_ERROR(compiler->SetNodeToken( + n->name(), xla::GetTupleElement(output_handle, computation_output))); + } return b->first_error(); } diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD index 47209d285f1a077fd80f779a406e6980892f1646..b3f050c52b3a71067a1cc7aa0cd18905e35e4f1c 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -34,6 +34,7 @@ tf_kernel_library( "dynamic_slice_ops.cc", "dynamic_stitch_op.cc", "elu_op.cc", + "empty_op.cc", "extract_image_patches_op.cc", "fake_param_op.cc", "fake_quantize_ops.cc", @@ -143,14 +144,27 @@ tf_kernel_library( "//tensorflow/compiler/xla/client/lib:qr", "//tensorflow/compiler/xla/client/lib:quantize", "//tensorflow/compiler/xla/client/lib:sorting", - "//tensorflow/compiler/xla/client/lib:triangular_solve", + "//tensorflow/core:bitwise_ops_op_lib", + "//tensorflow/core:control_flow_ops_op_lib", + "//tensorflow/core:data_flow_ops_op_lib", "//tensorflow/core:framework", + "//tensorflow/core:functional_ops_op_lib", "//tensorflow/core:image_ops_op_lib", "//tensorflow/core:lib", "//tensorflow/core:linalg_ops_op_lib", + "//tensorflow/core:list_ops_op_lib", + "//tensorflow/core:math_ops_op_lib", + "//tensorflow/core:nn_ops_op_lib", + "//tensorflow/core:no_op_op_lib", "//tensorflow/core:protos_all_cc", + "//tensorflow/core:random_ops_op_lib", + "//tensorflow/core:resource_variable_ops_op_lib", + "//tensorflow/core:sendrecv_ops_op_lib", + "//tensorflow/core:sparse_ops_op_lib", "//tensorflow/core:spectral_ops_op_lib", + "//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", diff --git a/tensorflow/compiler/tf2xla/kernels/arg_op.cc b/tensorflow/compiler/tf2xla/kernels/arg_op.cc index 795ea09831e183a26fb3498b9bbaf9c3adaef9ed..5554d7a377d38554058aa731770ee10e400bc535 100644 --- a/tensorflow/compiler/tf2xla/kernels/arg_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/arg_op.cc @@ -53,7 +53,11 @@ class XlaArgOp : public XlaOpKernel { const XlaExpression& arg = ctx->xla_context()->args()[index_]; OP_REQUIRES(ctx, arg.kind() != XlaExpression::Kind::kInvalid, errors::InvalidArgument("Invalid/missing argument expression")); - ctx->SetOutputExpression(0, arg); + if (ctx->expected_output_dtype(0) == DT_VARIANT) { + ctx->SetTensorListOutput(0, arg.handle()); + } else { + ctx->SetOutputExpression(0, arg); + } } private: @@ -63,6 +67,8 @@ class XlaArgOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(XlaArgOp); }; -REGISTER_XLA_OP(Name("_Arg").AllowResourceTypes().CompilationOnly(), XlaArgOp); +REGISTER_XLA_OP( + Name("_Arg").AllowResourceTypes().AllowVariantTypes().CompilationOnly(), + XlaArgOp); } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc index 0e2f335f3354e3ae6008bdc0ac0b80683fe479c1..f1d78c87527eb5f818dcf92209feabe33653a625 100644 --- a/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc @@ -18,6 +18,8 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.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/core/util/tensor_format.h" @@ -34,6 +36,7 @@ class FusedBatchNormOp : public XlaOpKernel { OP_REQUIRES( ctx, FormatFromString(data_format_str, &data_format_), errors::InvalidArgument("Invalid data format: ", data_format_str)); + is_on_gpu_ = ctx->device_type().type_string() == DEVICE_GPU_XLA_JIT; } void Compile(XlaOpKernelContext* ctx) override { @@ -71,7 +74,18 @@ class FusedBatchNormOp : public XlaOpKernel { // variance to the gradient. Here we maintain the same behavior by setting // them to the mean and variance calculated by BatchNormTraining. ctx->SetOutput(3, xla::GetTupleElement(output, 1)); - ctx->SetOutput(4, xla::GetTupleElement(output, 2)); + if (is_on_gpu_) { + // The last two outputs from the FusedBatchNorm training TensorFlow GPU + // op are implementation defined. For now we rely on the in-practice + // behavior of the op: + // output 3 is the mean + // output 4 is rsqrt(variance + epsilon) + xla::XlaOp variance = xla::GetTupleElement(output, 2); + ctx->SetOutput(4, xla::Rsqrt(xla::Add( + variance, xla::ScalarLike(variance, epsilon_)))); + } else { + ctx->SetOutput(4, xla::GetTupleElement(output, 2)); + } } else { xla::XlaOp output = xla::BatchNormInference( input, ctx->Input(1), ctx->Input(2), ctx->Input(3), ctx->Input(4), @@ -89,6 +103,7 @@ class FusedBatchNormOp : public XlaOpKernel { float epsilon_; TensorFormat data_format_; bool is_training_; + bool is_on_gpu_; }; REGISTER_XLA_OP(Name("FusedBatchNorm"), FusedBatchNormOp); @@ -104,6 +119,7 @@ class FusedBatchNormGradOp : public XlaOpKernel { OP_REQUIRES( ctx, FormatFromString(data_format_str, &data_format_), errors::InvalidArgument("Invalid data format: ", data_format_str)); + is_on_gpu_ = ctx->device_type().type_string() == DEVICE_GPU_XLA_JIT; } void Compile(XlaOpKernelContext* ctx) override { @@ -130,6 +146,22 @@ class FusedBatchNormGradOp : public XlaOpKernel { xla::XlaOp scale_backprop; xla::XlaOp offset_backprop; if (is_training_) { + if (is_on_gpu_) { + // The last two inputs to the FusedBatchNormGrad training TensorFlow GPU + // op are implementation defined. For now we rely on the in-practice + // behavior of the op: input 3 is the mean input 4 is rsqrt(variance + + // epsilon) + // + // The XLA op expects: + // input 3 is the mean + // input 4 is the variance + // + // so we adjust input 4 here. + xla::XlaOp one = xla::ScalarLike(var, 1.0f); + xla::XlaOp epsilon = xla::ScalarLike(var, epsilon_); + var = xla::Sub(one / (var * var), epsilon); + } + xla::XlaOp output = xla::BatchNormGrad(activations, scale, mean, var, grad_backprop, epsilon_, feature_index); @@ -158,9 +190,8 @@ class FusedBatchNormGradOp : public XlaOpKernel { offset_backprop = XlaHelpers::ConvertElementType(reduce, scale_dtype); // scratch1 = rsqrt(pop_var + epsilon) - auto neg_half = XlaHelpers::FloatLiteral(b, scale_dtype, -0.5); - auto scratch1 = xla::Pow( - xla::Add(var, xla::ConstantR0(b, epsilon_)), neg_half); + auto epsilon = XlaHelpers::FloatLiteral(b, scale_dtype, epsilon_); + auto scratch1 = xla::Rsqrt(xla::Add(var, epsilon)); // scratch2 = sum(y_backprop * (x - mean)) auto mul = @@ -187,6 +218,7 @@ class FusedBatchNormGradOp : public XlaOpKernel { TensorFormat data_format_; float epsilon_; bool is_training_; + bool is_on_gpu_; }; REGISTER_XLA_OP(Name("FusedBatchNormGrad"), FusedBatchNormGradOp); diff --git a/tensorflow/compiler/tf2xla/kernels/bias_ops.cc b/tensorflow/compiler/tf2xla/kernels/bias_ops.cc index e7f369b761f36a717ea5fb536780af91a8955b1e..33bdf9aec3167b0277f3c1db18c9e247ed9bb5d1 100644 --- a/tensorflow/compiler/tf2xla/kernels/bias_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/bias_ops.cc @@ -48,8 +48,11 @@ class BiasOp : public XlaOpKernel { OP_REQUIRES(ctx, TensorShapeUtils::IsVector(bias_shape), errors::InvalidArgument("Biases must be 1D: ", bias_shape.DebugString())); - int feature_dim = (data_format_ == FORMAT_NHWC) ? input_shape.dims() - 1 - : input_shape.dims() - 3; + + // feature_dim is the channel (C) dimension of the data. + int feature_dim = (data_format_ == FORMAT_NHWC) + ? input_shape.dims() - 1 + : /*data_format == FORMAT_NCHW*/ 1; OP_REQUIRES( ctx, feature_dim >= 0, errors::InvalidArgument("Input tensor does not have enough dimensions " @@ -91,9 +94,10 @@ class BiasAddGradOp : public XlaOpKernel { errors::InvalidArgument("Input tensor must be at least 2D: ", out_backprop_shape.DebugString())); + // feature_dim is the channel (C) dimension of the data. int feature_dim = (data_format_ == FORMAT_NHWC) ? out_backprop_shape.dims() - 1 - : out_backprop_shape.dims() - 3; + : /*data_format == FORMAT_NCHW*/ 1; OP_REQUIRES( ctx, feature_dim >= 0, errors::InvalidArgument("Input tensor does not have enough dimensions " diff --git a/tensorflow/compiler/tf2xla/kernels/binary_ops.cc b/tensorflow/compiler/tf2xla/kernels/binary_ops.cc index 5e9280c1fe692037b0a842a92ef5a8c28b854a54..ad6b334326a470442c8c0d79b725345d4165be10 100644 --- a/tensorflow/compiler/tf2xla/kernels/binary_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/binary_ops.cc @@ -20,7 +20,9 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/client_library.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/xla_data.pb.h" #include "tensorflow/core/framework/kernel_def_builder.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/types.h" @@ -165,12 +167,8 @@ XLA_MAKE_BINARY( xla::Div(xla::Mul(rhs, XlaHelpers::FloatLiteral(b, input_type(0), 0.5)), lhs, extend_dimensions)); -static xla::XlaOp Square(xla::XlaBuilder* builder, const xla::XlaOp& x) { - return xla::Mul(x, x); -} - XLA_MAKE_BINARY(SquaredDifference, - Square(b, xla::Sub(lhs, rhs, extend_dimensions))); + xla::Square(xla::Sub(lhs, rhs, extend_dimensions))); XLA_MAKE_BINARY(TruncateDiv, xla::Div(lhs, rhs, extend_dimensions)); XLA_MAKE_BINARY(TruncateMod, xla::Rem(lhs, rhs, extend_dimensions)); @@ -195,8 +193,8 @@ XLA_MAKE_BINARY(SoftplusGrad, // softsigngrad(gradients, features) = gradients / (1 + abs(features)) ** 2 XLA_MAKE_BINARY(SoftsignGrad, xla::Div(lhs, - Square(b, xla::Add(XlaHelpers::One(b, input_type(0)), - xla::Abs(rhs))))); + xla::Square(xla::Add(XlaHelpers::One(b, input_type(0)), + xla::Abs(rhs))))); XLA_MAKE_BINARY(TanhGrad, xla::Mul(rhs, xla::Sub(XlaHelpers::One(b, input_type(0)), @@ -204,6 +202,8 @@ XLA_MAKE_BINARY(TanhGrad, XLA_MAKE_BINARY(Pow, xla::Pow(lhs, rhs, extend_dimensions)); +XLA_MAKE_BINARY(NextAfter, xla::NextAfter(lhs, rhs)); + #undef XLA_MAKE_BINARY class ApproximateEqualOp : public XlaOpKernel { diff --git a/tensorflow/compiler/tf2xla/kernels/cast_op.cc b/tensorflow/compiler/tf2xla/kernels/cast_op.cc index 8cc2479dd555380da7500abe6b2aca380110333b..ca2152d6c103e05c06809d85d9529720ff112217 100644 --- a/tensorflow/compiler/tf2xla/kernels/cast_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/cast_op.cc @@ -12,6 +12,7 @@ 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/lib/util.h" #include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" @@ -19,6 +20,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/primitive_util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/framework/kernel_def_builder.h" namespace tensorflow { @@ -31,6 +33,7 @@ class CastOp : public XlaOpKernel { OP_REQUIRES_OK(ctx, ctx->GetAttr("DstT", &dst_dtype_)); OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(src_dtype_, &src_type_)); OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(dst_dtype_, &dst_type_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("Truncate", &use_truncation_)); } void Compile(XlaOpKernelContext* ctx) override { @@ -48,6 +51,36 @@ class CastOp : public XlaOpKernel { // imaginary part. output = xla::ConvertElementType(xla::Real(input), dst_type_); } else { + if (use_truncation_) { + OP_REQUIRES( + ctx, + xla::primitive_util::IsFloatingPointType(src_type_) && + xla::primitive_util::IsFloatingPointType(dst_type_), + errors::Unimplemented("Truncate attribute is only " + "implemented for floating point datatypes.")); + int mantissa_difference = + xla::primitive_util::SignificandWidth(src_type_) - + xla::primitive_util::SignificandWidth(dst_type_); + OP_REQUIRES(ctx, mantissa_difference > 0, + errors::Unimplemented( + "Truncate attribute is only implemented in cases where " + "dst datatype " + "has fewer mantissa bits than the src datatype")); + int src_bitwidth = xla::primitive_util::BitWidth(src_type_); + + // Bitcast to same-width integer, mask off the LSBs, bitcast back to the + // source datatype. + int64 mask = ~((1L << mantissa_difference) - 1); + xla::PrimitiveType same_width_int = + xla::primitive_util::UnsignedIntegralTypeForBitWidth(src_bitwidth); + OP_REQUIRES(ctx, same_width_int != xla::PRIMITIVE_TYPE_INVALID, + errors::Unimplemented("Unexpected type bitwidth")); + input = xla::BitcastConvertType( + xla::And( + xla::BitcastConvertType(input, same_width_int), + ::tensorflow::IntegerLiteral(builder, same_width_int, mask)), + src_type_); + } output = xla::ConvertElementType(input, dst_type_); } @@ -57,6 +90,7 @@ class CastOp : public XlaOpKernel { protected: DataType src_dtype_, dst_dtype_; xla::PrimitiveType src_type_, dst_type_; + bool use_truncation_; TF_DISALLOW_COPY_AND_ASSIGN(CastOp); }; @@ -79,8 +113,8 @@ class BitcastOp : public XlaOpKernel { if (src_dtype_ == dst_dtype_) { output = input; } else { - // The only complex type in XLA is C64, so error out if the bitcast has a - // complex source or destination type and the bitcast is not trivial. + // Error out if the bitcast has a complex source or destination type and + // the bitcast is not trivial. OP_REQUIRES(ctx, !xla::primitive_util::IsComplexType(src_type_) && !xla::primitive_util::IsComplexType(dst_type_), diff --git a/tensorflow/compiler/tf2xla/kernels/categorical_op.cc b/tensorflow/compiler/tf2xla/kernels/categorical_op.cc index 7199b9b6feb36dd45ef51f4c38463bc715fcc38a..c2b4c28d1566f5429c5d8109db94af0c3762b131 100644 --- a/tensorflow/compiler/tf2xla/kernels/categorical_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/categorical_op.cc @@ -99,8 +99,8 @@ class CategoricalOp : public XlaOpKernel { xla::PrimitiveType xla_output_type; OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(output_type(0), &xla_output_type)); - xla::XlaOp argmax = XlaHelpers::ArgMax(softmax_entries, xla_output_type, - /*axis=*/class_dimension); + xla::XlaOp argmax = xla::ArgMax(softmax_entries, xla_output_type, + /*axis=*/class_dimension); if (num_samples == 1) { argmax = xla::Reshape(argmax, {batch_size, 1}); } diff --git a/tensorflow/compiler/tf2xla/kernels/concat_op.cc b/tensorflow/compiler/tf2xla/kernels/concat_op.cc index cd7c7f4a82df7a65829787efcb1fd2f77870e945..91e4d9cea7cbf6075e30250587044174c4b8e7f4 100644 --- a/tensorflow/compiler/tf2xla/kernels/concat_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/concat_op.cc @@ -24,13 +24,13 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/bounds_check.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/const_op.cc b/tensorflow/compiler/tf2xla/kernels/const_op.cc index dff8af800229b9605bb93e0498bc5e5cf012f244..ff6c54e47c62f0555ef045e25051f6ec5a3c1d39 100644 --- a/tensorflow/compiler/tf2xla/kernels/const_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/const_op.cc @@ -83,6 +83,17 @@ class ConstOp : public XlaOpKernel { return; } break; + case DT_COMPLEX128: + if (proto_.scomplex_val_size() == 2) { + ctx->SetOutput( + 0, + xla::Broadcast(xla::ConstantR0( + b, xla::complex128(proto_.dcomplex_val(0), + proto_.dcomplex_val(1))), + shape.dim_sizes())); + return; + } + break; case DT_INT32: if (proto_.int_val_size() == 1) { ctx->SetOutput( diff --git a/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc b/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc index b0bc7640307149459a29e6b0b2e8e8132e4141c9..e8b270c67a23b876612ab1dba92a8ae7a46a392d 100644 --- a/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc +++ b/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc @@ -26,13 +26,13 @@ limitations under the License. #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/core/framework/bounds_check.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/numeric_op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_slice.h" -#include "tensorflow/core/kernels/bounds_check.h" #include "tensorflow/core/kernels/conv_grad_ops.h" #include "tensorflow/core/kernels/ops_util.h" #include "tensorflow/core/util/padding.h" @@ -203,7 +203,8 @@ Status ConvBackpropComputeDimensionsV2XlaShapes( StringPiece label, int num_spatial_dims, const xla::Shape& input_shape, const xla::Shape& filter_shape, const xla::Shape& out_backprop_shape, absl::Span dilations, const std::vector& strides, - Padding padding, TensorFormat data_format, ConvBackpropDimensions* dims) { + Padding padding, TensorFormat data_format, ConvBackpropDimensions* dims, + absl::Span explicit_paddings) { TensorShape input_tensor_shape, filter_tensor_shape, out_backprop_tensor_shape; TF_RETURN_IF_ERROR(XLAShapeToTensorShape(input_shape, &input_tensor_shape)); @@ -212,8 +213,8 @@ Status ConvBackpropComputeDimensionsV2XlaShapes( XLAShapeToTensorShape(out_backprop_shape, &out_backprop_tensor_shape)); return ConvBackpropComputeDimensionsV2( label, num_spatial_dims, input_tensor_shape, filter_tensor_shape, - out_backprop_tensor_shape, dilations, strides, padding, data_format, - dims); + out_backprop_tensor_shape, dilations, strides, padding, explicit_paddings, + data_format, dims); } } // anonymous namespace @@ -227,6 +228,10 @@ xla::StatusOr ConvOpAttrs::Create(int num_spatial_dims, TF_RETURN_IF_ERROR(ctx->GetAttr("dilations", &attrs.dilations)); TF_RETURN_IF_ERROR(ctx->GetAttr("strides", &attrs.strides)); TF_RETURN_IF_ERROR(ctx->GetAttr("padding", &attrs.padding)); + if (attrs.padding == EXPLICIT) { + TF_RETURN_IF_ERROR( + ctx->GetAttr("explicit_paddings", &attrs.explicit_paddings)); + } string data_format; TF_RETURN_IF_ERROR(ctx->GetAttr("data_format", &data_format)); @@ -298,6 +303,11 @@ xla::StatusOr MakeXlaForwardConvOp(StringPiece /*type_string*/, window_strides[i] = attrs.strides.at(dim); rhs_dilation[i] = attrs.dilations.at(dim); + if (attrs.padding == EXPLICIT) { + padding[i] = {attrs.explicit_paddings.at(dim * 2), + attrs.explicit_paddings.at(dim * 2 + 1)}; + } + int64 unused_output_size; TF_RETURN_IF_ERROR(GetWindowedOutputSizeVerboseV2( input_shape.dimensions(dim), filter_shape.dimensions(i), @@ -332,7 +342,7 @@ xla::StatusOr MakeXlaBackpropInputConvOp( TF_RETURN_IF_ERROR(ConvBackpropComputeDimensionsV2XlaShapes( type_string, attrs.num_spatial_dims, input_shape, expanded_filter_shape, out_backprop_shape, attrs.dilations, attrs.strides, attrs.padding, - attrs.data_format, &dims)); + attrs.data_format, &dims, attrs.explicit_paddings)); // The input gradients are computed by a convolution of the output // gradients and the filter, with some appropriate padding. See the @@ -415,7 +425,7 @@ xla::StatusOr MakeXlaBackpropFilterConvOp( TF_RETURN_IF_ERROR(ConvBackpropComputeDimensionsV2XlaShapes( type_string, attrs.num_spatial_dims, activations_shape, expanded_filter_shape, out_backprop_shape, attrs.dilations, attrs.strides, - attrs.padding, attrs.data_format, &dims)); + attrs.padding, attrs.data_format, &dims, attrs.explicit_paddings)); // The activations (inputs) form the LHS of the convolution. // Activations have shape: [batch, in_rows, in_cols, ..., in_depth] @@ -428,23 +438,14 @@ xla::StatusOr MakeXlaBackpropFilterConvOp( int n_dim = GetTensorBatchDimIndex(num_dims, attrs.data_format); int c_dim = GetTensorFeatureDimIndex(num_dims, attrs.data_format); - // The conversion logic below assumes that the data format is NHWC, so we also - // check that here. bool use_batch_group_count = - filter_tensor_shape.dim_size(num_dims - 1) == 1 && attrs.depthwise && - attrs.data_format == FORMAT_NHWC; + filter_tensor_shape.dim_size(num_dims - 1) == 1 && attrs.depthwise; std::vector> padding(attrs.num_spatial_dims); std::vector rhs_dilation(attrs.num_spatial_dims); std::vector window_strides(attrs.num_spatial_dims); std::vector ones(attrs.num_spatial_dims, 1); - // The activations (inputs) form the LHS of the convolution. - // Activations have shape: [batch, in_rows, in_cols, ..., in_depth] - // For the gradient computation, we flip the roles of the batch and - // feature dimensions. - // Each spatial entry has size in_depth * batch - // Swap n_dim and c_dim in the activations. dnums.set_input_batch_dimension(c_dim); dnums.set_input_feature_dimension(n_dim); @@ -473,12 +474,14 @@ xla::StatusOr MakeXlaBackpropFilterConvOp( int64 dim = GetTensorSpatialDimIndex(num_dims, attrs.data_format, i); dnums.add_input_spatial_dimensions(dim); dnums.add_kernel_spatial_dimensions(dim); + rhs_dilation[i] = dims.spatial_dims[i].stride; + window_strides[i] = attrs.dilations[dim]; // We will also need to pad the input with zeros such that after the // convolution, we get the right size for the filter. // The padded_in_rows should be such that when we convolve this with the // expanded_out_rows as a filter, we should get filter_rows back. - // + const int64 padded_in_size = dims.spatial_dims[i].expanded_output_size + (dims.spatial_dims[i].filter_size - 1) * attrs.dilations[dim]; @@ -499,6 +502,8 @@ xla::StatusOr MakeXlaBackpropFilterConvOp( // We apply negative padding in this case. const int64 pad_total = padded_in_size - dims.spatial_dims[i].input_size; + // + For the EXPLICIT padding, we pad the top/left side with the explicit + // padding and pad the bottom/right side with the remaining space. // + For the VALID padding, we don't pad anything on the top/left side // and pad the bottom/right side with the remaining space. // + For the SAME padding, we pad top/left side the same as bottom/right @@ -507,12 +512,12 @@ xla::StatusOr MakeXlaBackpropFilterConvOp( // In addition, if the padded input size is smaller than the input size, // we need to ignore some training elements of the input. We do this by // applying negative padding on the right/bottom. - const int64 pad_before = - attrs.padding == Padding::SAME ? std::max(pad_total / 2, 0) : 0; - + const int64 pad_before = attrs.padding == Padding::EXPLICIT + ? attrs.explicit_paddings[2 * dim] + : attrs.padding == Padding::SAME + ? std::max(pad_total / 2, 0) + : 0; padding[i] = {pad_before, pad_total - pad_before}; - rhs_dilation[i] = dims.spatial_dims[i].stride; - window_strides[i] = attrs.dilations[dim]; } // Besides padding the input, we will also expand output_rows to diff --git a/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.h b/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.h index 6e1b70a47850ae5c05939f8dfb7ec129c031df21..d893eca7f9ba07dded76eb215af4779080fa66b9 100644 --- a/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.h +++ b/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.h @@ -47,6 +47,7 @@ struct ConvOpAttrs { std::vector dilations; std::vector strides; Padding padding; + std::vector explicit_paddings; TensorFormat data_format; }; diff --git a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc index eafdba876ae9e2c38694f065cf83bb3725b8460e..52c3c2c4a903a8c51f6b511774bc0312d39df826 100644 --- a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc @@ -25,13 +25,13 @@ limitations under the License. #include "tensorflow/compiler/xla/client/lib/matrix.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/numeric_op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_slice.h" -#include "tensorflow/core/kernels/bounds_check.h" #include "tensorflow/core/kernels/conv_grad_ops.h" #include "tensorflow/core/kernels/ops_util.h" #include "tensorflow/core/util/padding.h" diff --git a/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc b/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc index 6e6ba21daf5bf3eab5bfc15378e77b6dd253da7c..b119997cf39e210ed8e0ae730a08829e72b238b4 100644 --- a/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc @@ -22,10 +22,10 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/kernels/bounds_check.h" namespace tensorflow { namespace { diff --git a/tensorflow/compiler/tf2xla/kernels/empty_op.cc b/tensorflow/compiler/tf2xla/kernels/empty_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..00d2ce7c12fdc96483612059d1c792c847df04f3 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/empty_op.cc @@ -0,0 +1,66 @@ +/* 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. +==============================================================================*/ + +// XLA-specific Empty Op. + +#include "tensorflow/compiler/tf2xla/type_util.h" +#include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/xla/client/lib/constants.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/core/framework/kernel_def_builder.h" +#include "tensorflow/core/framework/register_types.h" +#include "tensorflow/core/framework/tensor_shape.h" + +namespace tensorflow { +namespace { + +class EmptyOp : public XlaOpKernel { + public: + explicit EmptyOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype_)); + OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(dtype_, &type_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("init", &init_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + // The output of this Op is a tensor of shape 'shape' with each + // element set to the default value of 'dtype'. If 'init' is false then + // the result values may be left undefined, though we don't do that here. + const TensorShape shape_shape = ctx->InputShape("shape"); + OP_REQUIRES( + ctx, TensorShapeUtils::IsVector(shape_shape), + errors::InvalidArgument("shape must be a vector of int32, got shape ", + shape_shape.DebugString())); + + std::vector shape; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector("shape", &shape)); + + auto default_value = xla::Zero(ctx->builder(), type_); + auto result = xla::Broadcast(default_value, shape); + ctx->SetOutput(0, result); + } + + private: + DataType dtype_; + xla::PrimitiveType type_; + bool init_; +}; + +REGISTER_XLA_OP(Name("Empty").CompileTimeConstantInput("shape"), EmptyOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/fft_ops.cc b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc index 6df8b5367d2390e65995beb1583b225755e6ee9f..a623585aad3b1b8f1f096ca527e7694d74f1ba46 100644 --- a/tensorflow/compiler/tf2xla/kernels/fft_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc @@ -21,12 +21,12 @@ limitations under the License. #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/numeric_op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_slice.h" -#include "tensorflow/core/kernels/bounds_check.h" #include "tensorflow/core/kernels/conv_grad_ops.h" #include "tensorflow/core/kernels/ops_util.h" #include "tensorflow/core/util/padding.h" diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op.cc b/tensorflow/compiler/tf2xla/kernels/gather_op.cc index 41c31d0ed58fe9bc9bbde0bd58993c975f04fd60..6472045265e4d930a5da770a68f5c502192201ae 100644 --- a/tensorflow/compiler/tf2xla/kernels/gather_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/gather_op.cc @@ -167,13 +167,13 @@ class GatherOp : public XlaOpKernel { OP_REQUIRES_OK(context, context->ConstantInputAsIntScalar(2, &axis)); const auto params_dims = input_shape.dims(); - if (axis < 0) { - axis += params_dims; - } OP_REQUIRES( - context, 0 <= axis && axis < params_dims, + context, -params_dims <= axis && axis < params_dims, errors::InvalidArgument("Expected axis in the range [", -params_dims, ", ", params_dims, "), but got ", axis)); + if (axis < 0) { + axis += params_dims; + } } DataType index_type = input_type(1); diff --git a/tensorflow/compiler/tf2xla/kernels/identity_op.cc b/tensorflow/compiler/tf2xla/kernels/identity_op.cc index 19dd38c46ef154ea74bcbb6721dd04924702efcc..8b27e8e85a37bd5aa757b0cdd7e00e9fa3c0cf6e 100644 --- a/tensorflow/compiler/tf2xla/kernels/identity_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/identity_op.cc @@ -38,9 +38,13 @@ class IdentityOp : public XlaOpKernel { // XLA_* devices also register a "real" Identity operator so we suppress the // dummy operator using CompilationOnly(). -REGISTER_XLA_OP(Name("Identity").AllowResourceTypes().CompilationOnly(), - IdentityOp); -REGISTER_XLA_OP(Name("IdentityN").AllowResourceTypes().CompilationOnly(), +REGISTER_XLA_OP( + Name("Identity").AllowResourceTypes().AllowVariantTypes().CompilationOnly(), + IdentityOp); +REGISTER_XLA_OP(Name("IdentityN") + .AllowResourceTypes() + .AllowVariantTypes() + .CompilationOnly(), IdentityOp); REGISTER_XLA_OP(Name("PlaceholderWithDefault"), IdentityOp); REGISTER_XLA_OP(Name("PreventGradient"), IdentityOp); diff --git a/tensorflow/compiler/tf2xla/kernels/if_op.cc b/tensorflow/compiler/tf2xla/kernels/if_op.cc index 954ae0b596f33243fad1374473c689adb580f6a4..aa5637e2669555da17af8bb05ab08beeba6a89c3 100644 --- a/tensorflow/compiler/tf2xla/kernels/if_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/if_op.cc @@ -80,7 +80,7 @@ void XlaIfOp::Compile(XlaOpKernelContext* ctx) { arg.name = resource->name(); VLOG(2) << "Resource " << resource->name() << " type: " << DataTypeString(arg.type) - << " shape: " << arg.shape.DebugString() + << " shape: " << arg.HumanString() << " initialized: " << arg.initialized; num_resource_args++; @@ -89,7 +89,7 @@ void XlaIfOp::Compile(XlaOpKernelContext* ctx) { arg.type = input_types_[i]; arg.shape = ctx->InputShape(i + 1); VLOG(2) << "Arg type: " << DataTypeString(arg.type) - << " shape: " << arg.shape.DebugString(); + << " shape: " << arg.HumanString(); } } diff --git a/tensorflow/compiler/tf2xla/kernels/image_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_ops.cc index 96ddd42e2ae04d454e4fb85628d139e17a543d2e..92b20fe0ba5611ca5314cd954026f7b71ea75f84 100644 --- a/tensorflow/compiler/tf2xla/kernels/image_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/image_ops.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "absl/types/span.h" #include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h" #include "tensorflow/compiler/tf2xla/lib/util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" @@ -25,6 +26,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/types.pb.h" namespace tensorflow { namespace { @@ -185,19 +187,20 @@ class AdjustContrastOpV2 : public XlaOpKernel { factor_shape.DebugString())); xla::XlaBuilder* b = context->builder(); - xla::XlaOp input = context->Input(0); - xla::XlaOp factor = context->Input(1); - DataType type = context->input_type(0); + xla::XlaOp input = context->Input(0); + xla::XlaOp factor = XlaHelpers::ConvertElementType(context->Input(1), type); + const DataType accumulation_type = XlaHelpers::SumAccumulationType(type); auto converted = XlaHelpers::ConvertElementType(input, accumulation_type); auto reduce = xla::Reduce(converted, XlaHelpers::Zero(b, accumulation_type), *context->GetOrCreateAdd(accumulation_type), {height_dim, width_dim}); - auto output = XlaHelpers::ConvertElementType(reduce, type); - output = - xla::Div(output, XlaHelpers::FloatLiteral(b, type, height * width)); + + auto output = xla::Div( + reduce, XlaHelpers::FloatLiteral(b, accumulation_type, height * width)); + output = XlaHelpers::ConvertElementType(output, type); std::vector broadcast_dims(input_shape.dims() - 2); std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); @@ -233,8 +236,10 @@ class AdjustSaturationOp : public XlaOpKernel { channels, " channels.")); xla::XlaBuilder* b = context->builder(); - xla::XlaOp input = context->Input(0); - xla::XlaOp scale = context->Input(1); + xla::XlaOp input = + XlaHelpers::ConvertElementType(context->Input(0), DT_FLOAT); + xla::XlaOp scale = + XlaHelpers::ConvertElementType(context->Input(1), DT_FLOAT); DataType type = context->input_type(0); @@ -249,15 +254,17 @@ class AdjustSaturationOp : public XlaOpKernel { /*dimno=*/channel_dim); TensorShape channel_shape = input_shape; channel_shape.set_dim(channel_dim, 1); - auto hsv = RGBToHSV(context, b, {red, green, blue}, context->input_type(0), - channel_shape); + auto hsv = + RGBToHSV(context, b, {red, green, blue}, DT_FLOAT, channel_shape); - hsv[1] = xla::Clamp(XlaHelpers::Zero(b, type), xla::Mul(hsv[1], scale), - XlaHelpers::One(b, type)); + hsv[1] = xla::Clamp(XlaHelpers::Zero(b, DT_FLOAT), xla::Mul(hsv[1], scale), + XlaHelpers::One(b, DT_FLOAT)); - auto rgb = HSVToRGB(context->builder(), hsv, context->input_type(0)); + auto rgb = HSVToRGB(context->builder(), hsv, DT_FLOAT); - context->SetOutput(0, xla::ConcatInDim(b, rgb, channel_dim)); + auto output = XlaHelpers::ConvertElementType( + xla::ConcatInDim(b, rgb, channel_dim), type); + context->SetOutput(0, output); } }; REGISTER_XLA_OP(Name("AdjustSaturation"), AdjustSaturationOp); @@ -283,8 +290,10 @@ class AdjustHueOp : public XlaOpKernel { channels, " channels.")); xla::XlaBuilder* b = context->builder(); - xla::XlaOp input = context->Input(0); - xla::XlaOp delta = context->Input(1); + xla::XlaOp input = + XlaHelpers::ConvertElementType(context->Input(0), DT_FLOAT); + xla::XlaOp delta = + XlaHelpers::ConvertElementType(context->Input(1), DT_FLOAT); DataType type = context->input_type(0); @@ -299,20 +308,22 @@ class AdjustHueOp : public XlaOpKernel { /*dimno=*/channel_dim); TensorShape channel_shape = input_shape; channel_shape.set_dim(channel_dim, 1); - auto hsv = RGBToHSV(context, b, {red, green, blue}, context->input_type(0), - channel_shape); + auto hsv = + RGBToHSV(context, b, {red, green, blue}, DT_FLOAT, channel_shape); - auto zero = XlaHelpers::Zero(b, type); - auto one = XlaHelpers::One(b, type); + auto zero = XlaHelpers::Zero(b, DT_FLOAT); + auto one = XlaHelpers::One(b, DT_FLOAT); auto& hue = hsv[0]; hue = xla::Rem(xla::Add(hsv[0], delta), one); hue = xla::Select(xla::Lt(hue, zero), xla::Rem(xla::Add(one, hue), one), hue); - auto rgb = HSVToRGB(context->builder(), hsv, context->input_type(0)); + auto rgb = HSVToRGB(context->builder(), hsv, DT_FLOAT); - context->SetOutput(0, xla::ConcatInDim(b, rgb, channel_dim)); + auto output = XlaHelpers::ConvertElementType( + xla::ConcatInDim(b, rgb, channel_dim), type); + context->SetOutput(0, output); } }; REGISTER_XLA_OP(Name("AdjustHue"), AdjustHueOp); @@ -351,24 +362,26 @@ struct SuppressBodyFn { auto num_outputs_so_far = values[1]; auto iou_mask = values[2]; auto included_iou = values[3]; - auto zero_r1 = xla::ConstantR1(builder, {0}); + auto zero = xla::ConstantR0(builder, 0); // Determine if current elem is active using a slice. - auto row_idx_r1 = xla::Reshape(row_idx, {1}); - auto active_elem = xla::DynamicSlice(included_iou, row_idx_r1, {1}); + // TODO(b/118437727): The only reason we need an explicit vector is because + // some old GCCs can't deduce the right type for MakeConstSpan, and + // providing a single-value initializer list directly uses the wrong + // overload. Delete this once the deprecated overload is gone. + std::vector row_idx_vector = {row_idx}; + auto active_elem = xla::DynamicSlice(included_iou, row_idx_vector, {1}); active_elem = xla::Reshape(active_elem, {}); // Increment output count iff current elem is not suppressed. num_outputs_so_far = xla::Select( active_elem, num_outputs_so_far + xla::ConstantR0(builder, 1), num_outputs_so_far); // Slice out the row_idx. - auto starts = xla::ConcatInDim(builder, {row_idx_r1, zero_r1}, 0); - auto row_iou = xla::DynamicSlice(iou_mask, starts, {1, num_boxes}); + auto row_iou = xla::DynamicSlice(iou_mask, {row_idx, zero}, {1, num_boxes}); // Remove the diagonal from consideration. An elem cannot suppress // itself. - auto update_starts = xla::ConcatInDim(builder, {zero_r1, row_idx_r1}, 0); row_iou = xla::DynamicUpdateSlice( row_iou, xla::ConstantR2FromArray2D(builder, {{false}}), - update_starts); + {zero, row_idx}); // Create a suppression by inverting polarity. row_iou = xla::Reshape(row_iou, {num_boxes}); auto supp_mask = xla::Not(row_iou); diff --git a/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc index 5a10c52ba8b6d4fab73f0dda67cbd52fd625e76b..b96d45316f626e678a64392a4315979eeeb6e83c 100644 --- a/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc @@ -72,10 +72,10 @@ namespace { // from in_size to out_size. struct ResizeConvolutionDims { // Size of the kernel to use. - std::vector kernel_size; + std::vector kernel_size; // k // Stride of the convolution to use. - std::vector stride; + std::vector stride; // S }; ResizeConvolutionDims ComputeResizeConvolutionParameters( absl::Span in_size, absl::Span out_size, @@ -117,8 +117,10 @@ ResizeConvolutionDims ComputeResizeConvolutionParameters( // + dims.stride * (out_size - 1) int64 CalculateUpperPadding(int64 in_size, int64 out_size, int64 kernel_size, int64 stride) { - return (2 * kernel_size - 1) + (out_size - 1) * stride - (kernel_size - 1) - - 1 - (kernel_size * (in_size - 1)); + int64 padding = (2 * kernel_size - 1) + (out_size - 1) * stride - + (kernel_size - 1) - 1 - (kernel_size * (in_size - 1)); + + return padding; } // Form a 2D convolution kernel like: @@ -132,7 +134,7 @@ int64 CalculateUpperPadding(int64 in_size, int64 out_size, int64 kernel_size, // If the 2D kernel would be very large, the 1D kernel can be applied once in // each dimension due to the symmetry of the kernel along all axis to reduce the // computational intensity. -xla::XlaOp Make1DKernel(xla::XlaBuilder* builder, int64 n) { +xla::XlaOp MakeBilinear1DKernel(xla::XlaBuilder* builder, int64 n) { std::vector kernel(n * 2 - 1); for (int64 i = 0; i < n; ++i) { float v = (i + 1.0f) / n; @@ -142,43 +144,64 @@ xla::XlaOp Make1DKernel(xla::XlaBuilder* builder, int64 n) { return xla::ConstantR1(builder, kernel); } +// Unlike the bilinear kernel, which is triangular, the nearest neighbor +// kernel is a square. For example, a 1D kernel with n=3 would look like +// [0 1 1 1 0] +// and n=4 would look like +// [0 0 1 1 1 1 0]. +// Note that in the second case, the kernel is not symmetric and we default +// to the right (because an existing non TPU kernel +// for nearest neighbor resize already chose to default to the right, +// so we want to be consistent). +xla::XlaOp MakeNearestNeighbor1DKernel(xla::XlaBuilder* builder, int64 n) { + std::vector kernel(n * 2 - 1, 0.0f); + std::fill(&kernel[n / 2], &kernel[(3 * n) / 2], 1.0f); + + return xla::ConstantR1(builder, kernel); +} + // Kernels with more than 16 spatial elements are considered intense and the -// kernel should applied to each dimension independently. +// kernel should be applied to each dimension independently. const int64 kMax2DKernelSize = 16; -xla::XlaOp MakeBilinearResizeKernel(xla::XlaBuilder* builder, - absl::Span kernel_size, - int64 channels) { +xla::XlaOp MakeGeneralResizeKernel(xla::XlaBuilder* builder, + absl::Span kernel_size, + int64 channels, bool is_kernel_bilinear) { + auto make_kernel_func = + is_kernel_bilinear ? MakeBilinear1DKernel : MakeNearestNeighbor1DKernel; + auto depthwise_kernel = xla::Broadcast( xla::Zero(builder, xla::F32), {(2 * kernel_size[0] - 1), (2 * kernel_size[1] - 1), channels, 1}); return xla::Mul( - xla::Add(depthwise_kernel, Make1DKernel(builder, kernel_size[1]), + xla::Add(depthwise_kernel, make_kernel_func(builder, kernel_size[1]), /*broadcast_dimensions=*/{1}), - Make1DKernel(builder, kernel_size[0]), + make_kernel_func(builder, kernel_size[0]), /*broadcast_dimensions=*/{0}); } -xla::XlaOp MakeBilinearResizeKernelInDim(xla::XlaBuilder* builder, - absl::Span kernel_size, - int64 channels, int64 dim) { +xla::XlaOp MakeGeneralResizeKernelInDim(xla::XlaBuilder* builder, + absl::Span kernel_size, + int64 channels, int64 dim, + bool is_kernel_bilinear) { + auto make_kernel_func = + is_kernel_bilinear ? MakeBilinear1DKernel : MakeNearestNeighbor1DKernel; + auto depthwise_kernel = xla::Broadcast(xla::Zero(builder, xla::F32), {dim == 0 ? (2 * kernel_size[0] - 1) : 1, dim == 1 ? (2 * kernel_size[1] - 1) : 1, channels, 1}); - return xla::Add(depthwise_kernel, Make1DKernel(builder, kernel_size[dim]), + return xla::Add(depthwise_kernel, make_kernel_func(builder, kernel_size[dim]), /*broadcast_dimensions=*/{dim}); } -xla::XlaOp ResizeUsingDilationAndConvolution(xla::XlaBuilder* builder, - const xla::XlaOp& input, - const int num_spatial_dims, - std::vector in_size, - std::vector out_size, - const int64 channels, - const bool align_corners) { - // Picture for a 1x3 to 1x4 resize: +xla::XlaOp ResizeUsingDilationAndConvolution( + xla::XlaBuilder* builder, const xla::XlaOp& input, + const int num_spatial_dims, std::vector in_size, + std::vector out_size, const int64 channels, const bool align_corners, + bool is_kernel_bilinear) { + // Picture for a 1x3 to 1x4 bilinear resize: // stride = 2, kernel size = 3 // Input: // 3 6 9 @@ -264,8 +287,8 @@ xla::XlaOp ResizeUsingDilationAndConvolution(xla::XlaBuilder* builder, // Split convolutions into independent dimensions if they would be a very // large kernel. if (dims.kernel_size[0] * dims.kernel_size[1] < kMax2DKernelSize) { - xla::XlaOp kernel = - MakeBilinearResizeKernel(builder, dims.kernel_size, channels); + xla::XlaOp kernel = MakeGeneralResizeKernel(builder, dims.kernel_size, + channels, is_kernel_bilinear); output = xla::ConvGeneralDilated(input_data, kernel, dims.stride, /*padding=*/ @@ -275,8 +298,8 @@ xla::XlaOp ResizeUsingDilationAndConvolution(xla::XlaBuilder* builder, /*rhs_dilation=*/{1, 1}, dimension_numbers, /*feature_group_count=*/channels); } else { - xla::XlaOp kernel0 = - MakeBilinearResizeKernelInDim(builder, dims.kernel_size, channels, 0); + xla::XlaOp kernel0 = MakeGeneralResizeKernelInDim( + builder, dims.kernel_size, channels, 0, is_kernel_bilinear); output = xla::ConvGeneralDilated( input_data, kernel0, {dims.stride[0], 1}, /*padding=*/ @@ -284,8 +307,8 @@ xla::XlaOp ResizeUsingDilationAndConvolution(xla::XlaBuilder* builder, /*lhs_dilation=*/{dims.kernel_size[0], 1}, /*rhs_dilation=*/{1, 1}, dimension_numbers, /*feature_group_count=*/channels); - xla::XlaOp kernel1 = - MakeBilinearResizeKernelInDim(builder, dims.kernel_size, channels, 1); + xla::XlaOp kernel1 = MakeGeneralResizeKernelInDim( + builder, dims.kernel_size, channels, 1, is_kernel_bilinear); output = xla::ConvGeneralDilated( output, kernel1, {1, dims.stride[1]}, /*padding=*/ @@ -306,13 +329,11 @@ xla::XlaOp ResizeUsingDilationAndConvolution(xla::XlaBuilder* builder, return output; } -xla::XlaOp ResizeUsingDilationAndConvolutionGradOp(xla::XlaBuilder* builder, - const xla::XlaOp& grad, - const int num_spatial_dims, - std::vector in_size, - std::vector grad_size, - const int64 channels, - const bool align_corners) { +xla::XlaOp ResizeUsingDilationAndConvolutionGradOp( + xla::XlaBuilder* builder, const xla::XlaOp& grad, + const int num_spatial_dims, std::vector in_size, + std::vector grad_size, const int64 channels, + const bool align_corners, bool is_kernel_bilinear) { ResizeConvolutionDims dims = ComputeResizeConvolutionParameters(in_size, grad_size, align_corners); @@ -332,8 +353,8 @@ xla::XlaOp ResizeUsingDilationAndConvolutionGradOp(xla::XlaBuilder* builder, dimension_numbers.set_kernel_output_feature_dimension(num_spatial_dims); xla::XlaOp output; if (dims.kernel_size[0] * dims.kernel_size[1] < kMax2DKernelSize) { - xla::XlaOp kernel = - MakeBilinearResizeKernel(builder, dims.kernel_size, channels); + xla::XlaOp kernel = MakeGeneralResizeKernel(builder, dims.kernel_size, + channels, is_kernel_bilinear); // Broadcast the input kernel where the forward op expanded from a size == 1 // dimension to a size > 1 dimension. This has the effect of summing the @@ -355,14 +376,14 @@ xla::XlaOp ResizeUsingDilationAndConvolutionGradOp(xla::XlaBuilder* builder, /*rhs_dilation=*/{1, 1}, dimension_numbers, /*feature_group_count=*/channels); } else { - xla::XlaOp kernel0 = - MakeBilinearResizeKernelInDim(builder, dims.kernel_size, channels, 0); - xla::XlaOp kernel1 = - MakeBilinearResizeKernelInDim(builder, dims.kernel_size, channels, 1); - - // Broadcast the input kernel where the forward op expanded from a size == 1 - // dimension to a size > 1 dimension. This has the effect of summing the - // gradient contributions in that dimension. + xla::XlaOp kernel0 = MakeGeneralResizeKernelInDim( + builder, dims.kernel_size, channels, 0, is_kernel_bilinear); + xla::XlaOp kernel1 = MakeGeneralResizeKernelInDim( + builder, dims.kernel_size, channels, 1, is_kernel_bilinear); + + // Broadcast the input kernel where the forward op expanded from a + // size == 1 dimension to a size > 1 dimension. This has the effect of + // summing the gradient contributions in that dimension. if (in_size[0] == 1 && grad_size[0] > 1) { kernel0 = xla::Add(kernel0, xla::ConstantR1(builder, grad_size[0], 0), @@ -407,109 +428,139 @@ xla::XlaOp ResizeUsingDilationAndConvolutionGradOp(xla::XlaBuilder* builder, return output; } -class ResizeBilinearOp : public XlaOpKernel { - public: - explicit ResizeBilinearOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { - OP_REQUIRES_OK(ctx, ctx->GetAttr("align_corners", &align_corners_)); +void GeneralCompile(XlaOpKernelContext* ctx, bool align_corners_, + bool is_kernel_bilinear) { + xla::XlaBuilder* b = ctx->builder(); + + TensorShape input_shape = ctx->InputShape(0); + OP_REQUIRES(ctx, input_shape.dims() == 4, + errors::InvalidArgument("input must be 4-dimensional", + input_shape.DebugString())); + // First dimension always assumed to be batch + const int64 batch = input_shape.dim_size(0); + std::vector in_size = {input_shape.dim_size(1), + input_shape.dim_size(2)}; + // Last/4th dimension always assumed to be num channels + const int64 channels = input_shape.dim_size(3); + OP_REQUIRES(ctx, in_size[0] > 0 && in_size[1] > 0, + errors::InvalidArgument("input size must be positive, got [", + in_size[0], ",", in_size[1], "]")); + + std::vector out_size; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &out_size)); + OP_REQUIRES(ctx, out_size.size() == 2, + errors::InvalidArgument("output size must be length 2, got ", + out_size.size())); + OP_REQUIRES(ctx, out_size[0] > 0 && out_size[1] > 0, + errors::InvalidArgument("output size must be positive, got [", + out_size[0], ",", out_size[1], "]")); + + const int num_spatial_dims = 2; + + xla::XlaOp input = ctx->Input(0); + + // If in_size[i] > 1 and out_size[i] == 1, slice out the first input in + // dimension i. + bool slice_input = false; + for (int i = 0; i < num_spatial_dims; ++i) { + if (in_size[i] > 1 && out_size[i] == 1) { + // If in_size[i] > 1 but out_size[i] == 1, then we slice out the first + // entry before resizing. + slice_input = true; + in_size[i] = 1; + } + } + if (slice_input) { + input = xla::Slice(input, {0, 0, 0, 0}, + {batch, in_size[0], in_size[1], channels}, {1, 1, 1, 1}); } - void Compile(XlaOpKernelContext* ctx) override { - xla::XlaBuilder* b = ctx->builder(); + // Output is always type float. + input = xla::ConvertElementType(input, xla::F32); + + // Special Case: + // Instead of doing a ResizeUsingDilationAndConvolution directly, + // while (out_size[0]-1) = c * 2^x * (in_size[0]-1) for x>1 c>1, resize the + // image to 2*(in_size[0]-1)+1 x-times and then resize by scale c(int here). + // Instead of resizing directly we resize it iteratively. + // + // Since bilinear resize can be broken down as 2 sequential linear + // operations along different dimensions. + // Given sufficient numerical stability and a cxd is same as resizing axb -> exf -> cxd. + // This does not work in the case of align_corners_=false because of special + // padding requirements that cause multiple resizes to be very different + // from a single resize. + // + // This makes the convolutions kernels smaller and the operation faster. + xla::XlaOp output = input; + while (in_size != out_size) { + if (in_size[0] != 1 && in_size[1] != 1) { + std::vector k = { + (static_cast(out_size[0]) - 1) / ((in_size[0] - 1) * 2), + (static_cast(out_size[1]) - 1) / ((in_size[1] - 1) * 2)}; + if ((k[0] == std::floor(k[0])) && (k[1] == std::floor(k[1])) && + k[0] > 1 && k[1] > 1 && align_corners_) { + std::vector next_out_size = {(in_size[0] - 1) * 2 + 1, + (in_size[1] - 1) * 2 + 1}; + output = ResizeUsingDilationAndConvolution( + b, input, num_spatial_dims, in_size, next_out_size, channels, + align_corners_, is_kernel_bilinear); + input = output; + in_size = next_out_size; + } else { + output = ResizeUsingDilationAndConvolution( + b, input, num_spatial_dims, in_size, out_size, channels, + align_corners_, is_kernel_bilinear); + in_size = out_size; + } + } else { + output = ResizeUsingDilationAndConvolution( + b, input, num_spatial_dims, in_size, out_size, channels, + align_corners_, is_kernel_bilinear); + in_size = out_size; + } + } - TensorShape input_shape = ctx->InputShape(0); - OP_REQUIRES(ctx, input_shape.dims() == 4, - errors::InvalidArgument("input must be 4-dimensional", - input_shape.DebugString())); - const int64 batch = input_shape.dim_size(0); - std::vector in_size = {input_shape.dim_size(1), - input_shape.dim_size(2)}; - const int64 channels = input_shape.dim_size(3); - OP_REQUIRES(ctx, in_size[0] > 0 && in_size[1] > 0, - errors::InvalidArgument("input size must be positive, got [", - in_size[0], ",", in_size[1], "]")); + ctx->SetOutput(0, output); +} - std::vector out_size; - OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &out_size)); - OP_REQUIRES(ctx, out_size.size() == 2, - errors::InvalidArgument("output size must be length 2, got ", - out_size.size())); - OP_REQUIRES(ctx, out_size[0] > 0 && out_size[1] > 0, - errors::InvalidArgument("output size must be positive, got [", - out_size[0], ",", out_size[1], "]")); +class ResizeNearestNeighborOp : public XlaOpKernel { + public: + explicit ResizeNearestNeighborOp(OpKernelConstruction* ctx) + : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("align_corners", &align_corners_)); + OP_REQUIRES( + ctx, align_corners_ == true, + errors::Unimplemented("ResizeNearestNeighbor with align_corners=False " + "is not yet implemented")); + } - const int num_spatial_dims = 2; + void Compile(XlaOpKernelContext* ctx) override { + GeneralCompile(ctx, align_corners_, is_kernel_bilinear_); + } - xla::XlaOp input = ctx->Input(0); + private: + bool align_corners_ = true; + bool is_kernel_bilinear_ = false; +}; - // If in_size[i] > 1 and out_size[i] == 1, slice out the first input in - // dimension i. - bool slice_input = false; - for (int i = 0; i < num_spatial_dims; ++i) { - if (in_size[i] > 1 && out_size[i] == 1) { - // If in_size[i] > 1 but out_size[i] == 1, then we slice out the first - // entry before resizing. - slice_input = true; - in_size[i] = 1; - } - } - if (slice_input) { - input = - xla::Slice(input, {0, 0, 0, 0}, - {batch, in_size[0], in_size[1], channels}, {1, 1, 1, 1}); - } +REGISTER_XLA_OP(Name("ResizeNearestNeighbor").CompileTimeConstantInput("size"), + ResizeNearestNeighborOp); - // Output is always type float. - input = xla::ConvertElementType(input, xla::F32); - - // Special Case: - // Instead of doing a ResizeUsingDilationAndConvolution directly, - // while (out_size[0]-1) = c * 2^x * (in_size[0]-1) for x>1 c>1, resize the - // image to 2*(in_size[0]-1)+1 x-times and then resize by scale c(int here). - // Instead of resizing directly we resize it iteratively. - // - // Since bilinear resize can be broken down as 2 sequential linear - // operations along different dimensions. - // Given sufficient numerical stability and a cxd is same as resizing axb -> exf -> cxd. - // This does not work in the case of align_corners_=false because of special - // padding requirements that cause multiple resizes to be very different - // from a single resize. - // - // This makes the convolutions kernels smaller and the operation faster. - xla::XlaOp output = input; - while (in_size != out_size) { - if (in_size[0] != 1 && in_size[1] != 1) { - std::vector k = { - (static_cast(out_size[0]) - 1) / ((in_size[0] - 1) * 2), - (static_cast(out_size[1]) - 1) / ((in_size[1] - 1) * 2)}; - if ((k[0] == std::floor(k[0])) && (k[1] == std::floor(k[1])) && - k[0] > 1 && k[1] > 1 && align_corners_) { - std::vector next_out_size = {(in_size[0] - 1) * 2 + 1, - (in_size[1] - 1) * 2 + 1}; - output = ResizeUsingDilationAndConvolution(b, input, num_spatial_dims, - in_size, next_out_size, - channels, align_corners_); - input = output; - in_size = next_out_size; - } else { - output = ResizeUsingDilationAndConvolution(b, input, num_spatial_dims, - in_size, out_size, - channels, align_corners_); - in_size = out_size; - } - } else { - output = ResizeUsingDilationAndConvolution(b, input, num_spatial_dims, - in_size, out_size, channels, - align_corners_); - in_size = out_size; - } - } +class ResizeBilinearOp : public XlaOpKernel { + public: + explicit ResizeBilinearOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("align_corners", &align_corners_)); + } - ctx->SetOutput(0, output); + void Compile(XlaOpKernelContext* ctx) override { + GeneralCompile(ctx, align_corners_, is_kernel_bilinear_); } private: - bool align_corners_; + bool align_corners_ = true; + bool is_kernel_bilinear_ = true; }; REGISTER_XLA_OP(Name("ResizeBilinear").CompileTimeConstantInput("size"), @@ -581,19 +632,19 @@ class ResizeBilinearGradOp : public XlaOpKernel { (in_size[1] - 1) * 2 + 1}; output = ResizeUsingDilationAndConvolutionGradOp( b, grad, num_spatial_dims, in_size, next_grad_size, channels, - align_corners_); + align_corners_, true); grad = output; in_size = next_grad_size; } else { output = ResizeUsingDilationAndConvolutionGradOp( b, grad, num_spatial_dims, in_size, grad_size, channels, - align_corners_); + align_corners_, true); in_size = grad_size; } } else { output = ResizeUsingDilationAndConvolutionGradOp( b, grad, num_spatial_dims, in_size, grad_size, channels, - align_corners_); + align_corners_, true); in_size = grad_size; } } diff --git a/tensorflow/compiler/tf2xla/kernels/index_ops.cc b/tensorflow/compiler/tf2xla/kernels/index_ops.cc index 843b6bb4e658af16fd753c1a20b35dd3d18df027..c1539f48d4f729510b2d930de91666a7c31f1ef0 100644 --- a/tensorflow/compiler/tf2xla/kernels/index_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/index_ops.cc @@ -18,17 +18,16 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/kernels/index_ops.h" #include "tensorflow/compiler/tf2xla/type_util.h" -#include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/kernel_def_builder.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" -#include "tensorflow/core/kernels/bounds_check.h" namespace tensorflow { XlaArgMinMaxOp::XlaArgMinMaxOp(OpKernelConstruction* ctx, bool is_min) @@ -66,9 +65,9 @@ void XlaArgMinMaxOp::Compile(XlaOpKernelContext* ctx) { xla::XlaOp input = ctx->Input(0); xla::XlaOp output; if (is_min_) { - output = XlaHelpers::ArgMin(input, index_xla_type, axis); + output = xla::ArgMin(input, index_xla_type, axis); } else { - output = XlaHelpers::ArgMax(input, index_xla_type, axis); + output = xla::ArgMax(input, index_xla_type, axis); } ctx->SetOutput(0, output); diff --git a/tensorflow/compiler/tf2xla/kernels/index_ops_cpu.cc b/tensorflow/compiler/tf2xla/kernels/index_ops_cpu.cc index 3e7e8eae6ed406003be2843305bb6b173845d6cc..e4bbdef6480104a1051acfc647644deb65c80171 100644 --- a/tensorflow/compiler/tf2xla/kernels/index_ops_cpu.cc +++ b/tensorflow/compiler/tf2xla/kernels/index_ops_cpu.cc @@ -16,16 +16,16 @@ limitations under the License. // Native XLA implementations of indexing ops. #include "tensorflow/compiler/tf2xla/type_util.h" -#include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/kernel_def_builder.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" -#include "tensorflow/core/kernels/bounds_check.h" namespace tensorflow { namespace { @@ -74,7 +74,7 @@ class ArgMaxCustomCallOp : public XlaOpKernel { // shape isn't supported. if (!ctx->compiler()->options().allow_cpu_custom_calls || (input_dims != 1 && input_dims != 2)) { - xla::XlaOp output = XlaHelpers::ArgMax(ctx->Input(0), output_type, axis); + xla::XlaOp output = xla::ArgMax(ctx->Input(0), output_type, axis); ctx->SetOutput(0, output); return; } diff --git a/tensorflow/compiler/tf2xla/kernels/matmul_op.cc b/tensorflow/compiler/tf2xla/kernels/matmul_op.cc index 6440770c29894c951f010f6c1deb929f4fe79bbf..f36e0025250b3a196b31755a1ddf6620c415b6a3 100644 --- a/tensorflow/compiler/tf2xla/kernels/matmul_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/matmul_op.cc @@ -24,8 +24,8 @@ limitations under the License. namespace tensorflow { namespace { -constexpr std::array kMatmulTypes = { - {DT_HALF, DT_BFLOAT16, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64}}; +constexpr std::array kMatmulTypes = { + {DT_HALF, DT_BFLOAT16, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64, DT_COMPLEX128}}; class MatMulOp : public XlaOpKernel { public: diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc index 90c0ebefb24ec2c4378782e9b15d3f57c33032a4..5a6569c8954d1686dc9d7577a66feb720241ea13 100644 --- a/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_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/triangular_solve.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" namespace tensorflow { namespace { @@ -31,7 +32,10 @@ class MatrixTriangularSolveOp : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { auto result = xla::TriangularSolve( ctx->Input(0), ctx->Input(1), /*left_side=*/true, - /*lower=*/lower_, /*transpose_a=*/adjoint_, /*conjugate_a=*/adjoint_); + /*lower=*/lower_, /*unit_diagonal=*/false, + /*transpose_a=*/ + adjoint_ ? xla::TriangularSolveOptions::ADJOINT + : xla::TriangularSolveOptions::NO_TRANSPOSE); ctx->SetOutput(0, result); } diff --git a/tensorflow/compiler/tf2xla/kernels/pack_op.cc b/tensorflow/compiler/tf2xla/kernels/pack_op.cc index a9b519d8928cc2807831fd6b4f12e60b7d58ea55..426a0941df57f19072d1cb9f3fa3d0079db465c5 100644 --- a/tensorflow/compiler/tf2xla/kernels/pack_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/pack_op.cc @@ -24,12 +24,12 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/bounds_check.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/pooling_ops.cc b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc index 06c6cc37ec90192486ba15010bfeb763a9ffb987..23bb050a34d9246cdf73090aa6adfca054bf8bcf 100644 --- a/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc @@ -26,10 +26,10 @@ limitations under the License. #include "tensorflow/compiler/xla/client/xla_computation.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/kernels/bounds_check.h" #include "tensorflow/core/kernels/conv_grad_ops.h" #include "tensorflow/core/kernels/pooling_ops_common.h" diff --git a/tensorflow/compiler/tf2xla/kernels/random_ops.cc b/tensorflow/compiler/tf2xla/kernels/random_ops.cc index 2d92056e4f522f6206e7d632f0fa1e8b793fd6e3..01b047f732f0e9fb3b45b272e7886e2f8cf4fff4 100644 --- a/tensorflow/compiler/tf2xla/kernels/random_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/random_ops.cc @@ -160,17 +160,24 @@ class RandomShuffleOp : public XlaOpKernel { -> xla::StatusOr> { auto swaps = loop_vars[0]; auto indices = loop_vars[1]; - i = xla::Reshape(i, {1}); + // TODO(b/118437727): The absl::Span nonsense is only necessary because + // the deprecated overload creates ambiguity for the single-element span + // case. Remove it once the deprecated overload is gone. // temp = indices[i] - auto temp = xla::DynamicSlice(indices, i, {1}); + auto temp = + xla::DynamicSlice(indices, absl::Span({i}), {1}); // swap_index = swaps[i] - auto swap_index = xla::DynamicSlice(swaps, i, {1}); + auto swap_index = xla::Reshape( + xla::DynamicSlice(swaps, absl::Span({i}), {1}), {}); // swap_value = indices[swaps[i]] - auto swap_value = xla::DynamicSlice(indices, swap_index, {1}); + auto swap_value = xla::DynamicSlice( + indices, absl::Span({swap_index}), {1}); // indices[i] = indices[swaps[i]] - indices = xla::DynamicUpdateSlice(indices, swap_value, i); + indices = xla::DynamicUpdateSlice(indices, swap_value, + absl::Span({i})); // indices[swaps[i]] = temp - indices = xla::DynamicUpdateSlice(indices, temp, swap_index); + indices = xla::DynamicUpdateSlice( + indices, temp, absl::Span({swap_index})); return std::vector{swaps, indices}; }; // for i in range(n): diff --git a/tensorflow/compiler/tf2xla/kernels/retval_op.cc b/tensorflow/compiler/tf2xla/kernels/retval_op.cc index e4046c795577983bff1a8053743bf4d3a258e583..1f417037284c87753b219ea5ce1d4edce0ce6336 100644 --- a/tensorflow/compiler/tf2xla/kernels/retval_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/retval_op.cc @@ -37,10 +37,14 @@ class RetvalOp : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { const Tensor& input = ctx->op_kernel_context()->input(0); - OP_REQUIRES(ctx, input.dtype() == dtype_, - errors::InvalidArgument( - "Type mismatch: actual ", DataTypeString(input.dtype()), - " vs. expect ", DataTypeString(dtype_))); + // DT_VARIANT types represent Tensor Lists and are wrapped in a DT_UINT8 + // tensor so we skip the check here. + if (dtype_ != DT_VARIANT) { + OP_REQUIRES(ctx, input.dtype() == dtype_, + errors::InvalidArgument( + "Type mismatch: actual ", DataTypeString(input.dtype()), + " vs. expect ", DataTypeString(dtype_))); + } auto frame = ctx->call_frame(); if (frame) { // If 'frame' is non-null, this is an inner function call inside a JIT @@ -59,8 +63,9 @@ class RetvalOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(RetvalOp); }; -REGISTER_XLA_OP(Name("_Retval").AllowResourceTypes().CompilationOnly(), - RetvalOp); +REGISTER_XLA_OP( + Name("_Retval").AllowResourceTypes().AllowVariantTypes().CompilationOnly(), + RetvalOp); } // anonymous namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/scan_ops.cc b/tensorflow/compiler/tf2xla/kernels/scan_ops.cc index 4b9e1a578be2445091228953df7e5c5e82b42c28..daefdfc58a4957d9e685d25aa90da6218f2041ad 100644 --- a/tensorflow/compiler/tf2xla/kernels/scan_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/scan_ops.cc @@ -23,13 +23,13 @@ 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/bounds_check.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/partial_tensor_shape.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/bounds_check.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/select_op.cc b/tensorflow/compiler/tf2xla/kernels/select_op.cc index 9e4c57c9bf73369662274f6b783418e18ff860c2..aaf8c6075dd292e33e70683774a6c1bf374183e3 100644 --- a/tensorflow/compiler/tf2xla/kernels/select_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/select_op.cc @@ -20,8 +20,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/xla_builder.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/kernel_def_builder.h" -#include "tensorflow/core/kernels/bounds_check.h" namespace tensorflow { namespace { diff --git a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc index b1fa2915d59e4e5e2f2523e20e9a37898d087117..7a620d2a6518f8686ef570b33aac971d1dccb6c1 100644 --- a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc @@ -157,9 +157,11 @@ class LinSpaceOp : public XlaOpKernel { flat(0) = start; } else { const float step = (stop - start) / (num - 1); - for (int64 i = 0; i < num; ++i) { + for (int64 i = 0; i < num - 1; ++i) { flat(i) = start + step * i; } + // The last value in the sequence must be equal to stop. + flat(num - 1) = stop; } break; } @@ -171,9 +173,11 @@ class LinSpaceOp : public XlaOpKernel { flat(0) = start; } else { const double step = (stop - start) / (num - 1); - for (int64 i = 0; i < num; ++i) { + for (int64 i = 0; i < num - 1; ++i) { flat(i) = start + step * i; } + // The last value in the sequence must be equal to stop. + flat(num - 1) = stop; } break; } diff --git a/tensorflow/compiler/tf2xla/kernels/shape_op.cc b/tensorflow/compiler/tf2xla/kernels/shape_op.cc index 12830816ec16c9797f0fe4d8f3f13f5a8176161d..31d4cc131600f360c764ffa02831046c85d846e5 100644 --- a/tensorflow/compiler/tf2xla/kernels/shape_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/shape_op.cc @@ -20,10 +20,11 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/xla/client/lib/constants.h" #include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/kernel_def_builder.h" #include "tensorflow/core/framework/tensor_shape.h" -#include "tensorflow/core/kernels/bounds_check.h" namespace tensorflow { namespace { @@ -91,14 +92,20 @@ class SizeOp : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { const TensorShape input_shape = ctx->InputShape(0); - const int64 size = input_shape.num_elements(); - OP_REQUIRES(ctx, FastBoundsCheck(size, std::numeric_limits::max()), + OP_REQUIRES(ctx, + FastBoundsCheck(input_shape.num_elements(), + std::numeric_limits::max()), errors::InvalidArgument("Size does not work for tensors > " "int32 max.")); Tensor size_constant(DT_INT32, TensorShape({})); - size_constant.scalar()() = static_cast(size); - - ctx->SetConstantOutput(0, size_constant); + const int rank = input_shape.dims(); + xla::XlaBuilder* builder = ctx->builder(); + auto size = xla::One(builder, xla::U32); + for (int64 i = 0; i < rank; ++i) { + size = xla::Mul(size, xla::GetDimensionSize(ctx->Input(0), i)); + } + size = xla::ConvertElementType(size, xla::S32); + ctx->SetOutput(0, size); } }; diff --git a/tensorflow/compiler/tf2xla/kernels/shape_util.cc b/tensorflow/compiler/tf2xla/kernels/shape_util.cc index 76ea5f525598f511f295eb5a30f3cf603fbf57aa..b18e3f965c427aec456ce2b188dad79485df23cc 100644 --- a/tensorflow/compiler/tf2xla/kernels/shape_util.cc +++ b/tensorflow/compiler/tf2xla/kernels/shape_util.cc @@ -17,7 +17,7 @@ limitations under the License. #include -#include "tensorflow/core/kernels/bounds_check.h" +#include "tensorflow/core/framework/bounds_check.h" namespace tensorflow { diff --git a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc index 02d71a394280a67cc264450021f91ba475aef7fc..b6c96b1f582710e1cc39e6e1e0e800ef8170743d 100644 --- a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc @@ -24,13 +24,13 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/partial_tensor_shape.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/bounds_check.h" #include "tensorflow/core/kernels/concat_lib.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" @@ -146,9 +146,9 @@ class StackPushOp : public XlaOpKernel { xla::XlaOp value = ctx->Input(1); // start_indices of the DynamicUpdateSlice are [index, 0, 0, ..., 0]. - auto start_indices = - xla::Pad(xla::Reshape(index, {1}), xla::ConstantR0(b, 0), - xla::MakeEdgePaddingConfig({{0, elem_shape.dims()}})); + std::vector start_indices(elem_shape.dims() + 1, + xla::ConstantR0(b, 0)); + start_indices[0] = index; TensorShape slice_shape = elem_shape; slice_shape.InsertDim(0, 1LL); @@ -202,9 +202,9 @@ class StackPopOp : public XlaOpKernel { OP_REQUIRES_OK(ctx, resource->SetValue(xla::Tuple(b, {ta, index}))); // start_indices of the DynamicSlice are [index, 0, 0, ..., 0]. - auto start_indices = - xla::Pad(xla::Reshape(index, {1}), xla::ConstantR0(b, 0), - xla::MakeEdgePaddingConfig({{0, stack_shape.dims() - 1}})); + std::vector start_indices(stack_shape.dims(), + xla::ConstantR0(b, 0)); + start_indices[0] = index; auto slice_shape = stack_shape.dim_sizes(); slice_shape[0] = 1LL; diff --git a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc index 10d990b3213ab882cf44a4df20a977633de3fdab..e8846fbe88fa2a75244398ef0f601fd74e80ec50 100644 --- a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc @@ -288,19 +288,21 @@ class StridedSliceAssignOp : public XlaOpKernel { xla::XlaOp rhs = ctx->Input(4); absl::InlinedVector dimensions_to_reverse; - absl::InlinedVector slice_begin, slice_dims; + absl::InlinedVector slice_begin; + absl::InlinedVector slice_dims; for (int i = 0; i < begin.size(); ++i) { - // TODO(phawkins): implement strides != 1 + // TODO(b/121179231): implement strides != 1 OP_REQUIRES( ctx, strides[i] == 1 || strides[i] == -1, errors::Unimplemented("Strides != 1 or -1 are not yet implemented")); if (strides[i] > 0) { - slice_begin.push_back(begin[i]); + slice_begin.push_back(xla::ConstantR0(ctx->builder(), begin[i])); slice_dims.push_back(end[i] - begin[i]); } else { // Negative stride: swap begin and end, add 1 because the interval // is semi-open, and mark the dimension to be reversed. - slice_begin.push_back(end[i] + 1); + slice_begin.push_back( + xla::ConstantR0(ctx->builder(), end[i] + 1)); slice_dims.push_back(begin[i] - end[i]); dimensions_to_reverse.push_back(i); } @@ -311,14 +313,7 @@ class StridedSliceAssignOp : public XlaOpKernel { } rhs = xla::Reshape(rhs, slice_dims); - if (lhs_shape.dims() == 0) { - // TODO(b/38323843): DynamicUpdateSlice crashes on rank 0 inputs. Fix - // and remove this workaround. - lhs = rhs; - } else { - lhs = xla::DynamicUpdateSlice( - lhs, rhs, xla::ConstantR1(ctx->builder(), slice_begin)); - } + lhs = xla::DynamicUpdateSlice(lhs, rhs, slice_begin); OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, dtype_, lhs)); } diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc index 939d7e19515a1cb41e3e23e9d1fa957ae09ecab7..77a3e5c001e1c715f23ae5148f94dae2faa81acf 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc @@ -27,13 +27,13 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_resource.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/partial_tensor_shape.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/bounds_check.h" #include "tensorflow/core/kernels/concat_lib.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" @@ -123,7 +123,8 @@ Status GetTensorArrayShape(const XlaResource* resource, xla::XlaOp DynamicAddSlice(xla::XlaBuilder* builder, const xla::XlaOp& operand, const xla::XlaOp& update, absl::Span update_dims, - const xla::XlaOp& start_indices, DataType dtype) { + absl::Span start_indices, + DataType dtype) { xla::XlaOp current = xla::DynamicSlice(operand, start_indices, update_dims); xla::XlaOp sum = dtype == DT_BOOL ? xla::Or(current, update) : xla::Add(current, update); @@ -212,9 +213,9 @@ class TensorArrayWriteOp : public XlaOpKernel { xla::XlaOp flow = ctx->Input(3); // start_indices of the DynamicUpdateSlice are [index, 0, 0, ..., 0]. - auto start_indices = - xla::Pad(xla::Reshape(index, {1}), xla::ConstantR0(b, 0), - xla::MakeEdgePaddingConfig({{0, elem_shape.dims()}})); + std::vector start_indices(elem_shape.dims() + 1, + xla::ConstantR0(b, 0)); + start_indices[0] = index; TensorShape slice_shape = elem_shape; slice_shape.InsertDim(0, 1LL); @@ -263,9 +264,9 @@ class TensorArrayReadOp : public XlaOpKernel { xla::XlaOp index = ctx->Input(1); // start_indices of the DynamicSlice are [index, 0, 0, ..., 0]. - auto start_indices = - xla::Pad(xla::Reshape(index, {1}), xla::ConstantR0(b, 0), - xla::MakeEdgePaddingConfig({{0, ta_shape.dims() - 1}})); + std::vector start_indices(ta_shape.dims(), + xla::ConstantR0(b, 0)); + start_indices[0] = index; auto slice_shape = ta_shape.dim_sizes(); slice_shape[0] = 1LL; @@ -419,10 +420,10 @@ class TensorArrayScatterOp : public XlaOpKernel { auto slice = xla::Slice(value, value_starts, value_ends, value_strides); // start_indices of the DynamicUpdateSlice are [index, 0, 0, ..., 0]. - auto index = xla::Slice(indices, {i}, {i + 1}, {1}); - auto start_indices = - xla::Pad(xla::Reshape(index, {1}), xla::ConstantR0(b, 0), - xla::MakeEdgePaddingConfig({{0, elem_shape.dims()}})); + auto index = xla::Reshape(xla::Slice(indices, {i}, {i + 1}, {1}), {}); + std::vector start_indices(elem_shape.dims() + 1, + xla::ConstantR0(b, 0)); + start_indices[0] = index; ta = DynamicAddSlice(b, ta, slice, slice_dims, start_indices, dtype_); } } diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc index 747f591a1b1bea37981c8f292564b64f2785849d..65020012283d9c5f62e5e2fd11fc2bf1110e019a 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc @@ -14,6 +14,9 @@ limitations under the License. ==============================================================================*/ // XLA TensorList operators. +// Tensor lists are represented as tuple consisting of a pre-allocated list +// consisting of the tensors (and where dim 0 is the list index), along with a +// scalar telling us the current number of elements. #include #include @@ -24,13 +27,13 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/partial_tensor_shape.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/bounds_check.h" #include "tensorflow/core/kernels/concat_lib.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" @@ -50,6 +53,22 @@ Status GetTensorListShape(xla::XlaBuilder* builder, xla::XlaOp op, tensor_list_shape); } +class TensorListLengthOp : public XlaOpKernel { + public: + explicit TensorListLengthOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + + void Compile(XlaOpKernelContext* ctx) override { + xla::XlaOp tl = ctx->Input(0); + xla::XlaOp index = xla::GetTupleElement(tl, 1); + ctx->SetOutput(0, index); + } + + private: + TF_DISALLOW_COPY_AND_ASSIGN(TensorListLengthOp); +}; + +REGISTER_XLA_OP(Name("TensorListLength"), TensorListLengthOp); + class TensorListReserveOp : public XlaOpKernel { public: explicit TensorListReserveOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { @@ -67,9 +86,10 @@ class TensorListReserveOp : public XlaOpKernel { tensor_shape.AppendShape(element_shape); xla::XlaBuilder* b = ctx->builder(); - ctx->SetOutput(0, xla::Tuple(b, {xla::Broadcast(XlaHelpers::Zero(b, dtype_), - tensor_shape.dim_sizes()), - xla::ConstantR0(b, 0)})); + ctx->SetTensorListOutput( + 0, xla::Tuple(b, {xla::Broadcast(XlaHelpers::Zero(b, dtype_), + tensor_shape.dim_sizes()), + xla::ConstantR0(b, num_elements)})); } private: @@ -85,19 +105,41 @@ REGISTER_XLA_OP(Name("TensorListReserve") class EmptyTensorListOp : public XlaOpKernel { public: - explicit EmptyTensorListOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + explicit EmptyTensorListOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("element_dtype", &dtype_)); + } void Compile(XlaOpKernelContext* ctx) override { - ctx->CtxFailure( + TensorShape element_shape; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsShape(0, &element_shape)); + int64 max_num_elements; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(1, &max_num_elements)); + OP_REQUIRES( + ctx, max_num_elements >= 0, errors::InvalidArgument("XLA compilation requires a fixed tensor list " - "size. Use TensorListReserve instead.")); + "size. Set the max number of elements.")); + + TensorShape tensor_shape; + tensor_shape.AddDim(max_num_elements); + tensor_shape.AppendShape(element_shape); + + xla::XlaBuilder* b = ctx->builder(); + ctx->SetTensorListOutput( + 0, xla::Tuple(b, {xla::Broadcast(XlaHelpers::Zero(b, dtype_), + tensor_shape.dim_sizes()), + xla::ConstantR0(b, 0)})); } private: + DataType dtype_; + TF_DISALLOW_COPY_AND_ASSIGN(EmptyTensorListOp); }; -REGISTER_XLA_OP(Name("EmptyTensorList"), EmptyTensorListOp); +REGISTER_XLA_OP(Name("EmptyTensorList") + .CompileTimeConstantInput("element_shape") + .CompileTimeConstantInput("max_num_elements"), + EmptyTensorListOp); class TensorListElementShapeOp : public XlaOpKernel { public: @@ -139,6 +181,136 @@ class TensorListElementShapeOp : public XlaOpKernel { REGISTER_XLA_OP(Name("TensorListElementShape"), TensorListElementShapeOp); +class TensorListGetItemOp : public XlaOpKernel { + public: + explicit TensorListGetItemOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("element_dtype", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::XlaBuilder* b = ctx->builder(); + xla::XlaOp state = ctx->Input(0); + + TensorShape shape; + OP_REQUIRES_OK(ctx, GetTensorListShape(b, state, &shape)); + + xla::XlaOp ta = xla::GetTupleElement(state, 0); + xla::XlaOp index = ctx->Input(1); + + // start_indices of the DynamicSlice are [index, 0, 0, ..., 0]. + std::vector start_indices(shape.dims(), + xla::ConstantR0(b, 0)); + start_indices[0] = index; + auto slice_shape = shape.dim_sizes(); + slice_shape[0] = 1LL; + + xla::XlaOp read = xla::DynamicSlice(ta, start_indices, slice_shape); + // Remove the leading '1' dimension. + std::vector value_shape(slice_shape.begin() + 1, slice_shape.end()); + + ctx->SetOutput(0, xla::Reshape(read, value_shape)); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(TensorListGetItemOp); +}; + +REGISTER_XLA_OP(Name("TensorListGetItem"), TensorListGetItemOp); + +class TensorListStackOp : public XlaOpKernel { + public: + explicit TensorListStackOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("element_dtype", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::XlaOp state = ctx->Input(0); + xla::XlaOp ta = xla::GetTupleElement(state, 0); + ctx->SetOutput(0, ta); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(TensorListStackOp); +}; + +REGISTER_XLA_OP(Name("TensorListStack"), TensorListStackOp); + +class TensorListFromTensorOp : public XlaOpKernel { + public: + explicit TensorListFromTensorOp(OpKernelConstruction* ctx) + : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("element_dtype", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + TensorShape element_shape; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsShape(1, &element_shape)); + + const TensorShape tensor_shape = ctx->InputShape(0); + OP_REQUIRES(ctx, tensor_shape.dims() > 0, + errors::InvalidArgument("Input value must be at least a " + "vector but received shape: ", + tensor_shape.DebugString())); + const int num_elements = tensor_shape.dim_size(0); + + xla::XlaBuilder* b = ctx->builder(); + const xla::XlaOp tensor = ctx->Input(0); + + ctx->SetTensorListOutput( + 0, xla::Tuple(b, {tensor, xla::ConstantR0(b, num_elements)})); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(TensorListFromTensorOp); +}; + +REGISTER_XLA_OP( + Name("TensorListFromTensor").CompileTimeConstantInput("element_shape"), + TensorListFromTensorOp); + +class TensorListSetItemOp : public XlaOpKernel { + public: + explicit TensorListSetItemOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("element_dtype", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::XlaBuilder* b = ctx->builder(); + xla::XlaOp tl = ctx->Input(0); + TensorShape elem_shape = ctx->InputShape(2); + + xla::XlaOp ta = xla::GetTupleElement(tl, 0); + xla::XlaOp index = ctx->Input(1); + xla::XlaOp value = ctx->Input(2); + + // start_indices of the DynamicUpdateSlice are [index, 0, 0, ..., 0]. + std::vector start_indices(elem_shape.dims() + 1, + xla::ConstantR0(b, 0)); + start_indices[0] = index; + + TensorShape slice_shape = elem_shape; + slice_shape.InsertDim(0, 1LL); + auto update = xla::Reshape(value, slice_shape.dim_sizes()); + + ctx->SetTensorListOutput( + 0, xla::Tuple(b, {xla::DynamicUpdateSlice(ta, update, start_indices), + index + xla::ConstantR0(b, 1)})); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(TensorListSetItemOp); +}; + +REGISTER_XLA_OP(Name("TensorListSetItem"), TensorListSetItemOp); + class TensorListPushBackOp : public XlaOpKernel { public: explicit TensorListPushBackOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { @@ -147,25 +319,23 @@ class TensorListPushBackOp : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { xla::XlaBuilder* b = ctx->builder(); - xla::XlaOp list = ctx->Input(0); + xla::XlaOp tl = ctx->Input(0); TensorShape elem_shape = ctx->InputShape(1); - xla::XlaOp ta = xla::GetTupleElement(list, 0); - xla::XlaOp index = xla::GetTupleElement(list, 1); + xla::XlaOp ta = xla::GetTupleElement(tl, 0); + xla::XlaOp index = xla::GetTupleElement(tl, 1); xla::XlaOp value = ctx->Input(1); // start_indices of the DynamicUpdateSlice are [index, 0, 0, ..., 0]. - auto start_indices = - xla::Pad(xla::Reshape(index, {1}), xla::ConstantR0(b, 0), - xla::MakeEdgePaddingConfig({{0, elem_shape.dims()}})); + std::vector start_indices(elem_shape.dims() + 1, + xla::ConstantR0(b, 0)); + start_indices[0] = index; TensorShape slice_shape = elem_shape; slice_shape.InsertDim(0, 1LL); auto update = xla::Reshape(value, slice_shape.dim_sizes()); - // TODO(phawkins): We don't check the index is in bounds --- there is no - // error mechanism in XLA. - ctx->SetOutput( + ctx->SetTensorListOutput( 0, xla::Tuple(b, {xla::DynamicUpdateSlice(ta, update, start_indices), index + xla::ConstantR0(b, 1)})); } @@ -197,20 +367,17 @@ class TensorListPopBackOp : public XlaOpKernel { index = index - xla::ConstantR0(b, 1); // start_indices of the DynamicSlice are [index, 0, 0, ..., 0]. - auto start_indices = - xla::Pad(xla::Reshape(index, {1}), xla::ConstantR0(b, 0), - xla::MakeEdgePaddingConfig({{0, shape.dims() - 1}})); - + std::vector start_indices(shape.dims(), + xla::ConstantR0(b, 0)); + start_indices[0] = index; auto slice_shape = shape.dim_sizes(); slice_shape[0] = 1LL; - // TODO(phawkins): We don't check the index is in bounds --- there is no - // error mechanism in XLA. xla::XlaOp read = xla::DynamicSlice(ta, start_indices, slice_shape); // Remove the leading '1' dimension. std::vector value_shape(slice_shape.begin() + 1, slice_shape.end()); - ctx->SetOutput(0, xla::Tuple(b, {ta, index})); + ctx->SetTensorListOutput(0, xla::Tuple(b, {ta, index})); ctx->SetOutput(1, xla::Reshape(read, value_shape)); } diff --git a/tensorflow/compiler/tf2xla/kernels/transpose_op.cc b/tensorflow/compiler/tf2xla/kernels/transpose_op.cc index c9b324a243e4cc3ec64daa3ca0d285336a0d0154..4ac714306248302242902f20d45d2609ef2c7cd3 100644 --- a/tensorflow/compiler/tf2xla/kernels/transpose_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/transpose_op.cc @@ -19,14 +19,15 @@ limitations under the License. // helper. #include "tensorflow/core/kernels/transpose_op.h" +#include "tensorflow/compiler/tf2xla/lib/scatter.h" #include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.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/framework/bounds_check.h" #include "tensorflow/core/framework/kernel_def_builder.h" #include "tensorflow/core/framework/register_types.h" -#include "tensorflow/core/kernels/bounds_check.h" namespace tensorflow { namespace { @@ -128,29 +129,46 @@ class InvertPermutationOp : public XlaOpKernel { errors::InvalidArgument("permutation of nonnegative int32s " "must have <= int32 max elements")); - std::vector perm; - OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(0, &perm)); - - int size = perm.size(); + auto e = ctx->InputExpression(0); + auto tensor_or_status = e.ResolveConstant(ctx->compiler()->client()); + OP_REQUIRES_OK(ctx, tensor_or_status.status()); + // If the input is a constant, we also want the output to be a constant. + // Some models rely on the result of InvertPermutation being a constant. + // TODO(b/32495713): Remove this when we can check whether Scatter is + // constant. Right now, we always assume it is non-constant because we don't + // check the embedded computation. + if (tensor_or_status.ValueOrDie().has_value()) { + std::vector perm; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(0, &perm)); + + int size = perm.size(); + + std::vector output(size); + std::fill_n(output.data(), size, -1); + for (int i = 0; i < size; ++i) { + const int64 d = perm[i]; + OP_REQUIRES(ctx, FastBoundsCheck(d, size), + errors::InvalidArgument(d, " is not between 0 and ", size)); + OP_REQUIRES(ctx, output[d] == -1, + errors::InvalidArgument(d, " is duplicated in the input.")); + output[d] = i; + } - std::vector output(size); - std::fill_n(output.data(), size, -1); - for (int i = 0; i < size; ++i) { - const int64 d = perm[i]; - OP_REQUIRES(ctx, FastBoundsCheck(d, size), - errors::InvalidArgument(d, " is not between 0 and ", size)); - OP_REQUIRES(ctx, output[d] == -1, - errors::InvalidArgument(d, " is duplicated in the input.")); - output[d] = i; + ctx->SetOutput(0, xla::ConstantR1(ctx->builder(), output)); + } else { + auto indices = ctx->Input(0); + int size = ctx->InputShape(0).num_elements(); + auto iota = xla::Iota(ctx->builder(), xla::S32, size); + auto result = XlaScatter(iota, iota, indices, + /*indices_are_vectors=*/false, /*combiner=*/{}, + ctx->builder()); + OP_REQUIRES_OK(ctx, result.status()); + ctx->SetOutput(0, result.ValueOrDie()); } - - ctx->SetOutput(0, xla::ConstantR1(ctx->builder(), output)); } }; -REGISTER_XLA_OP(Name("InvertPermutation") - .TypeConstraint("T", DT_INT32) - .CompileTimeConstantInput("x"), +REGISTER_XLA_OP(Name("InvertPermutation").TypeConstraint("T", DT_INT32), InvertPermutationOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc index a0ea6422d732b00fc1b8cf855d9c9ad603b87c82..62b5cd32da59063f8ce07119fd085f91ec3a1bc4 100644 --- a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc @@ -65,11 +65,8 @@ XLAJIT_MAKE_UNARY(Exp, xla::Exp(x)); XLAJIT_MAKE_UNARY(Expm1, xla::Expm1(x)); XLAJIT_MAKE_UNARY(Floor, xla::Floor(x)); XLAJIT_MAKE_UNARY(IsFinite, xla::IsFinite(x)); -XLAJIT_MAKE_UNARY( - IsInf, - xla::Eq(xla::Abs(x), - xla::ScalarLike(x, std::numeric_limits::infinity()))); -XLAJIT_MAKE_UNARY(IsNan, xla::Ne(x, x)); +XLAJIT_MAKE_UNARY(IsInf, xla::IsInf(x)); +XLAJIT_MAKE_UNARY(IsNan, xla::IsNan(x)); // Return 1/x XLAJIT_MAKE_UNARY(Inv, xla::ScalarLike(x, 1.0) / x); XLAJIT_MAKE_UNARY(Reciprocal, xla::ScalarLike(x, 1.0) / x); @@ -92,8 +89,9 @@ xla::XlaOp Sigmoid(xla::XlaOp x) { } XLAJIT_MAKE_UNARY(Sigmoid, Sigmoid(x)); -// Returns 0 if x is 0, -1 if x < 0 and 1 if x > 0. -XLAJIT_MAKE_UNARY(Sign, xla::Sign(x)); +// Returns 0 if x is NaN, 0 if x is 0, -1 if x < 0 and 1 if x > 0. +XLAJIT_MAKE_UNARY(Sign, + xla::Select(xla::Ne(x, x), xla::ZerosLike(x), xla::Sign(x))); XLAJIT_MAKE_UNARY(Sinh, xla::Sinh(x)); // softplus(x) = log(1 + exp(x)) @@ -116,37 +114,11 @@ XLAJIT_MAKE_UNARY(Tanh, xla::Tanh(x)); XLAJIT_MAKE_UNARY(Real, xla::Real(x)); XLAJIT_MAKE_UNARY(Imag, xla::Imag(x)); +XLAJIT_MAKE_UNARY(Erf, xla::Erf(x)); +XLAJIT_MAKE_UNARY(Erfc, xla::Erfc(x)); #undef XLAJIT_MAKE_UNARY -// Erf/Erfc. For x in (-1, 1), the erf approximation is used; erfc polynomial -// is used outside of this range. -class ErfOp : public XlaOpKernel { - public: - explicit ErfOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} - void Compile(XlaOpKernelContext* ctx) override { - xla::XlaOp x = ctx->Input(0); - xla::XlaOp one = xla::ScalarLike(x, 1.0); - auto y = - xla::Select(xla::Gt(xla::Abs(x), one), one - xla::Erfc(x), xla::Erf(x)); - ctx->SetOutput(0, y); - } -}; -REGISTER_XLA_OP(Name("Erf"), ErfOp); - -class ErfcOp : public XlaOpKernel { - public: - explicit ErfcOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} - void Compile(XlaOpKernelContext* ctx) override { - xla::XlaOp x = ctx->Input(0); - xla::XlaOp one = xla::ScalarLike(x, 1.0); - auto y = - xla::Select(xla::Lt(xla::Abs(x), one), one - xla::Erf(x), xla::Erfc(x)); - ctx->SetOutput(0, y); - } -}; -REGISTER_XLA_OP(Name("Erfc"), ErfcOp); - class LgammaOp : public XlaOpKernel { public: explicit LgammaOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} diff --git a/tensorflow/compiler/tf2xla/kernels/unpack_op.cc b/tensorflow/compiler/tf2xla/kernels/unpack_op.cc index 8671632976023fded04c26a9780c1a67638b0916..2fc5619de737b8977e4249e4d2297a0303c339ce 100644 --- a/tensorflow/compiler/tf2xla/kernels/unpack_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/unpack_op.cc @@ -24,12 +24,12 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/bounds_check.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 941b04363f8386a7bdbe8c91ea34c9754592a52d..f49da9683b3622bdda708cc305306baafa1639df 100644 --- a/tensorflow/compiler/tf2xla/kernels/while_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/while_op.cc @@ -70,13 +70,20 @@ Status MakeXlaCompilerArgumentsFromInputs( arg.name = resource->name(); VLOG(2) << " resource " << resource->name() << " type: " << DataTypeString(arg.type) - << " shape: " << arg.shape.DebugString() + << " shape: " << arg.ShapeHumanString() << " initialized: " << arg.initialized; } else { arg.kind = XlaCompiler::Argument::kParameter; arg.type = ctx->input_type(i); - arg.shape = ctx->InputShape(i); + + xla::XlaBuilder* builder = ctx->builder(); + xla::XlaOp handle = ctx->Input(i); + auto shape_or_status = builder->GetShape(handle); + if (!shape_or_status.ok()) { + return shape_or_status.status(); + } + arg.shape = shape_or_status.ValueOrDie(); } } return Status::OK(); @@ -341,8 +348,11 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { VLOG(1) << "Done building while loop"; } -REGISTER_XLA_OP(Name("While").AllowResourceTypes(), XlaWhileOp); -REGISTER_XLA_OP(Name("StatelessWhile").AllowResourceTypes(), XlaWhileOp); -REGISTER_XLA_OP(Name("XlaWhile").AllowResourceTypes(), XlaWhileOp); +REGISTER_XLA_OP(Name("While").AllowResourceTypes().AllowVariantTypes(), + XlaWhileOp); +REGISTER_XLA_OP(Name("StatelessWhile").AllowResourceTypes().AllowVariantTypes(), + XlaWhileOp); +REGISTER_XLA_OP(Name("XlaWhile").AllowResourceTypes().AllowVariantTypes(), + XlaWhileOp); } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/util.cc b/tensorflow/compiler/tf2xla/lib/util.cc index c0bd172d17c192435ba8ee196f9def0491c0bf5c..06eda41611861060a1f1c4d028b96405d288efdb 100644 --- a/tensorflow/compiler/tf2xla/lib/util.cc +++ b/tensorflow/compiler/tf2xla/lib/util.cc @@ -54,6 +54,9 @@ xla::XlaOp FloatLiteral(xla::XlaBuilder* builder, xla::PrimitiveType type, case xla::C64: return xla::ConstantR0(builder, value); break; + case xla::C128: + return xla::ConstantR0(builder, value); + break; default: LOG(FATAL) << "unhandled element type " << type; } @@ -90,6 +93,9 @@ xla::XlaOp IntegerLiteral(xla::XlaBuilder* builder, xla::PrimitiveType type, case xla::C64: literal = xla::LiteralUtil::CreateR0(value); break; + case xla::C128: + literal = xla::LiteralUtil::CreateR0(value); + break; case xla::PRED: LOG(FATAL) << "pred element type is not integral"; case xla::S16: diff --git a/tensorflow/compiler/tf2xla/literal_util_test.cc b/tensorflow/compiler/tf2xla/literal_util_test.cc index 15f4c38da29507da9e092c1d5725b5f95a81d1b9..44bccfe6474d175beda392ca17dfbcb08c0b1b11 100644 --- a/tensorflow/compiler/tf2xla/literal_util_test.cc +++ b/tensorflow/compiler/tf2xla/literal_util_test.cc @@ -49,7 +49,7 @@ using Types = std::pair, std::pair, std::pair>; -TYPED_TEST_CASE(LiteralUtilTest, Types); +TYPED_TEST_SUITE(LiteralUtilTest, Types); TYPED_TEST(LiteralUtilTest, LiteralToQuantizedHostTensor) { using int_type = typename TypeParam::first_type; diff --git a/tensorflow/compiler/tf2xla/ops/BUILD b/tensorflow/compiler/tf2xla/ops/BUILD index 4dce0a2102cf9c782850ccc7af4f14b59bd51e53..7140b6a1227a53290c3747892a55886a7f48513b 100644 --- a/tensorflow/compiler/tf2xla/ops/BUILD +++ b/tensorflow/compiler/tf2xla/ops/BUILD @@ -4,7 +4,11 @@ package( licenses(["notice"]) # Apache 2.0 -load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py") +load( + "//tensorflow:tensorflow.bzl", + "tf_custom_op_library", + "tf_gen_op_wrapper_py", +) cc_library( name = "xla_ops", @@ -24,3 +28,14 @@ tf_gen_op_wrapper_py( ":xla_ops", ], ) + +tf_custom_op_library( + name = "_xla_ops.so", + srcs = [ + "xla_ops.cc", + ], + deps = [ + "//tensorflow/core:lib", + "@com_google_absl//absl/algorithm:container", + ], +) diff --git a/tensorflow/compiler/tf2xla/ops/xla_ops.cc b/tensorflow/compiler/tf2xla/ops/xla_ops.cc index ab77984684db4525f4d3f42b2c9c0f093c82ec45..af641131ed76a8d6a7291c360302fa17c94af014 100644 --- a/tensorflow/compiler/tf2xla/ops/xla_ops.cc +++ b/tensorflow/compiler/tf2xla/ops/xla_ops.cc @@ -369,7 +369,11 @@ REGISTER_OP("XlaKeyValueSort") .Output("sorted_values: V") .Attr("K: realnumbertype") .Attr("V: type") - .SetShapeFn(shape_inference::UnchangedShape) + .SetShapeFn([](shape_inference::InferenceContext* c) { + c->set_output(0, c->input(0)); + c->set_output(1, c->input(1)); + return Status::OK(); + }) .Doc(R"doc( Wraps the XLA Sort operator, documented at https://www.tensorflow.org/performance/xla/operation_semantics#sort diff --git a/tensorflow/compiler/tf2xla/python/BUILD b/tensorflow/compiler/tf2xla/python/BUILD index fef97b98c376d9df8bbfd9cb6651216895e46bf4..9abdb04d7736e8ff5225688af4759a522d3e7fc7 100644 --- a/tensorflow/compiler/tf2xla/python/BUILD +++ b/tensorflow/compiler/tf2xla/python/BUILD @@ -15,6 +15,7 @@ load( "//tensorflow/core:platform/default/build_config.bzl", "tf_py_clif_cc", ) +load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") tf_py_clif_cc( name = "xla_op_registry", @@ -27,9 +28,13 @@ tf_py_clif_cc( ], ) -py_library( +tf_custom_op_py_library( name = "xla", srcs = ["xla.py"], + dso = ["//tensorflow/compiler/tf2xla/ops:_xla_ops.so"], + kernels = [ + "//tensorflow/compiler/tf2xla/ops:xla_ops", + ], deps = [ "//tensorflow/compiler/tf2xla/ops:gen_xla_ops", "//tensorflow/compiler/xla:xla_data_proto_py", diff --git a/tensorflow/compiler/tf2xla/shape_util.cc b/tensorflow/compiler/tf2xla/shape_util.cc index cb04a75d95e4d4e7cfa921104ec052b95617def1..8997b2f5c68da480e9d4cb1f7ff8776690363392 100644 --- a/tensorflow/compiler/tf2xla/shape_util.cc +++ b/tensorflow/compiler/tf2xla/shape_util.cc @@ -44,6 +44,43 @@ Status PopulateInfeedLayoutVector(const xla::Shape& shape, return Status::OK(); } +// Populate the output layout unless the minor_to_major array contains all -1 +// value, in which case the layout is considered missing and the API returns +// false. +xla::StatusOr MakeLayout(absl::Span minor_to_major, + xla::Layout* layout) { + if (std::all_of(minor_to_major.begin(), minor_to_major.end(), + [](int64 dim) { return dim == -1; })) { + return false; + } + std::vector dim_present(minor_to_major.size(), false); + for (auto dim : minor_to_major) { + if (dim < 0 || dim >= minor_to_major.size()) { + return errors::InvalidArgument("Layout dimension out of range: dim=", dim, + " rank=", minor_to_major.size()); + } + if (dim_present[dim]) { + return errors::InvalidArgument("Repeated layout dimension: dim=", dim); + } + dim_present[dim] = true; + } + *layout = xla::LayoutUtil::MakeLayout(minor_to_major); + return true; +} + +Status AssignLayout( + absl::Span minor_to_major, + const std::function& layout_func, + xla::Shape* shape) { + xla::Layout layout; + TF_ASSIGN_OR_RETURN(bool has_layout, MakeLayout(minor_to_major, &layout)); + if (!has_layout && layout_func) { + layout = layout_func(*shape); + } + *shape->mutable_layout() = layout; + return Status::OK(); +} + } // namespace // Convert an XLA Shape into the equivalent TensorFlow shape. @@ -84,10 +121,64 @@ xla::Shape TensorShapeToXLAShape(xla::PrimitiveType type, return xla::ShapeUtil::MakeShapeWithLayout(type, dimensions, layout); } -xla::StatusOr> GetInfeedLayoutVector(const xla::Shape& shape) { +xla::StatusOr> GetShapeLayoutVector(const xla::Shape& shape) { std::vector layouts; TF_RETURN_IF_ERROR(PopulateInfeedLayoutVector(shape, &layouts)); return layouts; } +Status GetShapeWithLayout( + const xla::Shape& input_shape, absl::Span minor_to_major, + const std::function& layout_func, + xla::Shape* output_shape) { + if (input_shape.IsTuple()) { + int64 tuple_elements = xla::ShapeUtil::TupleElementCount(input_shape); + std::vector shapes; + shapes.reserve(tuple_elements); + size_t position = 0; + for (int64 i = 0; i < tuple_elements; ++i) { + const xla::Shape& shape = + xla::ShapeUtil::GetTupleElementShape(input_shape, i); + if (shape.IsTuple()) { + return errors::InvalidArgument( + "Nested tuples not supported: ", + xla::ShapeUtil::HumanString(input_shape)); + } + int64 rank = shape.rank(); + if (position + rank > minor_to_major.size()) { + return errors::InvalidArgument( + "Not enough layout attribute elements: position=", position, + " rank=", rank, " elements=", minor_to_major.size()); + } + shapes.push_back(shape); + TF_RETURN_IF_ERROR(AssignLayout( + absl::Span(minor_to_major).subspan(position, rank), + layout_func, &shapes.back())); + position += rank; + + VLOG(4) << "Shape[" << i + << "] = " << xla::ShapeUtil::HumanStringWithLayout(shapes.back()); + } + if (position != minor_to_major.size()) { + return errors::InvalidArgument( + "Too many elements passed in the layout attribute: position=", + position, " size=", minor_to_major.size()); + } + *output_shape = xla::ShapeUtil::MakeTupleShape(shapes); + } else { + int64 rank = input_shape.rank(); + if (rank != minor_to_major.size()) { + return errors::InvalidArgument( + "Wrong number of layout attribute elements: rank=", rank, + " elements=", minor_to_major.size()); + } + *output_shape = input_shape; + TF_RETURN_IF_ERROR(AssignLayout(minor_to_major, layout_func, output_shape)); + + VLOG(4) << "Shape[] = " + << xla::ShapeUtil::HumanStringWithLayout(*output_shape); + } + return Status::OK(); +} + } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/shape_util.h b/tensorflow/compiler/tf2xla/shape_util.h index cf52bf46e7c2a237d57f4c87e7d6efbf3fa9b1c2..e775c4462c3dc15cf4b8d9e8d8e7d9a61e024cd0 100644 --- a/tensorflow/compiler/tf2xla/shape_util.h +++ b/tensorflow/compiler/tf2xla/shape_util.h @@ -45,12 +45,23 @@ xla::Shape TensorShapeToXLAShape(xla::PrimitiveType type, const TensorShape& tensor_shape); // Given an XLA shape with layouts, builds a layout vector in the form able to -// be fed to an InfeedEnqueue/InfeedEnqueueTuple ops. +// be fed to ops like InfeedEnqueue/InfeedEnqueueTuple/XRTAllocateV2/.... // THe returned vector is a linearized sequence of the minor-to-major values of // the layouts held within the input shape. // In case the input shape is a tuple, the minor-to-major values will be in the // order of the tuple elements within the tuple shape. -xla::StatusOr> GetInfeedLayoutVector(const xla::Shape& shape); +// If a shape (or a subshape of a tuple shape) has missing layout, a rank long +// sequence of -1 values will be emittted. +xla::StatusOr> GetShapeLayoutVector(const xla::Shape& shape); + +// Given the input shape and a linearized sequence of the minor-to-major values +// of the layouts, create the output shape by rewriting the input shape layouts. +// If a layout is missing (has -1 values) for a matching tuple subshape, the +// layout_func will be called, if not nullptr. +Status GetShapeWithLayout( + const xla::Shape& input_shape, absl::Span minor_to_major, + const std::function& layout_func, + xla::Shape* output_shape); } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/side_effect_util.cc b/tensorflow/compiler/tf2xla/side_effect_util.cc index b62f8e9115229ac35c657d374c68336f1168ff77..412f31adbb7df52b2d6933be054cc6d40947dc44 100644 --- a/tensorflow/compiler/tf2xla/side_effect_util.cc +++ b/tensorflow/compiler/tf2xla/side_effect_util.cc @@ -26,6 +26,49 @@ const char kXlaTokenArgNodeName[] = "_xla_token_arg_node"; const char kXlaHasHostTransferAttrName[] = "_xla_has_host_transfer"; +Status SetDeviceOrdinalAttributeForNode(Node* node, int device_ordinal) { + if (!HasNodeAttr(node->def(), kXlaHasHostTransferAttrName)) { + return errors::InvalidArgument("Node ", node->DebugString(), + " does not have attribute ", + kXlaHasHostTransferAttrName); + } + + if (node->type_string() == "_XlaRecvAtHost" || + node->type_string() == "_XlaSendFromHost") { + node->ClearAttr("device_ordinal"); + node->AddAttr("device_ordinal", device_ordinal); + } else if (node->type_string() == "If") { + AttrValue device_ordinal_value; + device_ordinal_value.set_i(device_ordinal); + for (const string& attr_name : + std::vector{"then_branch", "else_branch"}) { + NameAttrList branch_func; + TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), attr_name, &branch_func)); + (*branch_func.mutable_attr())["device_ordinal"] = device_ordinal_value; + node->ClearAttr(attr_name); + node->AddAttr(attr_name, branch_func); + } + } else if (node->type_string() == "While") { + AttrValue device_ordinal_value; + device_ordinal_value.set_i(device_ordinal); + for (const string& attr_name : std::vector{"cond", "body"}) { + NameAttrList branch_func; + TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), attr_name, &branch_func)); + (*branch_func.mutable_attr())["device_ordinal"] = device_ordinal_value; + node->ClearAttr(attr_name); + node->AddAttr(attr_name, branch_func); + } + } else if (HasNodeAttr(node->def(), "device_ordinal")) { + // Function call node containing outside compilation. + node->ClearAttr("device_ordinal"); + node->AddAttr("device_ordinal", device_ordinal); + } else { + return errors::Internal("Unknown node type to set 'device_ordinal': ", + node->DebugString()); + } + return Status::OK(); +} + std::set CalculateTokenInputsForOutputToken(const Graph& g) { std::set results; Node* first_side_effecting_node_on_path = nullptr; diff --git a/tensorflow/compiler/tf2xla/side_effect_util.h b/tensorflow/compiler/tf2xla/side_effect_util.h index 7081b362c36c4785164b29003a5f89cd73bcf3af..75e1f253fb08ae61b0336a8783b7449c69197dd1 100644 --- a/tensorflow/compiler/tf2xla/side_effect_util.h +++ b/tensorflow/compiler/tf2xla/side_effect_util.h @@ -38,6 +38,10 @@ extern const char kXlaTokenArgNodeName[]; // This node have XlaRecvAtHost/XlaSendFromHost in its associated functions. extern const char kXlaHasHostTransferAttrName[]; +// Sets device ordinal attribute for nodes with attribute +// `kXlaHasHostTransferAttrName`. +Status SetDeviceOrdinalAttributeForNode(Node* node, int device_ordinal); + // Calculates side-effect dependencies for the graph's token output. // Returns a set of node names representing these dependencies. std::set CalculateTokenInputsForOutputToken(const Graph& g); diff --git a/tensorflow/compiler/tf2xla/tf2xla.cc b/tensorflow/compiler/tf2xla/tf2xla.cc index 9fac16a9700419b189bf5393c2b8bd7d76c6c1cc..28a4566c9d284fb8410a2d618f368c4dd2c1d893 100644 --- a/tensorflow/compiler/tf2xla/tf2xla.cc +++ b/tensorflow/compiler/tf2xla/tf2xla.cc @@ -243,7 +243,9 @@ Status CreateXlaArgs(const Graph& graph, XlaCompiler::Argument arg; arg.kind = XlaCompiler::Argument::kParameter; TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "T", &arg.type)); - TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), kShapeAttr, &arg.shape)); + TensorShape shape; + TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), kShapeAttr, &shape)); + arg.shape = shape; TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), kDebugNameAttr, &arg.name)); xla_args->push_back(arg); } @@ -252,7 +254,8 @@ Status CreateXlaArgs(const Graph& graph, // Converts the TensorFlow graph into an XLA computation, by executing the // graph symbolically, with each op building up the XLA HLO. -Status ConvertGraphToXla(std::unique_ptr graph, xla::Client* client, +Status ConvertGraphToXla(std::unique_ptr graph, + const tf2xla::Config& config, xla::Client* client, xla::XlaComputation* computation) { XlaOpRegistry::RegisterCompilationKernels(); for (Node* node : graph->nodes()) { @@ -262,6 +265,19 @@ Status ConvertGraphToXla(std::unique_ptr graph, xla::Client* client, std::vector xla_args; TF_RETURN_IF_ERROR(CreateXlaArgs(*graph, &xla_args)); + // Populate arguments with resource variables from the config. The variables + // get turned into inputs and outputs. + for (const tf2xla::Variable& variable : config.variable()) { + XlaCompiler::Argument arg; + arg.type = variable.type(); + arg.kind = XlaCompiler::Argument::kResource; + arg.shape = variable.shape(); + arg.name = variable.node_name(); + arg.resource_kind = XlaResource::kVariable; + arg.initialized = true; + xla_args.push_back(std::move(arg)); + } + // Compile the graph into an XLA computation. XlaCompiler::Options compiler_options; compiler_options.client = client; @@ -359,7 +375,8 @@ Status ConvertGraphDefToXla(const GraphDef& graph_def, xla::XlaComputation* computation) { std::unique_ptr graph; TF_RETURN_IF_ERROR(InitGraph(graph_def, config, &graph)); - TF_RETURN_IF_ERROR(ConvertGraphToXla(std::move(graph), client, computation)); + TF_RETURN_IF_ERROR( + ConvertGraphToXla(std::move(graph), config, client, computation)); return Status::OK(); } diff --git a/tensorflow/compiler/tf2xla/tf2xla.proto b/tensorflow/compiler/tf2xla/tf2xla.proto index 18c9089f5fa0e9792a4763d9bfac4c4e826eb5b2..5627af7452b99da594c1c214d0b556d8d70544d5 100644 --- a/tensorflow/compiler/tf2xla/tf2xla.proto +++ b/tensorflow/compiler/tf2xla/tf2xla.proto @@ -39,6 +39,15 @@ message Fetch { string name = 2; // Optional name for generated code. }; +// Variable represents a resource variable with the given name, shape and type. +message Variable { + string node_name = 1; + string name = + 2; // Optional name for generated code. If empty, node_name will be used. + TensorShapeProto shape = 3; + DataType type = 4; +} + // Config represents configuration information for tf2xla conversion. message Config { // Each feed is a positional input argument for the generated computation. @@ -47,4 +56,6 @@ message Config { // Each fetch is a positional output argument for the generated computation. // The order of each entry matches the order of each output argument. repeated Fetch fetch = 2; + // Each variable is a named input and output of the generated computation. + repeated Variable variable = 3; }; diff --git a/tensorflow/compiler/tf2xla/tf2xla_util.cc b/tensorflow/compiler/tf2xla/tf2xla_util.cc index 18d87727c500619bf386be7d8c7085724f44aba3..88c03a6056ac6484013c3fd32c9889899b5c15c5 100644 --- a/tensorflow/compiler/tf2xla/tf2xla_util.cc +++ b/tensorflow/compiler/tf2xla/tf2xla_util.cc @@ -122,7 +122,12 @@ Status ReplaceArgUsageWithConstNode( for (const auto& iter : const_input_index_to_node) { int arg_index = iter.first; - Node* const_node = g->CopyNode(iter.second); + NodeDef const_def = iter.second->def(); + const_def.set_name(g->NewName(const_def.name())); + Status s; + Node* const_node = g->AddNode(const_def, &s); + TF_RETURN_IF_ERROR(s); + Node* arg_node = arg_nodes[arg_index]; // Collect all usages of the _Arg node. @@ -265,6 +270,13 @@ Status PropagateConstIntoWhileNode(Graph* g, Node* while_node, } // Check if i-th retval's input comes from i-th arg directly. + // For resource variable input of While nodes, TF2XLA convention is to place + // them at the end of all inputs (after all data inputs), and *not* return + // them. So number of While node inputs might be larger than number of its + // outputs. + if (i >= body_func->signature().output_arg_size()) { + continue; + } const OpDef_ArgDef& output_arg = body_func->signature().output_arg(i); auto output_arg_input = body_func->ret().find(output_arg.name()); if (output_arg_input == body_func->ret().end()) { diff --git a/tensorflow/compiler/tf2xla/tf2xla_util_test.cc b/tensorflow/compiler/tf2xla/tf2xla_util_test.cc index 202e929315cacd4d6cdfc69d50639d8a427ec6c2..28b4744470e7d28863b5f7275f829b9bd59641e1 100644 --- a/tensorflow/compiler/tf2xla/tf2xla_util_test.cc +++ b/tensorflow/compiler/tf2xla/tf2xla_util_test.cc @@ -21,11 +21,13 @@ limitations under the License. #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/ops/data_flow_ops.h" #include "tensorflow/cc/ops/function_ops.h" +#include "tensorflow/cc/ops/functional_ops.h" #include "tensorflow/cc/ops/standard_ops.h" #include "tensorflow/compiler/tf2xla/sharding_util.h" #include "tensorflow/core/common_runtime/graph_optimizer.h" #include "tensorflow/core/common_runtime/process_function_library_runtime.h" #include "tensorflow/core/framework/function.h" +#include "tensorflow/core/framework/graph_to_functiondef.h" #include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/lib/core/status.h" @@ -329,5 +331,90 @@ TEST(CachedFunctionHandles, Basic) { TF_EXPECT_OK(cached_function_handles.ReleaseAllHandles()); } +TEST(PropagateConstIntoFunctionalNodes, WhileLoopWithResourceInput) { + FunctionLibraryDefinition fld(OpRegistry::Global(), {}); + { + // Cond graph & body graph. + Scope scope = Scope::NewRootScope().ExitOnError(); + auto pred = ops::_Arg(scope.WithOpName("pred"), DT_BOOL, 0); + auto input = ops::_Arg(scope.WithOpName("input"), DT_RESOURCE, 1); + auto ret = ops::_Retval(scope.WithOpName("ret"), pred, 0); + Graph graph(OpRegistry::Global()); + TF_ASSERT_OK(scope.ToGraph(&graph)); + FunctionDef cond_fdef; + TF_ASSERT_OK(GraphToFunctionDef(graph, "cond", &cond_fdef)); + TF_ASSERT_OK(fld.AddFunctionDef(cond_fdef)); + FunctionDef body_fdef; + TF_ASSERT_OK(GraphToFunctionDef(graph, "body", &body_fdef)); + TF_ASSERT_OK(fld.AddFunctionDef(body_fdef)); + } + Scope scope = Scope::NewRootScope().ExitOnError(); + auto pred = ops::Const(scope.WithOpName("pred"), false, TensorShape({})); + auto input = ops::Const(scope.WithOpName("input"), 0, TensorShape({})); + NameAttrList cond_fn, body_fn; + cond_fn.set_name("cond"); + body_fn.set_name("body"); + auto while_op = + ops::While(scope.WithOpName("while"), + std::initializer_list{pred, input}, cond_fn, body_fn); + Graph graph(OpRegistry::Global()); + TF_ASSERT_OK(scope.ToGraph(&graph)); + + TF_EXPECT_OK(PropagateConstIntoFunctionalNodes(&graph, &fld, &fld)); +} + +TEST(PropagateConstIntoFunctionalNodes, CopiedConstNodeHasUniqueName) { + FunctionLibraryDefinition fld(OpRegistry::Global(), {}); + { + // Cond graph & body graph. + Scope scope = Scope::NewRootScope().ExitOnError(); + auto pred = ops::_Arg(scope.WithOpName("arg0"), DT_BOOL, 0); + auto input = ops::_Arg(scope.WithOpName("arg1"), DT_BOOL, 1); + auto duplicate_name = ops::NoOp(scope.WithOpName("duplicate_name")); + auto ret = ops::_Retval(scope.WithOpName("ret"), pred, 0); + Graph graph(OpRegistry::Global()); + TF_ASSERT_OK(scope.ToGraph(&graph)); + FunctionDef cond_fdef; + TF_ASSERT_OK(GraphToFunctionDef(graph, "cond", &cond_fdef)); + TF_ASSERT_OK(fld.AddFunctionDef(cond_fdef)); + FunctionDef body_fdef; + TF_ASSERT_OK(GraphToFunctionDef(graph, "body", &body_fdef)); + TF_ASSERT_OK(fld.AddFunctionDef(body_fdef)); + } + Scope scope = Scope::NewRootScope().ExitOnError(); + auto pred = + ops::Const(scope.WithOpName("duplicate_name"), false, TensorShape({})); + auto input = ops::Const(scope.WithOpName("input"), false, TensorShape({})); + NameAttrList cond_fn, body_fn; + cond_fn.set_name("cond"); + body_fn.set_name("body"); + auto while_op = + ops::While(scope.WithOpName("while"), + std::initializer_list{pred, input}, cond_fn, body_fn); + Graph graph(OpRegistry::Global()); + TF_ASSERT_OK(scope.ToGraph(&graph)); + + TF_EXPECT_OK(PropagateConstIntoFunctionalNodes(&graph, &fld, &fld)); + + // Check that in rewritten body function, the NoOp node still has name + // "duplicate_name", and the copied Const node has name "duplicate_name/_0". + auto node_name_index = graph.BuildNodeNameIndex(); + Node* while_node = node_name_index["while"]; + ASSERT_NE(while_node, nullptr); + TF_ASSERT_OK(GetNodeAttr(while_node->def(), "body", &body_fn)); + const FunctionDef* rewritten_body_fn = fld.Find(body_fn.name()); + ASSERT_NE(rewritten_body_fn, nullptr); + std::unordered_map nodes; + for (const NodeDef& node_def : rewritten_body_fn->node_def()) { + nodes[node_def.name()] = node_def; + } + auto noop_def = nodes.find("duplicate_name"); + ASSERT_NE(noop_def, nodes.end()); + EXPECT_EQ(noop_def->second.op(), "NoOp"); + auto const_def = nodes.find("duplicate_name/_0"); + ASSERT_NE(const_def, nodes.end()); + EXPECT_EQ(const_def->second.op(), "Const"); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/type_util.cc b/tensorflow/compiler/tf2xla/type_util.cc index d00b1376620c0c9d112c7d7426758f6d3f25e86f..732f957d7329c93ad104dacf5190948fbfd7974b 100644 --- a/tensorflow/compiler/tf2xla/type_util.cc +++ b/tensorflow/compiler/tf2xla/type_util.cc @@ -69,6 +69,9 @@ Status DataTypeToPrimitiveType(DataType data_type, xla::PrimitiveType* type) { case tensorflow::DT_COMPLEX64: *type = xla::C64; return Status::OK(); + case tensorflow::DT_COMPLEX128: + *type = xla::C128; + return Status::OK(); default: return errors::InvalidArgument( "Unsupported type in DataTypeToPrimitiveType ", diff --git a/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h b/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h index c7341cf8b9e8d7a06fd304ae8766420d20f0c16e..de2e485a47c18ae8e58a06aba408dbb61a30d00a 100644 --- a/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h +++ b/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h @@ -59,45 +59,8 @@ class XlaCompiledCpuFunction { // AOT this is backed by data compiled into the object file. // // The contents of StaticData are XLA-internal implementation details and - // should not be relied on by clients. - // - // TODO(sanjoy): Come up with a cleaner way to express the contraint we want - // here: generated XlaCompiledCpuFunction subclasses should be able to create - // instances of StaticData but only XlaCompiledCpuFunction should be able to - // read from StaticData instances. + // should not be relied on by clients (and therefore are private). class StaticData { - public: - void set_raw_function(RawFunction raw_function) { - raw_function_ = raw_function; - } - void set_buffer_infos( - const cpu_function_runtime::BufferInfo* buffer_infos) { - buffer_infos_ = buffer_infos; - } - void set_num_buffers(size_t num_buffers) { num_buffers_ = num_buffers; } - void set_arg_index_table(const int32* arg_index_table) { - arg_index_table_ = arg_index_table; - } - void set_num_args(int64 num_args) { num_args_ = num_args; } - void set_result_index(size_t result_index) { result_index_ = result_index; } - void set_arg_names(const char** arg_names) { arg_names_ = arg_names; } - void set_result_names(const char** result_names) { - result_names_ = result_names; - } - void set_program_shape(const xla::ProgramShapeProto* program_shape) { - program_shape_ = program_shape; - } - const xla::HloProfilePrinterData* hlo_profile_printer_data() const { - return hlo_profile_printer_data_; - } - void set_hlo_profile_printer_data( - const xla::HloProfilePrinterData* hlo_profile_printer_data) { - hlo_profile_printer_data_ = hlo_profile_printer_data; - } - void set_profile_counters_size(int64 profile_counters_size) { - profile_counters_size_ = profile_counters_size; - } - private: // The raw function to call. RawFunction raw_function_; @@ -134,7 +97,8 @@ class XlaCompiledCpuFunction { // declared so we don't have access to that information here. int64 profile_counters_size_ = 0; - // Only XlaCompiledCpuFunction is allowed to read the above fields. + // Only XlaCompiledCpuFunction is allowed to read and write the above + // fields. friend class XlaCompiledCpuFunction; }; @@ -148,7 +112,7 @@ class XlaCompiledCpuFunction { RESULTS_PROFILES_AND_TEMPS_ONLY, }; - XlaCompiledCpuFunction( + explicit XlaCompiledCpuFunction( const StaticData& static_data, AllocMode alloc_mode = AllocMode::ARGS_RESULTS_PROFILES_AND_TEMPS); virtual ~XlaCompiledCpuFunction(); @@ -280,6 +244,76 @@ class XlaCompiledCpuFunction { return *hlo_profile_printer_data_; } + protected: + // --------------------------------------------------------------------------- + // Accessors for reading from and writing to instances of `StaticData`. + // + // Classes generated by tfcompile can call these because the generated classes + // inherit from `XlaCompiledCpuFunction`. `XlaJitCompiledCpuFunction` can + // call these because it is explicitly added as a friend. + + static void set_static_data_raw_function(StaticData* static_data, + RawFunction raw_function) { + static_data->raw_function_ = raw_function; + } + + static void set_static_data_buffer_infos( + StaticData* static_data, + const cpu_function_runtime::BufferInfo* buffer_infos) { + static_data->buffer_infos_ = buffer_infos; + } + + static void set_static_data_num_buffers(StaticData* static_data, + size_t num_buffers) { + static_data->num_buffers_ = num_buffers; + } + + static void set_static_data_arg_index_table(StaticData* static_data, + const int32* arg_index_table) { + static_data->arg_index_table_ = arg_index_table; + } + + static void set_static_data_num_args(StaticData* static_data, + int64 num_args) { + static_data->num_args_ = num_args; + } + + static void set_static_data_result_index(StaticData* static_data, + size_t result_index) { + static_data->result_index_ = result_index; + } + + static void set_static_data_arg_names(StaticData* static_data, + const char** arg_names) { + static_data->arg_names_ = arg_names; + } + + static void set_static_data_result_names(StaticData* static_data, + const char** result_names) { + static_data->result_names_ = result_names; + } + + static void set_static_data_program_shape( + StaticData* static_data, const xla::ProgramShapeProto* program_shape) { + static_data->program_shape_ = program_shape; + } + + static void set_static_data_hlo_profile_printer_data( + StaticData* static_data, + const xla::HloProfilePrinterData* hlo_profile_printer_data) { + static_data->hlo_profile_printer_data_ = hlo_profile_printer_data; + } + + static const xla::HloProfilePrinterData* + get_static_data_hlo_profile_printer_data(StaticData* static_data) { + return static_data->hlo_profile_printer_data_; + } + + static void set_static_data_profile_counters_size( + StaticData* static_data, int64 profile_counters_size) { + static_data->profile_counters_size_ = profile_counters_size; + } + private: const RawFunction raw_function_; const size_t result_index_; @@ -313,6 +347,10 @@ class XlaCompiledCpuFunction { const char** result_names_ = nullptr; const xla::ProgramShapeProto* program_shape_ = nullptr; const xla::HloProfilePrinterData* hlo_profile_printer_data_ = nullptr; + + // Add `XlaJitCompiledCpuFunction` as a friend so that it can access the + // `set_static_data_*` static methods above. + friend class XlaJitCompiledCpuFunction; }; } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc index ee461a3c07d4db514c7697e005a9371be4b54dd0..0833264523770dc43c6a784f8b3d731485f38e53 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler.cc @@ -31,6 +31,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/client/xla_computation.h" +#include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/common_runtime/device.h" #include "tensorflow/core/common_runtime/executor.h" #include "tensorflow/core/common_runtime/function.h" @@ -42,6 +43,8 @@ limitations under the License. #include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/graph/graph_constructor.h" #include "tensorflow/core/graph/node_builder.h" +#include "tensorflow/core/lib/core/error_codes.pb.h" +#include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/gtl/cleanup.h" #include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/platform/logging.h" @@ -57,7 +60,11 @@ Status CheckSignature(const DataTypeVector& types, " elements while function has ", types.size()); } for (int i = 0; i < types.size(); ++i) { - if (types[i] != args[i].type && types[i] != DT_RESOURCE) { + // Don't perform type checks on resource variables and tensor + // lists (DT_VARIANT) as we have to trick the type system in order to + // plumb them through. DT_VARIANTS are wrapped in a DT_UINT8 tensor. + if (types[i] != args[i].type && types[i] != DT_RESOURCE && + types[i] != DT_VARIANT) { return errors::Internal( "Argument ", i, " has declared type ", DataTypeString(args[i].type), " but function parameter has type ", DataTypeString(types[i])); @@ -192,6 +199,8 @@ Status BuildComputation( output.shape = output.constant_value.shape(); break; + case XlaExpression::Kind::kTensorList: + TF_FALLTHROUGH_INTENDED; case XlaExpression::Kind::kXlaOp: { output.is_constant = false; TF_ASSIGN_OR_RETURN(output.shape, retval.GetShape()); @@ -333,8 +342,21 @@ bool XlaCompiler::Argument::operator==( other.tensor_array_gradients)) { return false; } - if (shape != other.shape) { - return false; + if (absl::holds_alternative(shape)) { + if (!absl::holds_alternative(other.shape)) { + return false; + } + if (!xla::Shape::Equal()(absl::get(shape), + absl::get(other.shape))) { + return false; + } + } else { + if (!absl::holds_alternative(other.shape)) { + return false; + } + if (absl::get(shape) != absl::get(other.shape)) { + return false; + } } if (constant_value.shape() != other.constant_value.shape()) { return false; @@ -348,7 +370,7 @@ string XlaCompiler::Argument::HumanString() const { common = absl::StrCat(" name=", name); } absl::StrAppend(&common, " type=", DataTypeString(type), - " shape=", shape.DebugString()); + " shape=", ShapeHumanString()); switch (kind) { case kInvalid: return "invalid"; @@ -375,6 +397,23 @@ string XlaCompiler::Argument::HumanString() const { } } +std::vector XlaCompiler::Argument::DimensionSizes() const { + if (absl::holds_alternative(shape)) { + return xla::InlinedVectorToVector( + absl::get(shape).dim_sizes()); + } else { + return absl::get(shape).dimensions(); + } +} + +string XlaCompiler::Argument::ShapeHumanString() const { + if (absl::holds_alternative(shape)) { + return absl::get(shape).DebugString(); + } else { + return absl::get(shape).DebugString(); + } +} + XlaCompiler::XlaCompiler(XlaCompiler::Options options) : options_(options), initialization_status_(Status::OK()), @@ -462,8 +501,34 @@ std::unique_ptr XlaCompiler::GetGraph(const FunctionBody* fbody) { opts.set_do_function_inlining(true); opts.set_do_constant_folding(true); GraphOptimizer optimizer(opts); + // Do not constant fold nodes that output DT_VARIANT type tensors. + // XLA does not support Const nodes of Variant type since it needs + // to know the original ops to be able to compile them to the relevant + // XLA form. + // TODO(srbs): This filter is a little conservative. E.g. a subgraph of + // the form: + // Const + // | + // EmptyTensorList -> TensorListPushBack -> TensorListPopBack -> Op + // | + // (Discard popped list) + // + // Would have been reduced to "Const -> Op" without this filter. + // However since we are only allowed to specify the filter at the "Node" + // level there is no good way to allow the above behavior. So we + // disallow any sort of constant folding on Variant nodes for now. + auto cf_consider_fn = [](const Node* n) { + for (const auto& output_arg : n->op_def().output_arg()) { + if (output_arg.type() == DT_VARIANT) { + return false; + } + } + return true; + }; + GraphOptimizer::Options graph_optimizer_options; + graph_optimizer_options.cf_consider_fn = cf_consider_fn; optimizer.Optimize(flib_runtime_, flib_runtime_->env(), - /*device=*/nullptr, &graph, /*shape_map=*/nullptr); + /*device=*/nullptr, &graph, graph_optimizer_options); return graph; } @@ -548,11 +613,22 @@ Status XlaCompiler::XLAShapeForArgument(const XlaCompiler::Argument& arg, LOG(FATAL) << "Unreachable case"; case XlaCompiler::Argument::kParameter: { if (is_entry_computation) { - TF_ASSIGN_OR_RETURN( - *xla_shape, options_.shape_representation_fn(arg.shape, arg.type)); + TensorShape shape; + if (absl::holds_alternative(arg.shape)) { + shape = absl::get(arg.shape); + } else { + TF_RETURN_IF_ERROR( + XLAShapeToTensorShape(absl::get(arg.shape), &shape)); + } + TF_ASSIGN_OR_RETURN(*xla_shape, + options_.shape_representation_fn(shape, arg.type)); } else { - TF_RETURN_IF_ERROR( - TensorShapeToXLAShape(arg.type, arg.shape, xla_shape)); + if (absl::holds_alternative(arg.shape)) { + *xla_shape = absl::get(arg.shape); + } else { + TF_RETURN_IF_ERROR(TensorShapeToXLAShape( + arg.type, absl::get(arg.shape), xla_shape)); + } } return Status::OK(); } @@ -561,8 +637,10 @@ Status XlaCompiler::XLAShapeForArgument(const XlaCompiler::Argument& arg, switch (arg.resource_kind) { case XlaResource::kVariable: { - TF_ASSIGN_OR_RETURN(*xla_shape, options_.shape_representation_fn( - arg.shape, arg.type)); + TF_RET_CHECK(absl::holds_alternative(arg.shape)); + TF_ASSIGN_OR_RETURN(*xla_shape, + options_.shape_representation_fn( + absl::get(arg.shape), arg.type)); return Status::OK(); } @@ -571,9 +649,10 @@ Status XlaCompiler::XLAShapeForArgument(const XlaCompiler::Argument& arg, return errors::InvalidArgument( "Negative max_array_size in XLAShapeForArgument"); } + TF_RET_CHECK(absl::holds_alternative(arg.shape)); TensorShape shape; shape.AddDim(arg.max_array_size); - shape.AppendShape(arg.shape); + shape.AppendShape(absl::get(arg.shape)); TF_RETURN_IF_ERROR(TensorShapeToXLAShape(arg.type, shape, xla_shape)); if (!arg.tensor_array_gradients.empty()) { @@ -588,9 +667,10 @@ Status XlaCompiler::XLAShapeForArgument(const XlaCompiler::Argument& arg, return errors::InvalidArgument( "Negative max_array_size in XLAShapeForArgument"); } + TF_RET_CHECK(absl::holds_alternative(arg.shape)); TensorShape shape; shape.AddDim(arg.max_array_size); - shape.AppendShape(arg.shape); + shape.AppendShape(absl::get(arg.shape)); xla::Shape buffer_shape; TF_RETURN_IF_ERROR( TensorShapeToXLAShape(arg.type, shape, &buffer_shape)); @@ -620,14 +700,15 @@ Status XlaCompiler::BuildArguments( bool use_tuple_arg, xla::XlaBuilder* builder, XlaContext* context, const std::map& arg_cores, std::vector* arg_expressions, - std::vector* input_mapping, std::vector* input_shapes, + std::vector* input_to_args, std::vector* input_shapes, bool is_entry_computation) { arg_expressions->resize(args.size()); // Argument numbers of arguments and resources that are to be passed to the - // XLA computation as runtime parameters. - input_mapping->clear(); - input_mapping->reserve(args.size()); + // XLA computation as runtime parameters. `input_to_args[a] = b` means that + // the a'th XLA input corresponds to the b'th original arg indexes. + input_to_args->clear(); + input_to_args->reserve(args.size()); // Fills in constant arguments, and computes non-constant argument order. for (std::vector::size_type i = 0; i < args.size(); @@ -637,24 +718,25 @@ Status XlaCompiler::BuildArguments( switch (arg.kind) { case XlaCompiler::Argument::kResource: { TF_RET_CHECK(arg.resource_kind != XlaResource::kInvalid); + TF_RET_CHECK(absl::holds_alternative(arg.shape)); // TODO(phawkins): this code assumes that resource arguments do not // alias. XlaResource* resource = context->AddResource(absl::make_unique( - arg.resource_kind, i, arg.name, arg.type, arg.shape, - xla::XlaOp(), + arg.resource_kind, i, arg.name, arg.type, + absl::get(arg.shape), xla::XlaOp(), /*max_array_size=*/arg.max_array_size, /*tensor_array_gradients=*/arg.tensor_array_gradients, /*tensor_array_multiple_writes_aggregate=*/true)); arg_expression = XlaExpression::Resource(resource); if (arg.initialized) { - input_mapping->push_back(i); + input_to_args->push_back(i); } break; } case XlaCompiler::Argument::kParameter: case XlaCompiler::Argument::kToken: { - input_mapping->push_back(i); + input_to_args->push_back(i); break; } case XlaCompiler::Argument::kConstant: @@ -666,15 +748,23 @@ Status XlaCompiler::BuildArguments( } } - if (input_mapping->empty()) { + if (input_to_args->empty()) { return Status::OK(); } - std::vector arg_shapes(input_mapping->size()); - for (std::vector::size_type i = 0; i < input_mapping->size(); ++i) { + // `arg_to_inputs[c] = d` means that the c'th original arg index corresponds + // to the d'th XLA input. Note that the value -1 corresponds to constants, or + // other args that don't correspond to an input. + std::vector arg_to_inputs(args.size(), -1); + for (int i = 0; i < input_to_args->size(); i++) { + arg_to_inputs[input_to_args->at(i)] = i; + } + + std::vector arg_shapes(input_to_args->size()); + for (std::vector::size_type i = 0; i < input_to_args->size(); ++i) { // Computes the shapes of non-constant arguments. TF_RETURN_IF_ERROR(XLAShapeForArgument( - args[(*input_mapping)[i]], is_entry_computation, &arg_shapes[i])); + args[(*input_to_args)[i]], is_entry_computation, &arg_shapes[i])); } if (use_tuple_arg) { @@ -691,13 +781,13 @@ Status XlaCompiler::BuildArguments( builder->SetOpMetadata(arg_metadata); // Build parameter handles for non-constant arguments. - std::vector arg_handles(input_mapping->size()); + std::vector arg_handles(input_to_args->size()); if (use_tuple_arg) { xla::XlaOp tuple; if (is_entry_computation) { xla::OpSharding tuple_sharding; tuple_sharding.set_type(xla::OpSharding::Type::OpSharding_Type_TUPLE); - for (int64 parameter : *input_mapping) { + for (int64 parameter : *input_to_args) { auto it = arg_cores.find(parameter); const int core = it == arg_cores.end() ? 0 : it->second; *tuple_sharding.add_tuple_shardings() = @@ -709,7 +799,19 @@ Status XlaCompiler::BuildArguments( } else { tuple = xla::Parameter(builder, 0, (*input_shapes)[0], "arg_tuple"); } - for (std::vector::size_type i = 0; i < input_mapping->size(); ++i) { + + for (int i = 0; i < input_to_args->size(); ++i) { + const XlaCompiler::Argument& arg = args[input_to_args->at(i)]; + for (const auto& dim_and_arg_num : arg.dynamic_dim_to_arg_num_map) { + int dynamic_size_param_index = arg_to_inputs.at(dim_and_arg_num.second); + TF_RETURN_IF_ERROR(builder->SetDynamicBinding( + /*dynamic_size_param_num=*/0, {dynamic_size_param_index}, + /*target_param_num=*/0, /*target_param_index=*/{i}, + dim_and_arg_num.first)); + } + } + + for (std::vector::size_type i = 0; i < input_to_args->size(); ++i) { auto it = arg_cores.find(i); const int core = it == arg_cores.end() ? -1 : it->second; xla::XlaScopedShardingAssignment assign_sharding( @@ -718,7 +820,7 @@ Status XlaCompiler::BuildArguments( arg_handles[i] = xla::GetTupleElement(tuple, i); } } else { - for (std::vector::size_type i = 0; i < input_mapping->size(); ++i) { + for (std::vector::size_type i = 0; i < input_to_args->size(); ++i) { auto it = arg_cores.find(i); const int core = it == arg_cores.end() ? -1 : it->second; xla::XlaScopedShardingAssignment assign_sharding( @@ -727,6 +829,17 @@ Status XlaCompiler::BuildArguments( arg_handles[i] = xla::Parameter(builder, i, (*input_shapes)[i], absl::StrCat("arg", i)); } + + for (int i = 0; i < input_to_args->size(); ++i) { + const XlaCompiler::Argument& arg = args[input_to_args->at(i)]; + for (const auto& dim_and_arg_num : arg.dynamic_dim_to_arg_num_map) { + int dynamic_size_param_index = arg_to_inputs.at(dim_and_arg_num.second); + TF_RETURN_IF_ERROR(builder->SetDynamicBinding( + /*dynamic_size_param_num=*/dynamic_size_param_index, {}, + /*target_param_num=*/i, /*target_param_index=*/{}, + dim_and_arg_num.first)); + } + } } builder->ClearOpMetadata(); @@ -734,12 +847,12 @@ Status XlaCompiler::BuildArguments( // Fill in the handles in non-constant arguments, and reshape parameters // back to their correct shapes. VLOG(2) << "XLA computation inputs:"; - for (std::vector::size_type i = 0; i < input_mapping->size(); ++i) { - const XlaCompiler::Argument& arg = args[input_mapping->at(i)]; + for (std::vector::size_type i = 0; i < input_to_args->size(); ++i) { + const XlaCompiler::Argument& arg = args[input_to_args->at(i)]; VLOG(2) << " XLA arg " << i << " shape: " << xla::ShapeUtil::HumanString(arg_shapes[i]) - << " name: " << arg.name << " TF arg " << input_mapping->at(i); - XlaExpression& arg_expression = (*arg_expressions)[input_mapping->at(i)]; + << " name: " << arg.name << " TF arg " << input_to_args->at(i); + XlaExpression& arg_expression = (*arg_expressions)[input_to_args->at(i)]; switch (arg.kind) { case XlaCompiler::Argument::kResource: { TF_RET_CHECK(arg.initialized); @@ -756,7 +869,7 @@ Status XlaCompiler::BuildArguments( // return values of functions, and then reshape unconditionally. if (is_entry_computation) { arg_expression = XlaExpression::XlaOp( - xla::Reshape(arg_handles[i], arg.shape.dim_sizes()), arg.type); + xla::Reshape(arg_handles[i], arg.DimensionSizes()), arg.type); } else { arg_expression = XlaExpression::XlaOp(arg_handles[i], arg.type); } @@ -997,8 +1110,17 @@ Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, result->outputs.resize(context->retvals().size()); std::vector retvals = context->retvals(); if (options.resolve_compile_time_constants) { - TF_RETURN_IF_ERROR(ResolveConstantExpressionsToConstants( - client(), absl::Span(retvals))); + Status status = ResolveConstantExpressionsToConstants( + client(), absl::Span(retvals)); + + // If the HloEvaluator has not implemented an expression, just evaluate it + // at runtime. + if (status.code() == error::UNIMPLEMENTED) { + ConvertConstantsToExpressions(&builder, + absl::Span(retvals)); + } else { + TF_RETURN_IF_ERROR(status); + } } else { ConvertConstantsToExpressions(&builder, absl::Span(retvals)); } diff --git a/tensorflow/compiler/tf2xla/xla_compiler.h b/tensorflow/compiler/tf2xla/xla_compiler.h index 0d801b73a8c2651305328384377751254ecaa41d..ad3144b41bdf3fc8b75ab5230e8e128df2962884 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.h +++ b/tensorflow/compiler/tf2xla/xla_compiler.h @@ -19,6 +19,7 @@ limitations under the License. #include #include "absl/types/span.h" +#include "absl/types/variant.h" #include "tensorflow/compiler/tf2xla/host_compute_metadata.pb.h" #include "tensorflow/compiler/tf2xla/xla_compilation_device.h" #include "tensorflow/compiler/tf2xla/xla_expression.h" @@ -124,7 +125,8 @@ class XlaCompiler { DataType type = DT_INVALID; // The shape of the argument. For: - // * a parameter: the shape of the parameter. + // * a parameter: the shape of the parameter. We allow setting the xla shape + // if known. This helps avoid conversions to and from TensorShape. // * a constant: ignored; the shape given by constant_value is used // instead. // * an uninitialized resource: ignored. We don't yet know the shape of an @@ -133,7 +135,7 @@ class XlaCompiler { // * an initialized TensorArray or Stack resource: the shape of an entry in // the TensorArray/Stack. Note this is the size of a single entry, not the // XLA data structure that represents the complete stack/array. - TensorShape shape; + absl::variant shape; // The value of the argument, if it is a compile-time constant. Must be a // host-memory tensor. @@ -157,10 +159,20 @@ class XlaCompiler { // as `tensor_array_gradients`. std::set tensor_array_gradients; + // dynamic dims to arg number map. Empty if no dynamic shapes. + std::map dynamic_dim_to_arg_num_map; + bool is_pad_arg = false; + bool operator==(const Argument& other) const; // Returns a human-readable summary of the argument. string HumanString() const; + + // Returns the dimension sizes for either TensorShape or xla::Shape. + std::vector DimensionSizes() const; + + // Returns the human-readable string for either TensorShape or xla::Shape. + string ShapeHumanString() const; }; // Options pertaining to an individual call to CompileGraph() or @@ -420,7 +432,7 @@ class XlaCompiler { XlaContext* context, const std::map& arg_cores, std::vector* arg_expressions, - std::vector* input_mapping, + std::vector* input_to_args, std::vector* input_shapes, bool is_entry_computation); diff --git a/tensorflow/compiler/tf2xla/xla_compiler_test.cc b/tensorflow/compiler/tf2xla/xla_compiler_test.cc index 32a430184099e8bac28182ea4c7790ea1b8266de..492010f7317d32a8a620147cd2cd9356d4f13fde 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler_test.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler_test.cc @@ -82,7 +82,7 @@ namespace { // compiled kernels. class DummyResourceForTest : public ResourceBase { public: - string DebugString() override { return "dummy"; } + string DebugString() const override { return "dummy"; } void Increment() { ++value_; } int Get() { return value_; } diff --git a/tensorflow/compiler/tf2xla/xla_context.cc b/tensorflow/compiler/tf2xla/xla_context.cc index a69af70503376b6c0905deb8980abdc3254a6e47..3f787fd86c9f7366a7728dcf146a3797ba672bc3 100644 --- a/tensorflow/compiler/tf2xla/xla_context.cc +++ b/tensorflow/compiler/tf2xla/xla_context.cc @@ -61,7 +61,7 @@ void XlaContext::set_args(std::vector args) { XlaContext::XlaContext(XlaCompiler* compiler, xla::XlaBuilder* builder) : compiler_(compiler), builder_(builder) {} -string XlaContext::DebugString() { return "XLA JIT context"; } +string XlaContext::DebugString() const { return "XLA JIT context"; } void XlaContext::SetRetval(int index, const XlaExpression& expression) { if (retvals_.size() <= index) { @@ -76,7 +76,7 @@ XlaResource* XlaContext::AddResource(std::unique_ptr resource) { } const xla::XlaComputation* XlaContext::GetOrCreateMax(const DataType type) { - return LookupOrCreate(type, &max_func_, [this, type] { + return LookupOrCreate(type, &max_func_, [type] { const string type_string = DataTypeString(type); VLOG(1) << "Building Max() for " << type_string; xla::XlaBuilder b("max<" + type_string + ">"); @@ -92,7 +92,7 @@ const xla::XlaComputation* XlaContext::GetOrCreateMax(const DataType type) { } const xla::XlaComputation* XlaContext::GetOrCreateMin(const DataType type) { - return LookupOrCreate(type, &min_func_, [this, type] { + return LookupOrCreate(type, &min_func_, [type] { const string type_string = DataTypeString(type); VLOG(1) << "Building Min() for " << type_string; xla::XlaBuilder b("min<" + type_string + ">"); @@ -108,7 +108,7 @@ const xla::XlaComputation* XlaContext::GetOrCreateMin(const DataType type) { } const xla::XlaComputation* XlaContext::GetOrCreateAdd(const DataType type) { - return LookupOrCreate(type, &add_func_, [this, type] { + return LookupOrCreate(type, &add_func_, [type] { const string type_string = DataTypeString(type); VLOG(1) << "Building Add() for " << type_string; xla::XlaBuilder b("add<" + type_string + ">"); @@ -124,7 +124,7 @@ const xla::XlaComputation* XlaContext::GetOrCreateAdd(const DataType type) { } const xla::XlaComputation* XlaContext::GetOrCreateMul(const DataType type) { - return LookupOrCreate(type, &mul_func_, [this, type] { + return LookupOrCreate(type, &mul_func_, [type] { const string type_string = DataTypeString(type); VLOG(1) << "Building Mul() for " << type_string; xla::XlaBuilder b("mul<" + type_string + ">"); diff --git a/tensorflow/compiler/tf2xla/xla_context.h b/tensorflow/compiler/tf2xla/xla_context.h index 0767d1faac14cedb8666f6cc37175eb7b55f6158..eb4ad3fe6a14b42a4df2c73c71cb6df1331fd796 100644 --- a/tensorflow/compiler/tf2xla/xla_context.h +++ b/tensorflow/compiler/tf2xla/xla_context.h @@ -47,7 +47,7 @@ class XlaContext : public ResourceBase { XlaContext(XlaCompiler* compiler, xla::XlaBuilder* builder); // Virtual method defined by ResourceBase. - string DebugString() override; + string DebugString() const override; XlaCompiler* compiler() const { return compiler_; } diff --git a/tensorflow/compiler/tf2xla/xla_expression.cc b/tensorflow/compiler/tf2xla/xla_expression.cc index ca0309166b7c73d1a5a818091e2a30fa112a4de4..3d228c92adcbe3d093a4fe70d157e57ab3e80c80 100644 --- a/tensorflow/compiler/tf2xla/xla_expression.cc +++ b/tensorflow/compiler/tf2xla/xla_expression.cc @@ -46,6 +46,14 @@ XlaExpression XlaExpression::XlaOp(xla::XlaOp value, DataType dtype) { return e; } +XlaExpression XlaExpression::TensorList(xla::XlaOp tensor_list) { + XlaExpression e; + e.kind_ = Kind::kTensorList; + e.dtype_ = DT_VARIANT; + e.handle_ = tensor_list; + return e; +} + XlaExpression XlaExpression::Resource(XlaResource* resource) { XlaExpression e; e.kind_ = Kind::kResource; @@ -64,6 +72,8 @@ string XlaExpression::HumanString() const { return "xla_op"; case Kind::kResource: return "resource"; + case Kind::kTensorList: + return "tensor_list"; } } @@ -76,6 +86,8 @@ xla::XlaOp XlaExpression::AsXlaOp(xla::XlaBuilder* builder) const { HostTensorToBorrowingLiteral(constant_value_, &literal)); return xla::ConstantLiteral(builder, literal); } + case Kind::kTensorList: + TF_FALLTHROUGH_INTENDED; case Kind::kXlaOp: if (builder != handle_.builder()) { return errors::InvalidArgument( @@ -96,7 +108,10 @@ xla::StatusOr> XlaExpression::ResolveConstant( return {constant_value()}; case Kind::kXlaOp: break; + case Kind::kTensorList: + TF_FALLTHROUGH_INTENDED; case Kind::kResource: + TF_FALLTHROUGH_INTENDED; case Kind::kInvalid: return errors::InvalidArgument( "ResolveConstant called on XlaExpression: ", HumanString()); @@ -134,6 +149,8 @@ xla::StatusOr XlaExpression::GetShape() const { TF_RETURN_IF_ERROR(XLAShapeToTensorShape(xla_shape, &shape)); return shape; } + case Kind::kTensorList: + return TensorShape({}); case Kind::kResource: return TensorShape({}); case Kind::kInvalid: diff --git a/tensorflow/compiler/tf2xla/xla_expression.h b/tensorflow/compiler/tf2xla/xla_expression.h index bed6761d362a98d344003c1edea342e68c31ef07..ac0232d8924cf2c9e35ad3f0772a3a2adc18af87 100644 --- a/tensorflow/compiler/tf2xla/xla_expression.h +++ b/tensorflow/compiler/tf2xla/xla_expression.h @@ -32,11 +32,16 @@ namespace tensorflow { // * a constant tensor. // * an xla::XlaOp, representing a symbolic XLA value. // * a resource, e.g., a variable, represented as an XlaResource pointer. +// * a tensor list, represented by a tuple of tensors and the list length. // // Constant tensors are mostly an optimization to avoid passing large constants // to XLA, but are also sometimes used to represent tensors that have no XLA // representation, for example, DT_STRING tensors. A canonical use case might be // an error message string. +// +// Tensor lists are very similar to xla::XlaOp, however they require some +// specific logic around shape management since the tuples are not supported by +// TensorFlow. class XlaExpression { public: enum class Kind { @@ -44,6 +49,7 @@ class XlaExpression { kConstant, kXlaOp, kResource, + kTensorList, }; XlaExpression(); @@ -62,6 +68,9 @@ class XlaExpression { // be derived from the XLA type. static XlaExpression XlaOp(xla::XlaOp value, DataType dtype); + // Builds a tensor list expression. + static XlaExpression TensorList(xla::XlaOp tensor_list); + // Builds a resource expression. static XlaExpression Resource(XlaResource* resource); @@ -100,7 +109,8 @@ class XlaExpression { DataType dtype_ = DT_INVALID; - // The XLA handle of the expression's computation, if kind_ == kXlaOp. + // The XLA handle of the expression's computation, if kind_ == kXlaOp or + // a tuple expression if kind_ == kTensorList. xla::XlaOp handle_; // The value of the constant, if kind_ == kConstant. diff --git a/tensorflow/compiler/tf2xla/xla_helpers.cc b/tensorflow/compiler/tf2xla/xla_helpers.cc index 00035d24b7891060adbc5ff871356b7471ab92ef..04a5d934064a9083a41cc210b48df65bbc862fff 100644 --- a/tensorflow/compiler/tf2xla/xla_helpers.cc +++ b/tensorflow/compiler/tf2xla/xla_helpers.cc @@ -34,63 +34,6 @@ limitations under the License. namespace tensorflow { -namespace { - -xla::XlaOp ArgMinMax(xla::XlaOp input, xla::PrimitiveType output_type, int axis, - bool is_min) { - xla::XlaBuilder* builder = input.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - TF_ASSIGN_OR_RETURN(xla::Shape input_shape, builder->GetShape(input)); - xla::XlaOp init_value; - xla::XlaComputation reducer; - if (is_min) { - init_value = xla::MaxValue(builder, input_shape.element_type()); - reducer = - xla::CreateScalarMinComputation(input_shape.element_type(), builder); - } else { - init_value = xla::MinValue(builder, input_shape.element_type()); - reducer = - xla::CreateScalarMaxComputation(input_shape.element_type(), builder); - } - - xla::XlaOp input_max = xla::Reduce(input, init_value, reducer, - /*dimensions_to_reduce=*/{axis}); - std::vector broadcast_dims(input_shape.rank() - 1); - std::iota(broadcast_dims.begin(), broadcast_dims.begin() + axis, 0); - std::iota(broadcast_dims.begin() + axis, broadcast_dims.end(), axis + 1); - // Compute a mask that has 1s for elements equal to the maximum. - xla::XlaOp partial_mask = xla::ConvertElementType( - xla::Eq(input, input_max, broadcast_dims), output_type); - - // In order to make identity elements for a bitwise And, we: - // Left shift the 1 to the leftmost bit, yielding 0x10...0 - // Arithmetic right shift the 1 back to the rightmost bit, yielding - // 0xFF...F - int32 bits_in_type = - xla::ShapeUtil::ByteSizeOfPrimitiveType(output_type) * 8 - 1; - xla::XlaOp shift_amount = - xla::ConstantR0WithType(builder, output_type, bits_in_type); - xla::XlaOp full_mask = xla::ShiftRightArithmetic( - xla::ShiftLeft(partial_mask, shift_amount), shift_amount); - - // And with the vector [0, 1, 2, ...] to convert each 0xFF...F into its - // index. - - const int64 axis_size = xla::ShapeUtil::GetDimension(input_shape, axis); - xla::XlaOp iota = xla::Iota(builder, output_type, axis_size); - xla::XlaOp product = - xla::And(full_mask, iota, /*broadcast_dimensions=*/{axis}); - - // If there are multiple maximum elements, choose the one with the highest - // index. - return xla::Reduce(product, xla::MinValue(builder, output_type), - xla::CreateScalarMaxComputation(output_type, builder), - /*dimensions_to_reduce=*/{axis}); - }); -} - -} // namespace - xla::XlaOp XlaHelpers::Zero(xla::XlaBuilder* b, DataType data_type) { xla::PrimitiveType type; TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type)); @@ -148,16 +91,6 @@ static Tensor MakeLinspaceTensor(const TensorShape& shape, int64 depth) { return linspace; } -xla::XlaOp XlaHelpers::ArgMax(xla::XlaOp input, xla::PrimitiveType output_type, - int axis) { - return ArgMinMax(input, output_type, axis, /*is_min=*/false); -} - -xla::XlaOp XlaHelpers::ArgMin(xla::XlaOp input, xla::PrimitiveType output_type, - int axis) { - return ArgMinMax(input, output_type, axis, /*is_min=*/true); -} - Status XlaHelpers::OneHot(xla::XlaBuilder* builder, int64 depth, int axis, DataType index_type, const TensorShape& indices_shape, const xla::XlaOp& indices, const xla::XlaOp& on_value, diff --git a/tensorflow/compiler/tf2xla/xla_helpers.h b/tensorflow/compiler/tf2xla/xla_helpers.h index 4858dfee55a393d04cd2af83916eeb40820ee368..490923526bd3acd4b167ccb3faff1d6c9e631131 100644 --- a/tensorflow/compiler/tf2xla/xla_helpers.h +++ b/tensorflow/compiler/tf2xla/xla_helpers.h @@ -53,16 +53,6 @@ class XlaHelpers { absl::Span shape, xla::Literal* output); - // Returns the argmax of `input` along `axis`. `output_type` is the type to - // use for the output. - static xla::XlaOp ArgMax(xla::XlaOp input, xla::PrimitiveType output_type, - int axis); - - // Returns the argmin of `input` along `axis`. `output_type` is the type to - // use for the output. - static xla::XlaOp ArgMin(xla::XlaOp input, xla::PrimitiveType output_type, - int axis); - // Converts `indices` into a one-hot representation. `depth` is the size // of the new axis to add. `axis` is the position at which to add the new // axis. `indices_shape` is the shape of `indices`. `on_value` and diff --git a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc index fabbcd04fed96ad814d04c2df9394f43bfe0cf99..884dc45cb11b18ae557c3da3f4192b3805cb7980 100644 --- a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc +++ b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc @@ -135,24 +135,34 @@ XlaJitCompiledCpuFunction::Compile( jit->arg_index_table_ = std::move(arg_index_table); jit->program_shape_ = absl::make_unique(program_shape->ToProto()); - jit->static_data_.set_raw_function(raw_function); - jit->static_data_.set_buffer_infos(jit->buffer_infos_.data()); - jit->static_data_.set_num_buffers(jit->buffer_infos_.size()); - jit->static_data_.set_arg_index_table(jit->arg_index_table_.data()); - jit->static_data_.set_num_args(jit->arg_index_table_.size()); - jit->static_data_.set_result_index(result_index); + XlaCompiledCpuFunction::set_static_data_raw_function(&jit->static_data_, + raw_function); + XlaCompiledCpuFunction::set_static_data_buffer_infos( + &jit->static_data_, jit->buffer_infos_.data()); + XlaCompiledCpuFunction::set_static_data_num_buffers( + &jit->static_data_, jit->buffer_infos_.size()); + XlaCompiledCpuFunction::set_static_data_arg_index_table( + &jit->static_data_, jit->arg_index_table_.data()); + XlaCompiledCpuFunction::set_static_data_num_args( + &jit->static_data_, jit->arg_index_table_.size()); + XlaCompiledCpuFunction::set_static_data_result_index(&jit->static_data_, + result_index); // Optional metadata is collected and set below. CollectNames(config.feed(), &jit->nonempty_arg_names_, &jit->arg_names_); CollectNames(config.fetch(), &jit->nonempty_result_names_, &jit->result_names_); - jit->static_data_.set_arg_names(jit->arg_names_.data()); - jit->static_data_.set_result_names(jit->result_names_.data()); - jit->static_data_.set_program_shape(jit->program_shape_.get()); + XlaCompiledCpuFunction::set_static_data_arg_names(&jit->static_data_, + jit->arg_names_.data()); + XlaCompiledCpuFunction::set_static_data_result_names( + &jit->static_data_, jit->result_names_.data()); + XlaCompiledCpuFunction::set_static_data_program_shape( + &jit->static_data_, jit->program_shape_.get()); if (cpu_executable->hlo_profiling_enabled()) { - jit->static_data_.set_hlo_profile_printer_data( - &cpu_executable->hlo_profile_printer_data()); - jit->static_data_.set_profile_counters_size( + XlaCompiledCpuFunction::set_static_data_hlo_profile_printer_data( + &jit->static_data_, &cpu_executable->hlo_profile_printer_data()); + XlaCompiledCpuFunction::set_static_data_profile_counters_size( + &jit->static_data_, cpu_executable->hlo_profile_printer_data().profile_counters_size()); } diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.cc b/tensorflow/compiler/tf2xla/xla_op_kernel.cc index 58bd173e61aa3263fae4b494914707833c7a624f..78bc2c94425e00c2b26058daf609d71f1853664e 100644 --- a/tensorflow/compiler/tf2xla/xla_op_kernel.cc +++ b/tensorflow/compiler/tf2xla/xla_op_kernel.cc @@ -93,7 +93,7 @@ TensorShape XlaOpKernelContext::InputShape(absl::string_view name) { } DataType XlaOpKernelContext::input_type(int index) const { - return context_->input(index).dtype(); + return context_->input_dtype(index); } DataType XlaOpKernelContext::InputType(absl::string_view name) { @@ -229,7 +229,8 @@ Status XlaOpKernelContext::ConstantInputAsFloatScalar(int index, double* out) { static Status LiteralToInt64Vector(const xla::LiteralSlice& literal, std::vector* out) { if (literal.shape().rank() != 1) { - return errors::InvalidArgument("value is not 1D"); + return errors::InvalidArgument("value is not 1D, rank: ", + literal.shape().rank()); } int64 size = xla::ShapeUtil::ElementsIn(literal.shape()); if (literal.shape().element_type() == xla::S32) { @@ -353,8 +354,8 @@ Status ReadVariableInputTensor(const Tensor& tensor, DataType type, TF_RET_CHECK(variable != nullptr); TF_RET_CHECK(variable->kind() == XlaResource::kVariable); if (!variable->initialized()) { - return errors::InvalidArgument("Read of uninitialized variable ", - variable->name()); + return errors::FailedPrecondition("Read of uninitialized variable ", + variable->name()); } if (variable->type() != type) { return errors::InvalidArgument( @@ -456,6 +457,11 @@ void XlaOpKernelContext::SetConstantOutput(int index, const Tensor& constant) { SetOutputExpression(index, XlaExpression::Constant(constant)); } +void XlaOpKernelContext::SetTensorListOutput(int index, + const xla::XlaOp& handle) { + SetOutputExpression(index, XlaExpression::TensorList(handle)); +} + void XlaOpKernelContext::SetResourceOutput(int index, XlaResource* resource) { SetOutputExpression(index, XlaExpression::Resource(resource)); } diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.h b/tensorflow/compiler/tf2xla/xla_op_kernel.h index 1858844bc05a6e12abbf07af83cad816590ddd03..e44415f60bff82fb92d0cf4ec81935564a2f083a 100644 --- a/tensorflow/compiler/tf2xla/xla_op_kernel.h +++ b/tensorflow/compiler/tf2xla/xla_op_kernel.h @@ -168,6 +168,9 @@ class XlaOpKernelContext { // Returns an XlaExpression describing the value of 'index'. void SetOutputExpression(int index, const XlaExpression& expression); + // Sets output `index` to the Tensor List `handle`. + void SetTensorListOutput(int index, const xla::XlaOp& handle); + // Status handling. void SetStatus(const Status& status) { context_->SetStatus(status); } Status status() { return context_->status(); } diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.cc b/tensorflow/compiler/tf2xla/xla_op_registry.cc index 14237df69081016817fbd1a5332f22996e7f264d..26314034a18b2a77a3529f0c1af242e29ec69902 100644 --- a/tensorflow/compiler/tf2xla/xla_op_registry.cc +++ b/tensorflow/compiler/tf2xla/xla_op_registry.cc @@ -73,6 +73,11 @@ XlaOpRegistry::~XlaOpRegistry() = default; << " have incompatible allow_resource_types settings."; return false; } + if (x.allow_variant_types != y.allow_variant_types) { + LOG(WARNING) << "Registrations of " << x.name + << " have incompatible allow_variant_types settings."; + return false; + } if (!x.has_device_whitelist && !y.has_device_whitelist) { LOG(WARNING) << "Duplicate registrations of " << x.name << "with no device whitelists."; @@ -289,6 +294,9 @@ void XlaOpRegistry::RegisterCompilationKernels() { if (op_registration->allow_resource_types) { allowed_values->add_type(DT_RESOURCE); } + if (op_registration->allow_variant_types) { + allowed_values->add_type(DT_VARIANT); + } // Don't build KernelDefs that have unsatisfiable type constraints. if (allowed_values->type().empty()) { unsatisfiable_type_constraint = true; @@ -485,6 +493,11 @@ XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::AllowResourceTypes() { return *this; } +XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::AllowVariantTypes() { + registration_->allow_variant_types = true; + return *this; +} + XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::TypeConstraint( absl::string_view attr_name, DataType allowed) { std::set& types = diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.h b/tensorflow/compiler/tf2xla/xla_op_registry.h index 0bdd4a1085445420a5147756daac4a54f4725f11..c5e078a02d1ca6fdd8405ae6556a5205e387421e 100644 --- a/tensorflow/compiler/tf2xla/xla_op_registry.h +++ b/tensorflow/compiler/tf2xla/xla_op_registry.h @@ -47,13 +47,14 @@ extern const char* const DEVICE_XLA_GPU; constexpr std::array kFloatTypes = { {DT_HALF, DT_FLOAT, DT_DOUBLE, DT_BFLOAT16}}; -constexpr std::array kNumericTypes = { +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_BFLOAT16}}; + 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_BOOL}}; + DT_QINT32, DT_INT64, DT_HALF, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64, + DT_COMPLEX128, DT_BOOL}}; constexpr std::array kGpuAllTypes = { {DT_UINT8, DT_QUINT8, DT_UINT32, DT_UINT64, DT_INT8, DT_QINT8, DT_INT32, @@ -211,6 +212,10 @@ class XlaOpRegistry { // allow DT_RESOURCE. bool allow_resource_types = false; + // Should we allow variant types for type attributes? Used by While to + // allow TensorList which is of type DT_VARIANT. + bool allow_variant_types = false; + // Mapping from attribute name to a list of supported types. std::unordered_map> type_constraints; @@ -232,9 +237,9 @@ class XlaOpRegistry { // Returns true if registrations x and y can both be added to the registry. // This is always the case if they refer to different ops. If they refer to - // the same op name, they must: have the same values for compilation_only and - // allow_resource_types; use a device_whitelist; and their - // whitelists must not intersect. + // the same op name, they must: have the same values for compilation_only, + // allow_resource_types and allow_variant_types; use a device_whitelist; and + // their whitelists must not intersect. static bool IsCompatible(const OpRegistration& x, const OpRegistration& y); static Status CompileTimeConstantInputs(const NodeDef& node_def, @@ -292,6 +297,9 @@ class XlaOpRegistrationBuilder { // Allow DT_RESOURCE types for type parameters. XlaOpRegistrationBuilder& AllowResourceTypes(); + // Allow DT_VARIANT types for type parameters. + XlaOpRegistrationBuilder& AllowVariantTypes(); + // Mark 'input_name' as an argument whose value must be known at compile-time. XlaOpRegistrationBuilder& CompileTimeConstantInput( absl::string_view input_name); diff --git a/tensorflow/compiler/xla/BUILD b/tensorflow/compiler/xla/BUILD index 722d1376687efa1c04158e3fd9ce539aac9d0122..ee6f7d5956ede4af99498ca0df5de47150cc5e4d 100644 --- a/tensorflow/compiler/xla/BUILD +++ b/tensorflow/compiler/xla/BUILD @@ -150,9 +150,7 @@ cc_library( visibility = ["//visibility:public"], deps = [ ":status", - "//tensorflow/core:lib", - "//tensorflow/core:lib_internal", - "//tensorflow/stream_executor", + "//tensorflow/stream_executor/lib", ], ) @@ -194,7 +192,7 @@ cc_library( ":types", ":util", "//tensorflow/core:lib", - "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/time", ], ) @@ -717,6 +715,7 @@ cc_library( ":types", ":xla_data_proto", "//tensorflow/core:lib", + "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", ], @@ -741,6 +740,7 @@ cc_library( "//tensorflow/compiler/xla/service:hlo_evaluator", "//tensorflow/compiler/xla/service:shape_inference", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/memory", "@com_google_absl//absl/types:span", ], @@ -824,13 +824,13 @@ cc_library( "debug_options_parsers.h", ], hdrs = ["debug_options_flags.h"], + visibility = [":friends"], deps = [ ":parse_flags_from_env", "//tensorflow/compiler/xla:xla_proto", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", "@com_google_absl//absl/strings", ], ) diff --git a/tensorflow/compiler/xla/array.h b/tensorflow/compiler/xla/array.h index 58cc1575858201b4508d7340cb47e59c4f4c5783..529e7f77cec43f3158fcb59a53efa9a085d7422a 100644 --- a/tensorflow/compiler/xla/array.h +++ b/tensorflow/compiler/xla/array.h @@ -272,6 +272,15 @@ class Array { std::iota(&values_[0], &values_[0] + num_elements(), value); } + // Fills the array with a repeating sequence: + // [value, value + 1, ..., value + length - 1, value, ... ] + void FillRepeatedIota(const T& value, int64 length) { + for (int64 i = 0; i < num_elements(); i += length) { + std::iota(&values_[i], &values_[std::min(i + length, num_elements())], + value); + } + } + // Fills the array with the sequence i*multiplier for i=0,1,... void FillWithMultiples(const T& multiplier) { for (int64 i = 0; i < num_elements(); ++i) { @@ -280,11 +289,11 @@ class Array { } // Fills the array with random normal variables with the specified mean. - void FillRandom(const T& value, const double mean = 0.0, + void FillRandom(const T& stddev, const double mean = 0.0, const int seed = 12345) { std::mt19937 g(seed); std::normal_distribution distribution(mean, - static_cast(value)); + static_cast(stddev)); for (int64 i = 0; i < num_elements(); ++i) { values_[i] = static_cast(distribution(g)); } diff --git a/tensorflow/compiler/xla/client/BUILD b/tensorflow/compiler/xla/client/BUILD index 27c075e8f13f6777af4e837501d97a33034313f5..f5d56e8a9e1f3a05e1039f7cc90194407200f1ab 100644 --- a/tensorflow/compiler/xla/client/BUILD +++ b/tensorflow/compiler/xla/client/BUILD @@ -246,6 +246,7 @@ tf_cc_test( "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_matchers", diff --git a/tensorflow/compiler/xla/client/client.cc b/tensorflow/compiler/xla/client/client.cc index 43127cae1e5d81521003a28288e27d291e33c9b9..4f020bcec2756a328755d86ab04154d54f532465 100644 --- a/tensorflow/compiler/xla/client/client.cc +++ b/tensorflow/compiler/xla/client/client.cc @@ -278,53 +278,51 @@ StatusOr> Client::Execute( const XlaComputation& computation, absl::Span arguments, const ExecutionOptions* execution_options, ExecutionProfile* execution_profile) { - if (execution_options != nullptr && - execution_options->device_handles_size() > 1) { - std::vector computation_instances = { - XlaComputationInstance{ - computation, - std::vector(arguments.begin(), arguments.end()), - *execution_options, execution_profile}}; - TF_ASSIGN_OR_RETURN(auto results, ExecuteParallel(computation_instances)); - // The result selection is a bit hacky, but better than assuming it is - // device 0. - // - // TODO(b/118493728): Allow Execute to return one result per computation. - for (int64 i = 0; i < results.size(); i++) { - TF_ASSIGN_OR_RETURN(const Shape& shape, GetShape(*results[i])); - if (!ShapeUtil::IsEmptyTuple(shape)) { - VLOG(3) << "Fetching result from device " << i << ": " - << ShapeUtil::HumanString(shape); - return std::move(results[i]); - } + // Create an ExecutionOptions if necessary, or set its DeviceHandles. + absl::optional options_storage; + if (!execution_options || execution_options->device_handles().empty()) { + if (execution_options) { + options_storage.emplace(*execution_options); + } else { + options_storage.emplace(CreateDefaultExecutionOptions()); } - TF_RET_CHECK(!results.empty()); - VLOG(1) << "Defaulting to device 0 result"; - return std::move(results[0]); - } - - // The argument shapes affect how the computation is compiled. - std::vector arg_shapes(arguments.size()); - for (int i = 0; i < arguments.size(); i++) { - TF_ASSIGN_OR_RETURN(arg_shapes[i], GetShape(*arguments[i])); - } - - TF_ASSIGN_OR_RETURN(auto handle, - Compile(computation, arg_shapes, execution_options)); - - TF_ASSIGN_OR_RETURN(auto result, - Execute(handle, arguments, execution_profile)); - - if (execution_profile != nullptr) { - if (VLOG_IS_ON(1)) { - TF_ASSIGN_OR_RETURN( - auto execution_stats, - ExecutionStatsAsString(computation, *execution_profile)); - VLOG(1) << execution_stats; + execution_options = &*options_storage; + + TF_ASSIGN_OR_RETURN(auto device_handles, + GetDeviceHandles(/*device_count=*/1)); + TF_RET_CHECK(!device_handles.empty()); + *options_storage->add_device_handles() = std::move(device_handles[0]); + } + + std::vector computation_instances = { + XlaComputationInstance{ + computation, + std::vector(arguments.begin(), arguments.end()), + *execution_options, execution_profile}}; + + // Instead of invoking Compile() and Execute(), invoke + // Service::ExecuteParallel() to execute our one computation. Compile() + // caches the executable forever, which isn't what we want. + VLOG(1) << "Making ExecuteParallel request: " + << execution_options->DebugString(); + TF_ASSIGN_OR_RETURN(auto results, ExecuteParallel(computation_instances)); + VLOG(1) << "ExecuteParallel request done."; + + // The result selection is a bit hacky, but better than assuming it is + // device 0. + // + // TODO(b/118493728): Allow Execute to return one result per computation. + for (int64 i = 0; i < results.size(); i++) { + TF_ASSIGN_OR_RETURN(const Shape& shape, GetShape(*results[i])); + if (!ShapeUtil::IsEmptyTuple(shape)) { + VLOG(3) << "Fetching result from device " << i << ": " + << ShapeUtil::HumanString(shape); + return std::move(results[i]); } } - - return std::move(result); + TF_RET_CHECK(!results.empty()); + VLOG(1) << "Defaulting to device 0 result"; + return std::move(results[0]); } StatusOr>> Client::ExecuteParallel( diff --git a/tensorflow/compiler/xla/client/client.h b/tensorflow/compiler/xla/client/client.h index d0ac4703c632e0e01d3c8911594b46fedf28930d..eff8713ac340e82ee7633f1f078334ba73b67b2f 100644 --- a/tensorflow/compiler/xla/client/client.h +++ b/tensorflow/compiler/xla/client/client.h @@ -52,6 +52,12 @@ class Client { // need to live beyond this call.) // * If execution_options.device_handles should be empty. If you need // non-empty device handles, call 'Execute' instead. + // + // TODO(b/122731460): This call caches the resulting Executable in the Service + // *forever*. If you're only going to run the computation once, you may want + // to call the Execute(const XlaComputation&) overload. If you're going to + // run the computation more than once but you want control over when the + // Executable is unloaded, use the LocalClient API. StatusOr Compile( const XlaComputation& computation, absl::Span argument_shapes, @@ -76,6 +82,10 @@ class Client { // device is chosen by the service. // * If execution_profile is not nullptr then the pointed-to ExecutionProfile // will be filled with profile data from the execution. + // + // TODO(b/122731460): The given computation is compiled and then thrown away + // immediately after it's run. If you want control over how long the + // resulting Executable lives, use the LocalClient API. StatusOr> Execute( const XlaComputation& computation, absl::Span arguments, diff --git a/tensorflow/compiler/xla/client/executable_build_options.cc b/tensorflow/compiler/xla/client/executable_build_options.cc index 1f594e551af381d7537e947892cbf7e0b5b3b861..ec0e08975926f36c36c854f83a40b374b12a09a4 100644 --- a/tensorflow/compiler/xla/client/executable_build_options.cc +++ b/tensorflow/compiler/xla/client/executable_build_options.cc @@ -58,6 +58,12 @@ const Shape* ExecutableBuildOptions::result_layout() const { return result_layout_set_ ? &result_layout_ : nullptr; } +ExecutableBuildOptions& ExecutableBuildOptions::set_num_replicas( + int num_replicas) { + num_replicas_ = num_replicas; + return *this; +} + string ExecutableBuildOptions::ToString() const { string result_layout = "nullopt"; if (result_layout_set_) { @@ -65,8 +71,9 @@ string ExecutableBuildOptions::ToString() const { } return absl::StrFormat( "ExecutableBuildOptions{device_ordinal=%d, result_layout=%s, " - "generate_hlo_graph=%s}", - device_ordinal_, result_layout, debug_options().xla_generate_hlo_graph()); + "generate_hlo_graph=%s, num_replicas=%d}", + device_ordinal_, result_layout, debug_options().xla_generate_hlo_graph(), + num_replicas_); } } // namespace xla diff --git a/tensorflow/compiler/xla/client/executable_build_options.h b/tensorflow/compiler/xla/client/executable_build_options.h index a58090253bfac7779e4b61bc7231a0f0d945cc00..1d85fb34304b95d1fccdb0b0d6a7a65e739fae18 100644 --- a/tensorflow/compiler/xla/client/executable_build_options.h +++ b/tensorflow/compiler/xla/client/executable_build_options.h @@ -67,12 +67,18 @@ class ExecutableBuildOptions { // debugging. string ToString() const; + // The number of replicas of this computation that are to be executed. + // Defaults to 1. + int num_replicas() const { return num_replicas_; } + ExecutableBuildOptions& set_num_replicas(int num_replicas); + private: int device_ordinal_ = -1; Shape result_layout_; bool result_layout_set_ = false; absl::optional debug_options_; DeviceMemoryAllocator* device_allocator_ = nullptr; + int num_replicas_ = 1; }; } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/BUILD b/tensorflow/compiler/xla/client/lib/BUILD index 966a5819819c3494a20f67825c67d7496b35fa96..0abf546c14b4ff76d1421a93dd6a14270efab79f 100644 --- a/tensorflow/compiler/xla/client/lib/BUILD +++ b/tensorflow/compiler/xla/client/lib/BUILD @@ -34,6 +34,21 @@ cc_library( ], ) +xla_test( + name = "arithmetic_test", + srcs = ["arithmetic_test.cc"], + deps = [ + ":arithmetic", + "//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 = "cholesky", srcs = ["cholesky.cc"], @@ -50,7 +65,6 @@ cc_library( "//tensorflow/compiler/xla/client/lib:loops", "//tensorflow/compiler/xla/client/lib:matrix", "//tensorflow/compiler/xla/client/lib:slicing", - "//tensorflow/compiler/xla/client/lib:triangular_solve", "//tensorflow/core:lib", ], ) @@ -76,6 +90,39 @@ xla_test( ], ) +cc_library( + name = "comparators", + srcs = ["comparators.cc"], + hdrs = ["comparators.h"], + deps = [ + ":constants", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/client:xla_computation", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", + ], +) + +xla_test( + name = "comparators_test", + srcs = ["comparators_test.cc"], + deps = [ + ":comparators", + ":constants", + "//tensorflow/compiler/xla:shape_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", + "@com_google_absl//absl/container:inlined_vector", + ], +) + cc_library( name = "constants", srcs = ["constants.cc"], @@ -93,7 +140,6 @@ cc_library( xla_test( name = "constants_test", srcs = ["constants_test.cc"], - tags = ["enable_for_xla_interpreter"], deps = [ ":constants", "//tensorflow/compiler/xla:test", @@ -137,6 +183,7 @@ cc_library( srcs = ["math.cc"], hdrs = ["math.h"], deps = [ + ":arithmetic", ":constants", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -147,7 +194,23 @@ cc_library( xla_test( name = "math_test", srcs = ["math_test.cc"], - tags = ["enable_for_xla_interpreter"], + deps = [ + ":math", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_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", + ], +) + +xla_test( + name = "math_exhaustive_test", + srcs = ["math_exhaustive_test.cc"], + shard_count = 16, deps = [ ":math", "//tensorflow/compiler/xla:literal_util", @@ -167,13 +230,18 @@ cc_library( deps = [ ":arithmetic", ":constants", + ":slicing", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:xla_builder", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", ], ) @@ -181,16 +249,19 @@ cc_library( xla_test( name = "matrix_test", srcs = ["matrix_test.cc"], - tags = ["enable_for_xla_interpreter"], deps = [ ":matrix", ":slicing", + "//tensorflow/compiler/xla:status", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", "//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", + "@com_google_absl//absl/strings", ], ) @@ -229,7 +300,6 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:xla_builder", - "//tensorflow/core:lib", "@com_google_absl//absl/base", ], ) @@ -281,12 +351,7 @@ cc_library( srcs = ["slicing.cc"], hdrs = ["slicing.h"], deps = [ - "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:status_macros", - "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:xla_builder", "@com_google_absl//absl/types:span", ], @@ -295,13 +360,11 @@ cc_library( xla_test( name = "slicing_test", srcs = ["slicing_test.cc"], - tags = ["enable_for_xla_interpreter"], deps = [ ":slicing", "//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", @@ -313,6 +376,7 @@ cc_library( srcs = ["sorting.cc"], hdrs = ["sorting.h"], deps = [ + ":comparators", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", @@ -324,12 +388,10 @@ cc_library( xla_test( name = "sorting_test", srcs = ["sorting_test.cc"], - tags = ["enable_for_xla_interpreter"], deps = [ ":sorting", "//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", @@ -354,7 +416,6 @@ xla_test( srcs = ["quantize_test.cc"], # TODO(b/122119490): re-enable TAP after fixing. tags = [ - "enable_for_xla_interpreter", "notap", ], deps = [ @@ -391,51 +452,48 @@ cc_library( ) cc_library( - name = "triangular_solve", - srcs = ["triangular_solve.cc"], - hdrs = ["triangular_solve.h"], + name = "self_adjoint_eigen", + srcs = ["self_adjoint_eigen.cc"], + hdrs = ["self_adjoint_eigen.h"], deps = [ - "//tensorflow/compiler/xla:literal", + ":arithmetic", + ":constants", + ":loops", + ":math", + ":matrix", + ":slicing", + "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", - "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:xla_builder", - "//tensorflow/compiler/xla/client:xla_computation", - "//tensorflow/compiler/xla/client/lib:constants", - "//tensorflow/compiler/xla/client/lib:math", - "//tensorflow/compiler/xla/client/lib:matrix", - "//tensorflow/compiler/xla/client/lib:slicing", "//tensorflow/core:lib", ], ) xla_test( - name = "triangular_solve_test", - srcs = ["triangular_solve_test.cc"], - tags = [ - "enable_for_xla_interpreter", - "noasan", # sometimes times out, http://b/78650012 - ], + name = "self_adjoint_eigen_test", + size = "medium", + srcs = ["self_adjoint_eigen_test.cc"], + real_hardware_only = True, + shard_count = 10, + tags = ["optonly"], deps = [ - ":math", + ":arithmetic", + ":constants", ":matrix", - ":triangular_solve", + ":self_adjoint_eigen", "//tensorflow/compiler/xla:array2d", + "//tensorflow/compiler/xla:array3d", "//tensorflow/compiler/xla:literal", - "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:test", - "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:global_data", - "//tensorflow/compiler/xla/client:local_client", "//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:lib", "//tensorflow/core:test", ], ) diff --git a/tensorflow/compiler/xla/client/lib/arithmetic.cc b/tensorflow/compiler/xla/client/lib/arithmetic.cc index 33ff3971d7277e619db4bec99ffdf9c0c9d230ad..3b875135af29f142463ffd783bfeaadc61ada1af 100644 --- a/tensorflow/compiler/xla/client/lib/arithmetic.cc +++ b/tensorflow/compiler/xla/client/lib/arithmetic.cc @@ -123,4 +123,64 @@ XlaOp Any(XlaOp predicates) { }); } +namespace { + +XlaOp ArgMinMax(XlaOp input, PrimitiveType output_type, int axis, bool is_min) { + XlaBuilder* builder = input.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape input_shape, builder->GetShape(input)); + XlaOp init_value; + XlaComputation reducer; + if (is_min) { + init_value = MaxValue(builder, input_shape.element_type()); + reducer = CreateScalarMinComputation(input_shape.element_type(), builder); + } else { + init_value = MinValue(builder, input_shape.element_type()); + reducer = CreateScalarMaxComputation(input_shape.element_type(), builder); + } + + XlaOp input_max = Reduce(input, init_value, reducer, + /*dimensions_to_reduce=*/{axis}); + std::vector broadcast_dims(input_shape.rank() - 1); + std::iota(broadcast_dims.begin(), broadcast_dims.begin() + axis, 0); + std::iota(broadcast_dims.begin() + axis, broadcast_dims.end(), axis + 1); + // Compute a mask that has 1s for elements equal to the maximum. + XlaOp partial_mask = + ConvertElementType(Eq(input, input_max, broadcast_dims), output_type); + + // In order to make identity elements for a bitwise And, we: + // Left shift the 1 to the leftmost bit, yielding 0x10...0 + // Arithmetic right shift the 1 back to the rightmost bit, yielding + // 0xFF...F + int32 bits_in_type = + ShapeUtil::ByteSizeOfPrimitiveType(output_type) * 8 - 1; + XlaOp shift_amount = ConstantR0WithType(builder, output_type, bits_in_type); + XlaOp full_mask = ShiftRightArithmetic( + ShiftLeft(partial_mask, shift_amount), shift_amount); + + // And with the vector [0, 1, 2, ...] to convert each 0xFF...F into its + // index. + + const int64 axis_size = ShapeUtil::GetDimension(input_shape, axis); + XlaOp iota = Iota(builder, output_type, axis_size); + XlaOp product = And(full_mask, iota, /*broadcast_dimensions=*/{axis}); + + // If there are multiple maximum elements, choose the one with the highest + // index. + return Reduce(product, MinValue(builder, output_type), + CreateScalarMaxComputation(output_type, builder), + /*dimensions_to_reduce=*/{axis}); + }); +} + +} // namespace + +XlaOp ArgMax(XlaOp input, PrimitiveType output_type, int axis) { + return ArgMinMax(input, output_type, axis, /*is_min=*/false); +} + +XlaOp ArgMin(XlaOp input, PrimitiveType output_type, int axis) { + return ArgMinMax(input, output_type, axis, /*is_min=*/true); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/arithmetic.h b/tensorflow/compiler/xla/client/lib/arithmetic.h index 632e8cc8bc64fad236a0226c6e93079aadde7050..d4a7812c441c351b121e5d72faf9642b06728b18 100644 --- a/tensorflow/compiler/xla/client/lib/arithmetic.h +++ b/tensorflow/compiler/xla/client/lib/arithmetic.h @@ -57,6 +57,14 @@ XlaComputation CreateScalarOrComputation(PrimitiveType type, // Note: if predicates is zero-sized, Any() vacuously returns false. XlaOp Any(XlaOp predicates); +// Returns the argmax of `input` along `axis`. `output_type` is the type to +// use for the output. +XlaOp ArgMax(XlaOp input, PrimitiveType output_type, int axis); + +// Returns the argmin of `input` along `axis`. `output_type` is the type to +// use for the output. +XlaOp ArgMin(XlaOp input, PrimitiveType output_type, int axis); + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_ARITHMETIC_H_ diff --git a/tensorflow/compiler/xla/client/lib/arithmetic_test.cc b/tensorflow/compiler/xla/client/lib/arithmetic_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..a13839f9db89b9c07f2465867a503ef2193f8160 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/arithmetic_test.cc @@ -0,0 +1,67 @@ +/* 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/arithmetic.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" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { +namespace { + +using ArithmeticTest = ClientLibraryTestBase; + +XLA_TEST_F(ArithmeticTest, ArgMinR2Axis0) { + XlaBuilder builder(TestName()); + auto x = ConstantR2(&builder, {{1, 7, 4}, {6, 3, 5}, {8, 3, 3}}); + ArgMin(x, S32, /*axis=*/0); + + std::vector expected = {0, 2, 2}; + ComputeAndCompareR1(&builder, expected, {}); +} + +XLA_TEST_F(ArithmeticTest, ArgMinR2Axis1) { + XlaBuilder builder(TestName()); + auto x = ConstantR2(&builder, {{1, 7, 4}, {6, 3, 5}, {8, 3, 3}}); + ArgMin(x, S32, /*axis=*/1); + + std::vector expected = {0, 1, 2}; + ComputeAndCompareR1(&builder, expected, {}); +} + +XLA_TEST_F(ArithmeticTest, ArgMaxR2Axis0) { + XlaBuilder builder(TestName()); + auto x = ConstantR2(&builder, {{1, 7, 4}, {6, 3, 5}, {8, 3, 3}}); + ArgMax(x, S32, /*axis=*/0); + + std::vector expected = {2, 0, 1}; + ComputeAndCompareR1(&builder, expected, {}); +} + +XLA_TEST_F(ArithmeticTest, ArgMaxR2Axis1) { + XlaBuilder builder(TestName()); + auto x = ConstantR2(&builder, {{1, 7, 4}, {6, 3, 5}, {8, 3, 3}}); + ArgMax(x, S32, /*axis=*/1); + + std::vector expected = {1, 0, 0}; + ComputeAndCompareR1(&builder, expected, {}); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/cholesky.cc b/tensorflow/compiler/xla/client/lib/cholesky.cc index e100d47922efb6655e40a19b5c4cbb2190fac6d9..bb41f9932d1cc62b62d37fea2c10fbfeaa0bd15e 100644 --- a/tensorflow/compiler/xla/client/lib/cholesky.cc +++ b/tensorflow/compiler/xla/client/lib/cholesky.cc @@ -23,7 +23,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/lib/math.h" #include "tensorflow/compiler/xla/client/lib/matrix.h" #include "tensorflow/compiler/xla/client/lib/slicing.h" -#include "tensorflow/compiler/xla/client/lib/triangular_solve.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/primitive_util.h" @@ -68,29 +67,26 @@ XlaOp CholeskyUnblocked(XlaOp a, PrecisionConfig::Precision precision) { auto body_fn = [&](XlaOp i, absl::Span loop_vars, XlaBuilder* body_builder) -> StatusOr> { - Shape col_shape; - Shape row_shape; - for (int64 d : major_dims) { - row_shape.add_dimensions(d); - col_shape.add_dimensions(d); - } - row_shape.add_dimensions(1); - row_shape.add_dimensions(n); - row_shape.set_element_type(a_shape.element_type()); - auto mask_zeros_row = Zeros(body_builder, row_shape); - - col_shape.add_dimensions(n); - col_shape.add_dimensions(1); - col_shape.set_element_type(a_shape.element_type()); - auto mask_zeros_col = Zeros(body_builder, col_shape); - - std::vector mask_vector(n); - std::iota(mask_vector.begin(), mask_vector.end(), 0); - auto mask_range = ConstantR1(body_builder, mask_vector); + std::vector row_shape_dims(major_dims.begin(), major_dims.end()); + std::vector col_shape_dims(major_dims.begin(), major_dims.end()); + row_shape_dims.push_back(1); + row_shape_dims.push_back(n); + auto mask_zeros_row = + Zeros(body_builder, + ShapeUtil::MakeShape(a_shape.element_type(), row_shape_dims)); + + col_shape_dims.push_back(n); + col_shape_dims.push_back(1); + auto mask_zeros_col = + Zeros(body_builder, + ShapeUtil::MakeShape(a_shape.element_type(), col_shape_dims)); + auto mask_range_row = - Broadcast(Reshape(mask_range, {0}, {1, n}), major_dims); + Iota(body_builder, ShapeUtil::MakeShape(S32, row_shape_dims), + /*iota_dimension=*/n_dims - 1); auto mask_range_col = - Broadcast(Reshape(mask_range, {0}, {n, 1}), major_dims); + Iota(body_builder, ShapeUtil::MakeShape(S32, col_shape_dims), + /*iota_dimension=*/n_dims - 2); auto body_a = loop_vars[0]; auto body_l = loop_vars[1]; @@ -197,12 +193,12 @@ XlaOp Cholesky(XlaOp a, int64 block_size, // l[i+k:, i:i+k] = // trsm_right_transpose(l[i:i+k, i:i+k], a[i+k:, i:i+k]) auto panel = SliceInMinorDims(a, {i + k, i}, {n, i + k}); - auto update = TriangularSolve(factorized, panel, - /*left_side=*/false, - /*lower=*/true, - /*transpose_a=*/true, - /*conjugate_a=*/false, - /*block_size=*/block_size); + auto update = + TriangularSolve(factorized, panel, + /*left_side=*/false, + /*lower=*/true, + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::TRANSPOSE); l = UpdateSliceInMinorDims(l, update, {i + k, i}); } } diff --git a/tensorflow/compiler/xla/client/lib/cholesky_test.cc b/tensorflow/compiler/xla/client/lib/cholesky_test.cc index ba9580a3d32225625acc1447344b7d2c16c5d8a5..095dd4fbf8b7c90047c4428b50c626c16e9c1e94 100644 --- a/tensorflow/compiler/xla/client/lib/cholesky_test.cc +++ b/tensorflow/compiler/xla/client/lib/cholesky_test.cc @@ -157,10 +157,10 @@ XLA_TEST_P(RandomCholeskyTest, Random) { xla::ErrorSpec(1e-4, 1e-4)); } -INSTANTIATE_TEST_CASE_P(RandomCholeskyTestInstance, RandomCholeskyTest, - ::testing::Values(CholeskyTestCase{1, 1}, - CholeskyTestCase{1, 2}, - CholeskyTestCase{10, 5}, - CholeskyTestCase{2, 20})); +INSTANTIATE_TEST_SUITE_P(RandomCholeskyTestInstance, RandomCholeskyTest, + ::testing::Values(CholeskyTestCase{1, 1}, + CholeskyTestCase{1, 2}, + CholeskyTestCase{10, 5}, + CholeskyTestCase{2, 20})); } // namespace diff --git a/tensorflow/compiler/xla/client/lib/comparators.cc b/tensorflow/compiler/xla/client/lib/comparators.cc new file mode 100644 index 0000000000000000000000000000000000000000..c620c9841a5146618e3a142adeb3fe2da525950a --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/comparators.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/client/lib/comparators.h" + +#include +#include +#include + +#include "absl/strings/str_cat.h" +#include "absl/types/span.h" +#include "tensorflow/compiler/xla/client/lib/constants.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/client/xla_computation.h" +#include "tensorflow/compiler/xla/primitive_util.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { +namespace { + +using XlaOpGenerator = XlaOp (*)(const XlaOp&, const XlaOp&, + absl::Span); + +XlaOp BitcastConvertFloatingPointToIntegral(const XlaOp& value, + int64 bit_width) { + PrimitiveType signed_type; + PrimitiveType unsigned_type; + XlaOp max_value; + switch (bit_width) { + case 16: + max_value = + ConstantR0(value.builder(), + static_cast(std::numeric_limits::max())); + signed_type = S16; + unsigned_type = U16; + break; + case 32: + max_value = + ConstantR0(value.builder(), + static_cast(std::numeric_limits::max())); + signed_type = S32; + unsigned_type = U32; + break; + case 64: + max_value = + ConstantR0(value.builder(), + static_cast(std::numeric_limits::max())); + signed_type = S64; + unsigned_type = U64; + break; + default: + return value.builder()->ReportError( + InvalidArgument("Invalid bit width %lld for Comparator floating " + "point parameter.", + bit_width)); + } + // Switch from a floating point value to a integer value in such a way that + // when using the integer value to compare, we get the same result for normal + // values, and -Nan is treated as the smallest value, and Nan is treated as + // the largest value. + // If f is a float, and + // x = bit_cast(f); + // y = x < 0 ? numeric_limits::max() - x : x; + // then y is ordered as an int32 such that finite values have the obvious + // order, -0 is ordered before 0, and -NaN and NaN appear at the beginning + // and end of the ordering. + // Note that in order to avoid -x to overflow, we calculate + // numeric_limits::max() - x as unsigned, and then convert back to + // signed. + auto signed_value = BitcastConvertType(value, signed_type); + auto unsigned_value = BitcastConvertType(value, unsigned_type); + auto flipped_value = + BitcastConvertType(Sub(max_value, unsigned_value), signed_type); + auto is_negative = Lt(signed_value, Zero(value.builder(), signed_type)); + return Select(is_negative, flipped_value, signed_value); +} + +XlaComputation CreateScalarComparisonComputation( + const string& name, const std::vector& operand_types, + XlaBuilder* builder, XlaOpGenerator generator) { + // Create a default computation where we compare only the first two + // parameters of type 'operand_types[0]'. + auto b = builder->CreateSubBuilder(name); + if (operand_types.empty()) { + b->ReportError(InvalidArgument("operand_types should not be empty")); + return b->BuildAndNoteError(); + } + + int64 parameter_count = 0; + XlaOp first_lhs_param; + XlaOp first_rhs_param; + + // For each type in 'operand_types' we create two parameters of this type. The + // idea is that this computation can be used by n-ary Sort, and potentially + // should support comparing also the other operands of sort. In this default + // computation, however, we will not actually use any parameters except the + // first two. + for (auto operand_type : operand_types) { + auto scalar_shape = ShapeUtil::MakeShape(operand_type, {}); + auto lhs_param = Parameter(b.get(), parameter_count * 2, scalar_shape, + absl::StrCat("p.", parameter_count, ".lhs")); + auto rhs_param = Parameter(b.get(), parameter_count * 2 + 1, scalar_shape, + absl::StrCat("p.", parameter_count, ".rhs")); + if (parameter_count == 0) { + first_lhs_param = lhs_param; + first_rhs_param = rhs_param; + } + ++parameter_count; + } + if (primitive_util::IsFloatingPointType(operand_types[0])) { + PrimitiveType compare_type = operand_types[0]; + // Special-case handling for BF16. We currently do not support direct + // comparisons with BF16, so we convert to F32 and then use the F32 + // comparison logic. + if (compare_type == BF16) { + compare_type = F32; + first_lhs_param = ConvertElementType(first_lhs_param, F32); + first_rhs_param = ConvertElementType(first_rhs_param, F32); + } + int64 bit_width = primitive_util::BitWidth(compare_type); + first_lhs_param = + BitcastConvertFloatingPointToIntegral(first_lhs_param, bit_width); + first_rhs_param = + BitcastConvertFloatingPointToIntegral(first_rhs_param, bit_width); + } + generator(first_lhs_param, first_rhs_param, {}); + return b->BuildAndNoteError(); +} +} // namespace + +// Creates a scalar less-than computation and returns it. +XlaComputation CreateScalarLtComputation( + const std::vector& operand_types, XlaBuilder* builder) { + return CreateScalarComparisonComputation("compare-less-than", operand_types, + builder, Lt); +} + +// Creates a scalar greater-than computation and returns it. +XlaComputation CreateScalarGtComputation( + const std::vector& operand_types, XlaBuilder* builder) { + return CreateScalarComparisonComputation("compare-greater-than", + operand_types, builder, Gt); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/comparators.h b/tensorflow/compiler/xla/client/lib/comparators.h new file mode 100644 index 0000000000000000000000000000000000000000..cbcfc227dd495537f59bf0a9090bad8ade15da62 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/comparators.h @@ -0,0 +1,47 @@ +/* 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_CLIENT_LIB_COMPARATORS_H_ +#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_COMPARATORS_H_ + +#include + +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/client/xla_computation.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { + +// Creates a scalar less-than computation and returns it. The created +// computation has 2 * 'operand_types.size()' many parameters, where parameters +// 2 * i and 2 * i + 1 are a scalar with primitive type 'operand_types[i]'. The +// computation compares the first two parameters. For floating point types, a +// total order is created where +// -NaN < -infinity < ... < -0 < 0 < ... < infinity < NaN +XlaComputation CreateScalarLtComputation( + const std::vector& operand_types, XlaBuilder* builder); + +// Creates a scalar greater-than computation and returns it. The created +// computation has 2 * 'operand_types.size()' many parameters, where parameters +// 2 * i and 2 * i + 1 are a scalar with primitive type 'operand_types[i]'. The +// computation compares the first two parameters. For floating point types, a +// total order is created where +// NaN > infinity > ... > 0 > -0 > ... > -infinity > -NaN +XlaComputation CreateScalarGtComputation( + const std::vector& operand_types, XlaBuilder* builder); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_COMPARATORS_H_ diff --git a/tensorflow/compiler/xla/client/lib/comparators_test.cc b/tensorflow/compiler/xla/client/lib/comparators_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..598956803b34702b1e095a342648d348fa350b29 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/comparators_test.cc @@ -0,0 +1,149 @@ +/* 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/comparators.h" + +#include +#include + +#include "absl/container/inlined_vector.h" +#include "tensorflow/compiler/xla/client/lib/constants.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/primitive_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" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { +namespace { + +class ComparatorsTest : public ClientLibraryTestBase { + public: + ComparatorsTest() : builder_(TestName()) {} + XlaBuilder* builder() { return &builder_; } + + private: + XlaBuilder builder_; +}; + +template < + PrimitiveType type, + typename T = typename primitive_util::PrimitiveTypeToNative::type> +void BuildComparatorAndComparisons(ComparatorsTest* test, + bool compare_less_than, + absl::InlinedVector* expected) { + auto compare = compare_less_than + ? CreateScalarLtComputation({type}, test->builder()) + : CreateScalarGtComputation({type}, test->builder()); + + auto negative_nan = ConstantR0( + test->builder(), -T(std::numeric_limits::quiet_NaN())); + auto positive_nan = ConstantR0(test->builder(), + T(std::numeric_limits::quiet_NaN())); + auto negative_zero = ConstantR0(test->builder(), T(-0.)); + auto positive_zero = ConstantR0(test->builder(), T(0.)); + auto negative_infinity = MinValue(test->builder(), type); + auto positive_infinity = MaxValue(test->builder(), type); + + // List the values in the expected sorting order from smallest to largest. + std::vector all_constants{negative_nan, negative_infinity, + negative_zero, positive_zero, + positive_infinity, positive_nan}; + + // Do pairwise comparisons. + std::vector all_comparisons; + for (const XlaOp& lhs_constant : all_constants) { + for (const XlaOp& rhs_constant : all_constants) { + all_comparisons.push_back(Broadcast( + Call(test->builder(), compare, {lhs_constant, rhs_constant}), {1})); + } + } + + // Concantenate the comparison results. + ConcatInDim(test->builder(), all_comparisons, 0); + + // If we use less-than comparisons, we expect the comparison to result in true + // if the lhs value to be compared appears earlier in 'all_constants' than the + // rhs value. Likewise, if we use greater-than comparisons, we expect the + // comparison to return true if the rhs value appears earlier in + // 'all_constants' than the lhs value. + expected->clear(); + for (int i = 0; i < all_constants.size(); ++i) { + for (int j = 0; j < all_constants.size(); ++j) { + expected->push_back(compare_less_than ? i < j : i > j); + } + } +} + +XLA_TEST_F(ComparatorsTest, CompareLtBF16) { + absl::InlinedVector expected; + BuildComparatorAndComparisons(this, /*compare_less_than=*/true, + &expected); + ComputeAndCompareR1(builder(), expected, {}); +} + +XLA_TEST_F(ComparatorsTest, CompareGtBF16) { + absl::InlinedVector expected; + BuildComparatorAndComparisons(this, /*compare_less_than=*/false, + &expected); + ComputeAndCompareR1(builder(), expected, {}); +} + +XLA_TEST_F(ComparatorsTest, CompareLtF16) { + absl::InlinedVector expected; + BuildComparatorAndComparisons(this, /*compare_less_than=*/true, + &expected); + ComputeAndCompareR1(builder(), expected, {}); +} + +XLA_TEST_F(ComparatorsTest, CompareGtF16) { + absl::InlinedVector expected; + BuildComparatorAndComparisons(this, /*compare_less_than=*/false, + &expected); + ComputeAndCompareR1(builder(), expected, {}); +} + +XLA_TEST_F(ComparatorsTest, CompareLtF32) { + absl::InlinedVector expected; + BuildComparatorAndComparisons(this, /*compare_less_than=*/true, + &expected); + ComputeAndCompareR1(builder(), expected, {}); +} + +XLA_TEST_F(ComparatorsTest, CompareGtF32) { + absl::InlinedVector expected; + BuildComparatorAndComparisons(this, /*compare_less_than=*/false, + &expected); + ComputeAndCompareR1(builder(), expected, {}); +} + +XLA_TEST_F(ComparatorsTest, CompareLtF64) { + absl::InlinedVector expected; + BuildComparatorAndComparisons(this, /*compare_less_than=*/true, + &expected); + ComputeAndCompareR1(builder(), expected, {}); +} + +XLA_TEST_F(ComparatorsTest, CompareGtF64) { + absl::InlinedVector expected; + BuildComparatorAndComparisons(this, /*compare_less_than=*/false, + &expected); + ComputeAndCompareR1(builder(), expected, {}); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/constants.cc b/tensorflow/compiler/xla/client/lib/constants.cc index 1ada7b4a964ccf7ca400b937abbe425bef083468..d0efd79bbebb705cec086853af4521e365999591 100644 --- a/tensorflow/compiler/xla/client/lib/constants.cc +++ b/tensorflow/compiler/xla/client/lib/constants.cc @@ -100,4 +100,28 @@ XlaOp MaxFiniteValue(XlaBuilder* builder, PrimitiveType type) { } } +XlaOp NanValue(XlaBuilder* builder, PrimitiveType type) { + return builder->ReportErrorOrReturn([&]() -> StatusOr { + switch (type) { + case F16: + return ConstantR0( + builder, Eigen::NumTraits::quiet_NaN()); + case BF16: + return ConstantR0( + builder, bfloat16(std::numeric_limits::quiet_NaN())); + case F32: + return ConstantR0(builder, + std::numeric_limits::quiet_NaN()); + case F64: + return ConstantR0(builder, + std::numeric_limits::quiet_NaN()); + default: + return InvalidArgument( + "Operand to NanValue was %s, but must be a real-valued " + "floating-point type.", + PrimitiveType_Name(type)); + } + }); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/constants.h b/tensorflow/compiler/xla/client/lib/constants.h index 81624614c1e3599dfe116eb61d9e2edcd5230684..77e7ca60c378c48cc66523baa452d525b8144847 100644 --- a/tensorflow/compiler/xla/client/lib/constants.h +++ b/tensorflow/compiler/xla/client/lib/constants.h @@ -56,6 +56,8 @@ XlaOp ConstantR0WithType(XlaBuilder* builder, PrimitiveType type, T value) { return ConstantR0(builder, static_cast(value)); case C64: return ConstantR0(builder, static_cast(value)); + case C128: + return ConstantR0(builder, static_cast(value)); case U8: return ConstantR0(builder, static_cast(value)); case U32: @@ -88,6 +90,27 @@ XlaOp ScalarLike(XlaOp prototype, T value) { }); } +// Returns an array or scalar containing copies of `value` cast to the same +// run-type type as `prototype` and broadcast to the same dimensions as +// `prototype`. +// +// If `prototype` is not a scalar or array, returns an error. +template +XlaOp FullLike(XlaOp prototype, T value) { + XlaBuilder* builder = prototype.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(prototype)); + if (ShapeUtil::IsScalar(shape) || shape.IsArray()) { + return Broadcast(ScalarLike(prototype, value), shape.dimensions()); + } else { + return InvalidArgument( + "Prototype shape for BroadcastConstantLike must be a scalar or " + "array, but was %s", + shape.ToString()); + } + }); +} + // Returns a scalar with value '0' of 'type'. XlaOp Zero(XlaBuilder* builder, PrimitiveType type); @@ -119,6 +142,9 @@ XlaOp MaxValue(XlaBuilder* builder, PrimitiveType type); // Returns the maximum representable finite value for 'type'. XlaOp MaxFiniteValue(XlaBuilder* builder, PrimitiveType type); +// Returns a nan for the given type. Only valid for real-valued fp types. +XlaOp NanValue(XlaBuilder* builder, PrimitiveType type); + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_CONSTANTS_H_ diff --git a/tensorflow/compiler/xla/client/lib/constants_test.cc b/tensorflow/compiler/xla/client/lib/constants_test.cc index f4320f65c1f76d4d4c384110b39d6606773aaf01..180175b7495b32250af8ae77c8c7fba804703885 100644 --- a/tensorflow/compiler/xla/client/lib/constants_test.cc +++ b/tensorflow/compiler/xla/client/lib/constants_test.cc @@ -155,5 +155,12 @@ XLA_TEST_F(ConstantsTest, MaxValueF32) { {}); } +XLA_TEST_F(ConstantsTest, NanValueF32) { + XlaBuilder builder(TestName()); + NanValue(&builder, F32); + ComputeAndCompareR0(&builder, std::numeric_limits::quiet_NaN(), + {}); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/math.cc b/tensorflow/compiler/xla/client/lib/math.cc index 36fdda39b4124b9100c6054160f9c17bdf787d6f..19d98d100191fcba590d64c643d76b2bc5d5e5c5 100644 --- a/tensorflow/compiler/xla/client/lib/math.cc +++ b/tensorflow/compiler/xla/client/lib/math.cc @@ -13,14 +13,100 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +// This macro is required to make MSVC defines math constants in math.h +#define _USE_MATH_DEFINES +#include + #include "tensorflow/compiler/xla/client/lib/math.h" +#include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/lib/constants.h" +#include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" namespace xla { +// 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) { + auto& b = *operand.builder(); + TF_ASSIGN_OR_RETURN(auto shape, b.GetShape(operand)); + auto elem_ty = shape.element_type(); + if (!primitive_util::IsFloatingPointType(elem_ty)) { + return InvalidArgument( + "Operands to %s must be real-valued floating-point, but got %s", + op_name, PrimitiveType_Name(elem_ty)); + } + return Status::OK(); +} + +XlaOp IsPosInf(XlaOp operand) { + auto& b = *operand.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("IsPosInf", operand)); + TF_ASSIGN_OR_RETURN(auto shape, b.GetShape(operand)); + // Note that this is only correct for floating-point types. If we wanted it + // to be correct for all types, we'd need to Gt(MaxFiniteValue). + return Eq(operand, MaxValue(&b, shape.element_type())); + }); +} + +XlaOp IsNegInf(XlaOp operand) { + auto& b = *operand.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("IsNegInf", operand)); + TF_ASSIGN_OR_RETURN(auto shape, b.GetShape(operand)); + // Note that this is only correct for floating-point types. If we wanted it + // to be correct for all types, we'd need to Lt(MinFiniteValue). + return Eq(operand, MinValue(&b, shape.element_type())); + }); +} + +XlaOp IsInf(XlaOp operand) { + auto& b = *operand.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("IsInf", operand)); + return IsPosInf(Abs(operand)); + }); +} + +XlaOp IsNan(XlaOp operand) { + auto& b = *operand.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("IsNan", operand)); + return Ne(operand, operand); + }); +} + +XlaOp IsNegZero(XlaOp operand) { + auto& b = *operand.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("IsNegZero", operand)); + TF_ASSIGN_OR_RETURN(auto shape, b.GetShape(operand)); + + // The bitwise representation of -0 in bfloat16 and IEEE 754 is 0x80...0 + // (sign bit on, all other bits off). + switch (shape.element_type()) { + case F64: + return Eq(BitcastConvertType(operand, U64), + ConstantR0WithType(&b, U64, uint64{1} << 63)); + case F32: + return Eq(BitcastConvertType(operand, U32), + ConstantR0WithType(&b, U32, uint32{1} << 31)); + case F16: + case BF16: + // Not all XLA backends handle U16 well, so we convert to F32/U32. + // TODO(jlebar): It would be nice if we could stay in (B)F16/U16 for + // backends that *do* support it. + return Eq(BitcastConvertType(ConvertElementType(operand, F32), U32), + ConstantR0WithType(&b, U32, uint32{1} << 31)); + default: + LOG(FATAL) << "Expected real fp type."; + } + }); +} + XlaOp Sqrt(XlaOp operand) { return Pow(operand, ScalarLike(operand, 0.5)); } XlaOp Rsqrt(XlaOp operand) { return Pow(operand, ScalarLike(operand, -0.5)); } @@ -29,44 +115,6 @@ XlaOp Square(XlaOp operand) { return operand * operand; } XlaOp Reciprocal(XlaOp operand) { return ScalarLike(operand, 1.0) / operand; } -namespace { - -// Polynomials for computing erf/erfc. Originally from cephes. -// Note we use float for compatibility across devices, at the cost of some -// precision for 64 bit computations. -// -// Coefficients are in descending order. -std::array kErfcPCoefficient = { - 2.46196981473530512524E-10, 5.64189564831068821977E-1, - 7.46321056442269912687E0, 4.86371970985681366614E1, - 1.96520832956077098242E2, 5.26445194995477358631E2, - 9.34528527171957607540E2, 1.02755188689515710272E3, - 5.57535335369399327526E2}; -std::array kErfcQCoefficient = { - 1.00000000000000000000E0, 1.32281951154744992508E1, - 8.67072140885989742329E1, 3.54937778887819891062E2, - 9.75708501743205489753E2, 1.82390916687909736289E3, - 2.24633760818710981792E3, 1.65666309194161350182E3, - 5.57535340817727675546E2}; -std::array kErfcRCoefficient = { - 5.64189583547755073984E-1, 1.27536670759978104416E0, - 5.01905042251180477414E0, 6.16021097993053585195E0, - 7.40974269950448939160E0, 2.97886665372100240670E0}; -std::array kErfcSCoefficient = { - 1.00000000000000000000E0, 2.26052863220117276590E0, - 9.39603524938001434673E0, 1.20489539808096656605E1, - 1.70814450747565897222E1, 9.60896809063285878198E0, - 3.36907645100081516050E0}; -std::array kErfTCoefficient = { - 9.60497373987051638749E0, 9.00260197203842689217E1, - 2.23200534594684319226E3, 7.00332514112805075473E3, - 5.55923013010394962768E4}; -std::array kErfUCoefficient = { - 1.00000000000000000000E0, 3.35617141647503099647E1, - 5.21357949780152679795E2, 4.59432382970980127987E3, - 2.26290000613890934246E4, 4.92673942608635921086E4}; -} // namespace - // Evaluate the polynomial given coefficients and `x`. // N.B. Coefficients should be supplied in decreasing order. XlaOp EvaluatePolynomial(XlaOp x, absl::Span coefficients) { @@ -77,27 +125,86 @@ XlaOp EvaluatePolynomial(XlaOp x, absl::Span coefficients) { return poly; } -// Compute an approximation of the error function complement (1 - erf(x)). -XlaOp Erfc(XlaOp x) { +// Computes an approximation of the error function complement (1 - erf(x)). +// +// Precondition: abs(x) >= 1. Otherwise, use ErfImpl. +// +// This follows Cephes's f32 implementation of erfc, and so it may have errors +// for double precision. +// +// See also these alternate implementations of erf and erfc: +// +// https://stackoverflow.com/questions/35148198 +// https://stackoverflow.com/questions/35966695 +// +static XlaOp ErfcImpl(XlaOp x) { + // Coefficients for erfc(f32), from Cephes. + // + // erfc(x) = exp(-x^2) P(1/x), 1 < x < 2 + static std::array kErfcPCoefficient{ + +2.326819970068386E-2, -1.387039388740657E-1, +3.687424674597105E-1, + -5.824733027278666E-1, +6.210004621745983E-1, -4.944515323274145E-1, + +3.404879937665872E-1, -2.741127028184656E-1, +5.638259427386472E-1, + }; + // erfc(x) = exp(-x^2) 1/x P(1/x^2), 2 < x < 14 + static std::array kErfcRCoefficient{ + -1.047766399936249E+1, +1.297719955372516E+1, -7.495518717768503E+0, + +2.921019019210786E+0, -1.015265279202700E+0, +4.218463358204948E-1, + -2.820767439740514E-1, +5.641895067754075E-1, + }; + XlaOp abs_x = Abs(x); XlaOp z = Exp(-x * x); + XlaOp q = ScalarLike(x, 1) / abs_x; + XlaOp y = q * q; + XlaOp p = Select(Lt(abs_x, ScalarLike(x, 2.0)), + EvaluatePolynomial(y, kErfcPCoefficient), + EvaluatePolynomial(y, kErfcRCoefficient)); + y = z * q * p; + return Select(Lt(x, ScalarLike(x, 0)), ScalarLike(x, 2.0) - y, y); +} - XlaOp pp = EvaluatePolynomial(abs_x, kErfcPCoefficient); - XlaOp pq = EvaluatePolynomial(abs_x, kErfcQCoefficient); - XlaOp pr = EvaluatePolynomial(abs_x, kErfcRCoefficient); - XlaOp ps = EvaluatePolynomial(abs_x, kErfcSCoefficient); - - XlaOp y = Select(Lt(abs_x, ScalarLike(x, 8.0)), z * pp / pq, z * pr / ps); +// Compute a polynomial approximation of the error function. +// +// Precondition: abs(x) <= 1. Otherwise, use ErfcImpl. +// +// This follows Cephes's f32 implementation of erf, so it may have errors for +// double precision. +static XlaOp ErfImpl(XlaOp x) { + // Coefficients for by erf(f32), from Cephes. + // + // erf(x) = x P(x^2), 0 < x < 1 + static std::array kErfTCoefficient{ + +7.853861353153693E-5, -8.010193625184903E-4, +5.188327685732524E-3, + -2.685381193529856E-2, +1.128358514861418E-1, -3.761262582423300E-1, + +1.128379165726710E+0, + }; + + return x * EvaluatePolynomial(x * x, kErfTCoefficient); +} - return Select(Lt(x, ScalarLike(x, 0.0)), ScalarLike(x, 2.0) - y, y); +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)); + }); } -// Compute a polynomial approximation of the error function. XlaOp Erf(XlaOp x) { - XlaOp z = x * x; - XlaOp pt = EvaluatePolynomial(z, kErfTCoefficient); - XlaOp pu = EvaluatePolynomial(z, kErfUCoefficient); - return x * pt / pu; + auto& b = *x.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("Erf", 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)); + }); } // Approximation for the inverse error function from @@ -113,37 +220,30 @@ XlaOp Erf(XlaOp x) { // } // return p*x XlaOp ErfInv(XlaOp x) { - XlaBuilder* b = x.builder(); - return b->ReportErrorOrReturn([&]() -> StatusOr { - TF_ASSIGN_OR_RETURN(Shape shape, b->GetShape(x)); - constexpr int kDegree = 9; - constexpr std::array w_less_than_5_constants = { - 2.81022636e-08f, 3.43273939e-07f, -3.5233877e-06f, - -4.39150654e-06f, 0.00021858087f, -0.00125372503f, - -0.00417768164f, 0.246640727f, 1.50140941f}; - constexpr std::array w_greater_than_5_constants = { - -0.000200214257f, 0.000100950558f, 0.00134934322f, - -0.00367342844f, 0.00573950773f, -0.0076224613f, - 0.00943887047f, 1.00167406f, 2.83297682f}; + constexpr int kDegree = 9; + constexpr std::array w_less_than_5_constants = { + 2.81022636e-08f, 3.43273939e-07f, -3.5233877e-06f, + -4.39150654e-06f, 0.00021858087f, -0.00125372503f, + -0.00417768164f, 0.246640727f, 1.50140941f}; + constexpr std::array w_greater_than_5_constants = { + -0.000200214257f, 0.000100950558f, 0.00134934322f, + -0.00367342844f, 0.00573950773f, -0.0076224613f, + 0.00943887047f, 1.00167406f, 2.83297682f}; - auto one = ScalarLike(x, 1.0); - auto w = -Log((one - x) * (one + x)); - - auto lt = Lt(w, ScalarLike(x, 5.0)); - auto coefficient = [&](int i) { - return Select(lt, - Broadcast(ScalarLike(x, w_less_than_5_constants[i]), - AsInt64Slice(shape.dimensions())), - Broadcast(ScalarLike(x, w_greater_than_5_constants[i]), - AsInt64Slice(shape.dimensions()))); - }; - w = Select(lt, w - ScalarLike(x, 2.5), Sqrt(w) - ScalarLike(x, 3.0)); - auto p = coefficient(0); - for (int i = 1; i < kDegree; ++i) { - p = coefficient(i) + p * w; - } - return p * x; - }); + auto one = ScalarLike(x, 1.0); + auto w = -Log((one - x) * (one + x)); + + auto lt = Lt(w, ScalarLike(x, 5.0)); + auto coefficient = [&](int i) { + return Select(lt, FullLike(x, w_less_than_5_constants[i]), + FullLike(x, w_greater_than_5_constants[i])); + }; + w = Select(lt, w - ScalarLike(x, 2.5), Sqrt(w) - ScalarLike(x, 3.0)); + auto p = coefficient(0); + for (int i = 1; i < kDegree; ++i) { + p = coefficient(i) + p * w; + } + return p * x; } namespace { @@ -170,49 +270,86 @@ 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) { - XlaOp one_half = ScalarLike(input, 0.5); - XlaOp one = ScalarLike(input, 1); - - XlaOp pi = ScalarLike(input, M_PI); - XlaOp log_pi = ScalarLike(input, std::log(M_PI)); - XlaOp log_sqrt_two_pi = ScalarLike(input, (std::log(2) + std::log(M_PI)) / 2); - - XlaOp lanczos_gamma_plus_one_half = ScalarLike(input, kLanczosGamma + 0.5); - XlaOp log_lanczos_gamma_plus_one_half = - ScalarLike(input, std::log(kLanczosGamma + 0.5)); - - XlaOp base_lanczos_coeff = ScalarLike(input, kBaseLanczosCoeff); - - // If the input is less than 0.5 use Gauss's reflection formula: - // gamma(x) = pi / sin(pi * x) * gamma(1 - x) - XlaOp need_to_reflect = Lt(Real(input), one_half); - XlaOp z = Select(need_to_reflect, -input, input - one); - - XlaOp x = base_lanczos_coeff; - for (int i = 0; i < kLanczosCoefficients.size(); ++i) { - XlaOp lanczos_coefficient = ScalarLike(input, kLanczosCoefficients[i]); - XlaOp index = ScalarLike(input, i); - x = x + lanczos_coefficient / (z + index + one); - } + auto& b = *input.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("Lgamma", input)); + + XlaOp one_half = ScalarLike(input, 0.5); + XlaOp one = ScalarLike(input, 1); + + XlaOp pi = ScalarLike(input, M_PI); + XlaOp log_pi = ScalarLike(input, std::log(M_PI)); + XlaOp log_sqrt_two_pi = + ScalarLike(input, (std::log(2) + std::log(M_PI)) / 2); + + XlaOp lanczos_gamma_plus_one_half = ScalarLike(input, kLanczosGamma + 0.5); + XlaOp log_lanczos_gamma_plus_one_half = + ScalarLike(input, std::log(kLanczosGamma + 0.5)); + + XlaOp base_lanczos_coeff = ScalarLike(input, kBaseLanczosCoeff); + + // If the input is less than 0.5 use Euler's reflection formula: + // gamma(x) = pi / (sin(pi * x) * gamma(1 - x)) + XlaOp need_to_reflect = Lt(input, one_half); + XlaOp z = Select(need_to_reflect, -input, input - one); + + XlaOp x = base_lanczos_coeff; + for (int i = 0; i < kLanczosCoefficients.size(); ++i) { + XlaOp lanczos_coefficient = ScalarLike(input, kLanczosCoefficients[i]); + XlaOp index = ScalarLike(input, i); + x = x + lanczos_coefficient / (z + index + one); + } - // To improve accuracy on platforms with less-precise log implementations, - // compute log(lanczos_gamma_plus_one_half) at compile time and use log1p on - // the device. - // log(t) = log(kLanczosGamma + 0.5 + z) - // = log(kLanczosGamma + 0.5) + log1p(z / (kLanczosGamma + 0.5)) - XlaOp t = lanczos_gamma_plus_one_half + z; - 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); - - // If z = a + 0j, the analytic continuation of log reduces to taking the - // absolute value of the real part. - // Re(log(z)) = Re(log|z| + arg(z)j) - // = log|a| - XlaOp reflection = log_pi - Log(Abs(Sin(pi * input))) - log_y; - XlaOp result = Select(need_to_reflect, reflection, log_y); - return result; + // To improve accuracy on platforms with less-precise log implementations, + // compute log(lanczos_gamma_plus_one_half) at compile time and use log1p on + // the device. + // log(t) = log(kLanczosGamma + 0.5 + z) + // = log(kLanczosGamma + 0.5) + log1p(z / (kLanczosGamma + 0.5)) + XlaOp t = lanczos_gamma_plus_one_half + z; + 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 reflected value, used when x < 0.5: + // + // lgamma(x) = log(pi) - lgamma(1-x) - log(abs(sin(pi * x))). + // + // (The abs is because lgamma is the log of the absolute value of the gamma + // function.) + // + // We have to be careful when computing the final term above. gamma(x) goes + // to +/-inf at every integer x < 0, and this is controlled by the + // sin(pi * x) term. The slope is large, so precision is particularly + // 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. + // + // 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)). + // + XlaOp abs_input = Abs(input); + XlaOp reflection_denom = Log(Sin(pi * (abs_input - Floor(abs_input)))); + + // Avoid computing -inf - inf, which is nan. If reflection_denom is +/-inf, + // then it "wins" and the result is +/-inf. + XlaOp reflection = + Select(IsFinite(reflection_denom), log_pi - reflection_denom - log_y, + -reflection_denom); + XlaOp result = Select(need_to_reflect, reflection, log_y); + + // 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); + }); } // Compute the Digamma function using Lanczos' approximation from "A Precision @@ -223,69 +360,84 @@ 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) { - XlaOp zero = ScalarLike(input, 0); - XlaOp one_half = ScalarLike(input, 0.5); - XlaOp one = ScalarLike(input, 1); - - XlaOp pi = ScalarLike(input, M_PI); - - XlaOp lanczos_gamma = ScalarLike(input, kLanczosGamma); - XlaOp lanczos_gamma_plus_one_half = ScalarLike(input, kLanczosGamma + 0.5); - XlaOp log_lanczos_gamma_plus_one_half = - ScalarLike(input, std::log(kLanczosGamma + 0.5)); - - XlaOp base_lanczos_coeff = ScalarLike(input, kBaseLanczosCoeff); - - // If the input is less than 0.5 use Gauss's reflection formula: - // digamma(x) = digamma(1 - x) - pi * cot(pi * x) - XlaOp need_to_reflect = Lt(Real(input), one_half); - XlaOp z = Select(need_to_reflect, -input, input - one); - - XlaOp num = zero; - XlaOp denom = base_lanczos_coeff; - for (int i = 0; i < kLanczosCoefficients.size(); ++i) { - XlaOp lanczos_coefficient = ScalarLike(input, kLanczosCoefficients[i]); - XlaOp index = ScalarLike(input, i); - num = num - lanczos_coefficient / ((z + index + one) * (z + index + one)); - denom = denom + lanczos_coefficient / (z + index + one); - } + auto& b = *input.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("Digamma", input)); + + XlaOp zero = ScalarLike(input, 0); + XlaOp one_half = ScalarLike(input, 0.5); + XlaOp one = ScalarLike(input, 1); + + XlaOp pi = ScalarLike(input, M_PI); + + XlaOp lanczos_gamma = ScalarLike(input, kLanczosGamma); + XlaOp lanczos_gamma_plus_one_half = ScalarLike(input, kLanczosGamma + 0.5); + XlaOp log_lanczos_gamma_plus_one_half = + ScalarLike(input, std::log(kLanczosGamma + 0.5)); + + XlaOp base_lanczos_coeff = ScalarLike(input, kBaseLanczosCoeff); + + // If the input is less than 0.5 use Euler's reflection formula: + // digamma(x) = digamma(1 - x) - pi * cot(pi * x) + XlaOp need_to_reflect = Lt(input, one_half); + XlaOp z = Select(need_to_reflect, -input, input - one); + + XlaOp num = zero; + XlaOp denom = base_lanczos_coeff; + for (int i = 0; i < kLanczosCoefficients.size(); ++i) { + XlaOp lanczos_coefficient = ScalarLike(input, kLanczosCoefficients[i]); + XlaOp index = ScalarLike(input, i); + num = num - lanczos_coefficient / ((z + index + one) * (z + index + one)); + denom = denom + lanczos_coefficient / (z + index + one); + } - // To improve accuracy on platforms with less-precise log implementations, - // compute log(lanczos_gamma_plus_one_half) at compile time and use log1p on - // the device. - // log(t) = log(kLanczosGamma + 0.5 + z) - // = log(kLanczosGamma + 0.5) + log1p(z / (kLanczosGamma + 0.5)) - XlaOp t = lanczos_gamma_plus_one_half + z; - XlaOp log_t = - log_lanczos_gamma_plus_one_half + 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); - XlaOp result = Select(need_to_reflect, reflection, y); - return result; + // To improve accuracy on platforms with less-precise log implementations, + // compute log(lanczos_gamma_plus_one_half) at compile time and use log1p on + // the device. + // log(t) = log(kLanczosGamma + 0.5 + z) + // = log(kLanczosGamma + 0.5) + log1p(z / (kLanczosGamma + 0.5)) + XlaOp t = lanczos_gamma_plus_one_half + z; + XlaOp log_t = log_lanczos_gamma_plus_one_half + + 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); + }); } // Implements Banker's rounding: numbers that are equidistant between two // integers are rounded towards even. XlaOp RoundToEven(XlaOp x) { - auto half = ScalarLike(x, 0.5); - auto one = ScalarLike(x, 1.0); - auto two = ScalarLike(x, 2.0); - - auto round_val = Floor(x); - auto fraction = x - round_val; - auto nearest_even_int = round_val - two * Floor(half * x); - auto is_odd = Eq(nearest_even_int, one); - return Select(Or(Gt(fraction, half), And(Eq(fraction, half), is_odd)), - round_val + one, round_val); + auto& b = *x.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + // Reject non-real non-fp inputs (What does it even mean to round a complex + // number? Do you round each component equally? In that case, you should + // just ask for that explicitly.) + TF_RETURN_IF_ERROR(EnsureOperandIsRealFp("RoundToEven", x)); + + auto half = ScalarLike(x, 0.5); + auto one = ScalarLike(x, 1.0); + auto two = ScalarLike(x, 2.0); + + auto round_val = Floor(x); + auto fraction = x - round_val; + auto nearest_even_int = round_val - two * Floor(half * x); + auto is_odd = Eq(nearest_even_int, one); + return Select(Or(Gt(fraction, half), And(Eq(fraction, half), is_odd)), + round_val + one, round_val); + }); } // Trigonometric functions. -// acos(x) = 2 * atan(sqrt(1 - x^2) / (1 + x)) +// acos(x) = 2 * atan(sqrt(1 - x^2) / (1 + x)) if x != -1 +// pi if x == -1 XlaOp Acos(XlaOp x) { - return ScalarLike(x, 2.0) * - Atan2(Sqrt(ScalarLike(x, 1.0) - x * x), ScalarLike(x, 1.0) + x); + return Select(Ne(x, FullLike(x, -1)), + ScalarLike(x, 2.0) * Atan2(Sqrt(ScalarLike(x, 1.0) - x * x), + ScalarLike(x, 1.0) + x), + FullLike(x, M_PI)); } // asin(x) = 2 * atan(x / (1 + sqrt(1 - x^2))) @@ -323,9 +475,88 @@ XlaOp MaybeConjugate(XlaOp x, bool conjugate) { XlaBuilder* builder = x.builder(); return builder->ReportErrorOrReturn([&]() -> StatusOr { TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); - auto perform_conj = shape.element_type() == C64 && conjugate; + auto perform_conj = + primitive_util::IsComplexType(shape.element_type()) && conjugate; return perform_conj ? Conj(x) : x; }); } +XlaOp NextAfter(XlaOp from, XlaOp to) { + auto builder = from.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(auto shape, builder->GetShape(from)); + int bitwidth = primitive_util::BitWidth(shape.element_type()); + auto int_type = primitive_util::UnsignedIntegralTypeForBitWidth(bitwidth); + auto from_as_int = BitcastConvertType(from, int_type); + auto to_as_int = BitcastConvertType(to, int_type); + + // The result is NaN if either "from" or "to" are NaN. + auto from_is_nan = Ne(from, from); + auto to_is_nan = Ne(to, to); + auto nan_input = Or(from_is_nan, to_is_nan); + auto result_for_nan = + Broadcast(ScalarLike(from, std::numeric_limits::quiet_NaN()), + shape.dimensions()); + result_for_nan = BitcastConvertType(result_for_nan, int_type); + + // The sign bit is the MSB. + const int64 sign_mask = int64{1} << (bitwidth - 1); + // Discard the sign bit to make the result non-negative. + auto from_abs = And(from_as_int, ScalarLike(from_as_int, ~sign_mask)); + auto to_abs = And(to_as_int, ScalarLike(to_as_int, ~sign_mask)); + + // When both "from" and "to" are equal, the result is "to". + // N.B. It would not make a difference if we chose the result to be "from". + auto from_and_to_are_equal = Eq(from_as_int, to_as_int); + auto result_for_equal = to_as_int; + + // When both "from" and "to" are both 0, the result is "to". This ensures we + // get a zero signed like "to". + auto from_is_zero = Eq(from_abs, ZerosLike(from_abs)); + auto to_is_zero = Eq(to_abs, ZerosLike(to_abs)); + auto result_for_both_zero = to_as_int; + + auto from_sign = And(from_as_int, ScalarLike(from_as_int, sign_mask)); + auto to_sign = And(to_as_int, ScalarLike(to_as_int, sign_mask)); + + // If from == 0 && to != 0, we need to return the smallest subnormal number + // signed like "to". + auto result_for_from_zero_to_non_zero = + Or(to_sign, ScalarLike(from_as_int, 1)); + + // If the sign of "from" and "to" disagree: + // - we need to make the magnitude of "from" smaller so that it is closer to + // zero. + // + // Otherwise the signs agree: + // - "from" with a magnitude larger than "to" means we need to make the + // magnitude smaller. + // - "from" with a magnitude smaller than "to" means we need to make the + // magnitude larger. + // - "from" with the same magnitude and sign as "to" has already been + // handled. + auto signs_disagree = Ne(from_sign, to_sign); + auto from_magnitude_larger_than_to = Gt(from_abs, to_abs); + auto result_has_smaller_magnitude = + Or(from_magnitude_larger_than_to, signs_disagree); + auto magnitude_adjustment = + Select(result_has_smaller_magnitude, + Broadcast(ScalarLike(from_as_int, -1), shape.dimensions()), + Broadcast(ScalarLike(from_as_int, 1), shape.dimensions())); + auto result = Add(from_as_int, magnitude_adjustment); + // Handle from == ±0. + result = Select(from_is_zero, + Select(to_is_zero, result_for_both_zero, + result_for_from_zero_to_non_zero), + result); + // Handle from == to. + result = Select(from_and_to_are_equal, result_for_equal, result); + // Handle isnan(from) || isnan(to). + result = Select(nan_input, result_for_nan, result); + + // Cast back to the original type. + return BitcastConvertType(result, shape.element_type()); + }); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/math.h b/tensorflow/compiler/xla/client/lib/math.h index 17612bf9fdc0f1eabb338671c93c025c5b268872..b036fa299d92988439dfecbb5415865071d5577d 100644 --- a/tensorflow/compiler/xla/client/lib/math.h +++ b/tensorflow/compiler/xla/client/lib/math.h @@ -20,6 +20,23 @@ limitations under the License. namespace xla { +// Determines whether operand is +/-inf or nan. +// +// Raises an error if called on integral or complex values. +XlaOp IsPosInf(XlaOp operand); +XlaOp IsNegInf(XlaOp operand); +XlaOp IsInf(XlaOp operand); +XlaOp IsNan(XlaOp operand); + +// Determines whether operand is equal to -0. +// +// Raises an error for integral or complex values. +XlaOp IsNegZero(XlaOp operand); + +// Returns the next number after 'from' in the direction of 'to' the same way +// std::nextafter(from, to) would. +XlaOp NextAfter(XlaOp from, XlaOp to); + // Computes the square root of 'operand'. XlaOp Sqrt(XlaOp operand); @@ -32,7 +49,7 @@ XlaOp Square(XlaOp operand); // Computes the reciprocal of 'operand'. XlaOp Reciprocal(XlaOp operand); -// Evaluates a polynomial given coefficients and `x`. +// Evaluates a polynomial given coefficients and 'x'. // N.B. Coefficients should be supplied in decreasing order. XlaOp EvaluatePolynomial(XlaOp x, absl::Span coefficients); @@ -86,7 +103,7 @@ XlaOp Cosh(XlaOp x); // Computes the hyperbolic sine of 'x'. XlaOp Sinh(XlaOp x); -// Applies a complex conjugation operation if `a` is complex and `conjugate` +// Applies a complex conjugation operation if 'a' is complex and 'conjugate' // is true, otherwise returns its argument. xla::XlaOp MaybeConjugate(xla::XlaOp x, bool conjugate); diff --git a/tensorflow/compiler/xla/client/lib/math_exhaustive_test.cc b/tensorflow/compiler/xla/client/lib/math_exhaustive_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..f4ed0db56f7026d2c397c5beb1cc7ea3e4b06fee --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/math_exhaustive_test.cc @@ -0,0 +1,187 @@ +/* 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 ae2ea225d1aadd7b3a794eabeca866c498f34760..bdfb0575f573716b54cf9116d155d8a3a55056e8 100644 --- a/tensorflow/compiler/xla/client/lib/math_test.cc +++ b/tensorflow/compiler/xla/client/lib/math_test.cc @@ -16,6 +16,7 @@ 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/primitive_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" @@ -30,6 +31,138 @@ class MathTest : public ClientLibraryTestBase { ErrorSpec error_spec_{0.0001}; }; +// Write TYPED_TESTs within the class definition so that we don't have to litter +// "this->" everywhere. +template +class MathTypedTest : public MathTest { + public: + void TestLogEdgeCases() { + SetFastMathDisabled(true); + + XlaBuilder b(TestName()); + Log(AddParam(LiteralUtil::CreateR1({T{0.0}, T{-0.0}}), &b)); + ComputeAndCompareR1(&b, + {-std::numeric_limits::infinity(), + -std::numeric_limits::infinity()}, + {}, error_spec_); + } + + void TestLog1pEdgeCases() { + SetFastMathDisabled(true); + + XlaBuilder b(TestName()); + Log1p(AddParam(LiteralUtil::CreateR1({T{0.0}, T{-0.0}, T{-1.0}}), &b)); + ComputeAndCompareR1( + &b, {T{0.0}, T{-0.0}, -std::numeric_limits::infinity()}, {}, + error_spec_); + } + + void TestIsInfOrNan() { + SetFastMathDisabled(true); + + XlaBuilder b(TestName()); + auto x = + ConstantR1(&b, { + T{0}, + T{100}, + T{-1000}, + T{std::numeric_limits::max()}, + T{std::numeric_limits::lowest()}, + T{std::numeric_limits::infinity()}, + T{-std::numeric_limits::infinity()}, + T{std::numeric_limits::quiet_NaN()}, + T{std::numeric_limits::signaling_NaN()}, + }); + Tuple(&b, {IsFinite(x), IsInf(x), IsPosInf(x), IsNegInf(x), IsNan(x)}); + + auto expected = LiteralUtil::MakeTupleOwned( + LiteralUtil::CreateR1( + {true, true, true, true, true, false, false, false, false}), + LiteralUtil::CreateR1( + {false, false, false, false, false, true, true, false, false}), + LiteralUtil::CreateR1( + {false, false, false, false, false, true, false, false, false}), + LiteralUtil::CreateR1( + {false, false, false, false, false, false, true, false, false}), + LiteralUtil::CreateR1( + {false, false, false, false, false, false, false, true, true})); + ComputeAndCompareLiteral(&b, expected, {}); + } + + void TestIsNegZero() { + SetFastMathDisabled(true); + XlaBuilder b(TestName()); + T inf(std::numeric_limits::infinity()); + T nan(std::numeric_limits::quiet_NaN()); + IsNegZero(AddParam( + LiteralUtil::CreateR1({T{-0.0}, T{0}, T{1}, T{-1}, inf, -inf, nan}), + &b)); + + ComputeAndCompareLiteral( + &b, + LiteralUtil::CreateR1( + {true, false, false, false, false, false, false}), + {}, error_spec_); + } +}; + +// TODO(b/123355973): Add bfloat16 to TestTypes once it's working. +#ifdef XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16 +using TestTypes = ::testing::Types; +#else +using TestTypes = ::testing::Types; +#endif + +TYPED_TEST_CASE(MathTypedTest, TestTypes); + +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(); } + +// Check that certain ops only support real, floating-point inputs. +// +// TODO(jlebar): Expand this test to cover more ops. +XLA_TEST_F(MathTest, RealFpOnlyOps) { + for (int64 i = PrimitiveType_MIN; i <= PrimitiveType_MAX; ++i) { + auto ty = static_cast(i); + SCOPED_TRACE(PrimitiveType_Name(ty)); + Shape shape; + if (primitive_util::IsArrayType(ty)) { + shape = ShapeUtil::MakeShape(ty, {42}); + } else if (ty == PrimitiveType::TUPLE) { + shape = ShapeUtil::MakeTupleShape({}); + } else if (ty == PrimitiveType::OPAQUE) { + shape = ShapeUtil::MakeOpaqueShape(); + } else if (ty == PrimitiveType::TOKEN) { + shape = ShapeUtil::MakeTokenShape(); + } else { + continue; + } + + for (const auto& test : + std::vector, string>>({ + {IsFinite, "is_finite"}, + {IsInf, "is_inf"}, + {IsPosInf, "is_pos_inf"}, + {IsNegInf, "is_neg_inf"}, + {IsNan, "is_nan"}, + {Erf, "erf"}, + {Erfc, "erfc"}, + {Lgamma, "lgamma"}, + {Digamma, "digamma"}, + {RoundToEven, "round_to_even"}, + })) { + SCOPED_TRACE(test.second); + XlaBuilder b(TestName()); + XlaOp p = Parameter(&b, 0, shape, "p0"); + test.first(p); + + EXPECT_EQ(b.first_error().ok(), primitive_util::IsFloatingPointType(ty)); + } + } +} + XLA_TEST_F(MathTest, SqrtF32) { XlaBuilder builder(TestName()); Literal zero_literal = LiteralUtil::Zero(PrimitiveType::F32); @@ -106,6 +239,27 @@ XLA_TEST_F(MathTest, Lgamma) { ComputeAndCompareR1(&builder, expected, {}, error_spec_); } +XLA_TEST_F(MathTest, LgammaF16) { + SetFastMathDisabled(true); + + XlaBuilder b(TestName()); + + // These seemingly arbitrary inputs came from debugging the lgamma + // implementation against a test which tried all possible f16 values. + auto x = ConstantR1(&b, { + half(-7360.0), + half(-4066.0), + half(-5.9605e-08), + }); + Lgamma(x); + std::vector expected = { + std::numeric_limits::infinity(), + std::numeric_limits::infinity(), + half(16.64), + }; + ComputeAndCompareR1(&b, expected, {}, ErrorSpec{0.1}); +} + XLA_TEST_F(MathTest, Digamma) { XlaBuilder builder(TestName()); auto x = ConstantR1(&builder, {1.0, 0.5, 1 / 3.0, 0.25, 1 / 6.0, 0.125, @@ -148,5 +302,40 @@ XLA_TEST_F(MathTest, RoundToEven) { ComputeAndCompareR1(&builder, expected, {}, error_spec_); } +XLA_TEST_F(MathTest, ErfRejectsComplexInputs) { + XlaBuilder b(TestName()); + auto x = ConstantR1>(&b, {{0, 0}}); + Erf(x); + EXPECT_FALSE(b.Build().status().ok()); +} + +XLA_TEST_F(MathTest, ErfcRejectsComplexInputs) { + XlaBuilder b(TestName()); + auto x = ConstantR1>(&b, {{0, 0}}); + Erfc(x); + EXPECT_FALSE(b.Build().status().ok()); +} + +XLA_TEST_F(MathTest, LgammaRejectsComplexInputs) { + XlaBuilder b(TestName()); + auto x = ConstantR1>(&b, {{0, 0}}); + Lgamma(x); + EXPECT_FALSE(b.Build().status().ok()); +} + +XLA_TEST_F(MathTest, DigammaRejectsComplexInputs) { + XlaBuilder b(TestName()); + auto x = ConstantR1>(&b, {{0, 0}}); + Digamma(x); + EXPECT_FALSE(b.Build().status().ok()); +} + +XLA_TEST_F(MathTest, RoundToEvenRejectsComplexInputs) { + XlaBuilder b(TestName()); + auto x = ConstantR1>(&b, {{0, 0}}); + RoundToEven(x); + EXPECT_FALSE(b.Build().status().ok()); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/matrix.cc b/tensorflow/compiler/xla/client/lib/matrix.cc index dcec2139e47fc86d81a8877b4dccc43eb2b7207f..a055a8e625c680cf5232896c95cd35b78cb172bc 100644 --- a/tensorflow/compiler/xla/client/lib/matrix.cc +++ b/tensorflow/compiler/xla/client/lib/matrix.cc @@ -15,17 +15,26 @@ limitations under the License. #include "tensorflow/compiler/xla/client/lib/matrix.h" +#include #include #include +#include "absl/algorithm/container.h" +#include "absl/container/flat_hash_set.h" +#include "absl/strings/ascii.h" +#include "absl/strings/str_split.h" +#include "absl/strings/string_view.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/lib/constants.h" +#include "tensorflow/compiler/xla/client/lib/slicing.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" namespace xla { @@ -37,7 +46,7 @@ XlaOp IdentityMatrix(XlaBuilder* builder, PrimitiveType type, int64 m, return ConvertElementType(indicator, type); } -XlaOp GetMatrixDiagonal(XlaOp x) { +XlaOp GetMatrixDiagonal(XlaOp x, int k) { XlaBuilder* builder = x.builder(); return builder->ReportErrorOrReturn([&]() -> StatusOr { TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); @@ -45,10 +54,13 @@ XlaOp GetMatrixDiagonal(XlaOp x) { TF_RET_CHECK(n_dims >= 2); const int64 m = shape.dimensions(n_dims - 2); const int64 n = shape.dimensions(n_dims - 1); + + auto offset = ConstantR0WithType(builder, S32, k); + absl::Span major_dims = AsInt64Slice(shape.dimensions()).subspan(/*pos=*/0, /*len=*/n_dims - 2); - auto a = Iota(builder, U32, n); - auto b = Iota(builder, U32, m); + auto a = Iota(builder, S32, n); + auto b = Iota(builder, S32, m) + offset; auto indicator = Eq(b, Broadcast(a, {m}), /*broadcast_dimensions=*/{0}); auto mask = Broadcast(indicator, major_dims); @@ -58,9 +70,21 @@ XlaOp GetMatrixDiagonal(XlaOp x) { primitive_util::IsIntegralType(shape.element_type()) ? CreateScalarOrComputation(shape.element_type(), builder) : CreateScalarAddComputation(shape.element_type(), builder); - - return Reduce(Select(mask, x, Zeros(builder, shape)), ScalarLike(x, 0), - reducer, {m >= n ? n_dims - 2 : n_dims - 1}); + // k == 0, we can save one slice op. + if (k == 0) { + return Reduce(Select(mask, x, Zeros(builder, shape)), ScalarLike(x, 0), + reducer, {m >= n ? n_dims - 2 : n_dims - 1}); + } else if (k > 0) { + auto result = Reduce(Select(mask, x, Zeros(builder, shape)), + ScalarLike(x, 0), reducer, {n_dims - 2}); + return SliceInMinorDims(result, {std::min(k, n)}, + {std::min(m + k, n)}); + } else { + auto result = Reduce(Select(mask, x, Zeros(builder, shape)), + ScalarLike(x, 0), reducer, {n_dims - 1}); + return SliceInMinorDims(result, {std::min(-k, m)}, + {std::min(m, n - k)}); + } }); } @@ -91,77 +115,224 @@ XlaOp UpperTriangle(XlaOp x) { return Triangle(x, false); } XlaOp LowerTriangle(XlaOp x) { return Triangle(x, true); } -XlaOp BatchDot(XlaOp x, XlaOp y, PrecisionConfig::Precision precision) { +Status ValidateEinsumNumericDimensions(absl::Span x_config, + absl::Span y_config, + absl::Span output_config) { + for (auto dim : output_config) { + if (absl::c_linear_search(x_config, dim) || + absl::c_linear_search(y_config, dim)) { + if (absl::c_count(output_config, dim) > 1) { + return InvalidArgument("Einsum has repeated output dimension."); + } + continue; + } + return InvalidArgument( + "Einsum has output dimension without corresponding input dimension."); + } + for (auto dim : x_config) { + if (absl::c_linear_search(y_config, dim) || + absl::c_linear_search(output_config, dim)) { + if (absl::c_count(x_config, dim) > 1) { + return InvalidArgument("Einsum has repeated lhs dimension."); + } + continue; + } + return InvalidArgument( + "Einsum has lhs dimension without corresponding rhs or output " + "dimension."); + } + for (auto dim : y_config) { + if (absl::c_linear_search(x_config, dim) || + absl::c_linear_search(output_config, dim)) { + if (absl::c_count(y_config, dim) > 1) { + return InvalidArgument("Einsum has repeated rhs dimension."); + } + continue; + } + return InvalidArgument( + "Einsum has rhs dimension without corresponding lhs or output " + "dimension."); + } + return Status::OK(); +} + +xla::XlaOp Einsum(xla::XlaOp x, absl::Span x_config, xla::XlaOp y, + absl::Span y_config, + absl::Span output_config, + xla::PrecisionConfig::Precision precision) { XlaBuilder* builder = x.builder(); return builder->ReportErrorOrReturn([&]() -> StatusOr { - TF_ASSIGN_OR_RETURN(Shape x_shape, builder->GetShape(x)); - TF_ASSIGN_OR_RETURN(Shape y_shape, builder->GetShape(y)); + TF_RETURN_IF_ERROR( + ValidateEinsumNumericDimensions(x_config, y_config, output_config)); + const int64 x_rank = x_config.size(); + const int64 y_rank = y_config.size(); + const int64 output_rank = output_config.size(); + absl::flat_hash_set x_map; + absl::flat_hash_set y_map; + absl::flat_hash_set output_map; + + auto find = [&](const absl::flat_hash_set& map, int64 d) { + return map.count(d) != 0; + }; - // Check that both tensors have the same number of dimensions. There must be - // at least two (the batch dimensions can be empty). - if (x_shape.rank() != y_shape.rank()) { - return InvalidArgument( - "Arguments to BatchDot have different ranks: %s vs. %s", - ShapeUtil::HumanString(x_shape), ShapeUtil::HumanString(y_shape)); + auto insert = [&](absl::flat_hash_set& map, char d) { + CHECK(!find(map, d)); + map.insert(d); + }; + + for (auto d : x_config) { + insert(x_map, d); } - const int ndims = x_shape.rank(); - if (ndims < 2) { - return InvalidArgument( - "Arguments to BatchDot must have rank >= 2: got %d", ndims); + + for (auto d : y_config) { + insert(y_map, d); } - // The batch dimensions must be equal and the matrix dimensions must be - // valid. - std::vector batch_dimension_numbers; - for (int i = 0; i < ndims - 2; ++i) { - if (x_shape.dimensions(i) != y_shape.dimensions(i)) { - return InvalidArgument( - "Dimension %d of inputs to BatchDot must be equal: shapes %s vs %s", - i, ShapeUtil::HumanString(x_shape), - ShapeUtil::HumanString(y_shape)); - } - batch_dimension_numbers.push_back(i); + for (auto d : output_config) { + insert(output_map, d); } - int x_inner_dim = ndims - 1; - int y_inner_dim = ndims - 2; - if (x_shape.dimensions(x_inner_dim) != y_shape.dimensions(y_inner_dim)) { - return InvalidArgument( - "Dimensions %d and %d of arguments to BatchDot must be equal: " - "shapes %s vs %s", - x_inner_dim, y_inner_dim, ShapeUtil::HumanString(x_shape), - ShapeUtil::HumanString(y_shape)); + DotDimensionNumbers dnums; + std::vector lhs_outer_dims; + auto is_batch_dim = [&](int64 d) { + return find(x_map, d) && find(y_map, d) && find(output_map, d); + }; + auto is_contracting = [&](int64 d) { + return find(x_map, d) && find(y_map, d); + }; + auto rhs_dimension_number = [&](int64 d) { + return absl::c_find(y_config, d) - y_config.begin(); + }; + for (int64 i = 0; i < x_rank; ++i) { + auto dim_name = x_config[i]; + if (is_batch_dim(dim_name)) { + dnums.add_lhs_batch_dimensions(i); + dnums.add_rhs_batch_dimensions(rhs_dimension_number(dim_name)); + } else if (is_contracting(dim_name)) { + dnums.add_lhs_contracting_dimensions(i); + dnums.add_rhs_contracting_dimensions(rhs_dimension_number(dim_name)); + } else { + lhs_outer_dims.push_back(i); + } } - // Check for zero lhs/rhs dim size. - if (ShapeUtil::IsZeroElementArray(x_shape) || - ShapeUtil::IsZeroElementArray(y_shape)) { - std::vector dimensions(batch_dimension_numbers.size()); - for (int i = 0; i < batch_dimension_numbers.size(); ++i) { - dimensions[i] = x_shape.dimensions(batch_dimension_numbers[i]); + std::vector rhs_outer_dims; + for (int64 i = 0; i < y_rank; ++i) { + auto dim_name = y_config[i]; + if (!is_batch_dim(dim_name) && !is_contracting(dim_name)) { + rhs_outer_dims.push_back(i); } - int x_outer_dim = ndims - 2; - int y_outer_dim = ndims - 1; - dimensions.push_back(x_shape.dimensions(x_outer_dim)); - dimensions.push_back(y_shape.dimensions(y_outer_dim)); - return Broadcast( - ConstantLiteral(builder, LiteralUtil::Zero(x_shape.element_type())), - dimensions); + } + + auto output_dimension_number = [&](char d) { + return absl::c_find(output_config, d) - output_config.begin(); + }; + + std::vector output_dims; + output_dims.reserve(output_rank); + for (auto d : dnums.lhs_batch_dimensions()) { + output_dims.push_back(output_dimension_number(x_config[d])); + } + for (auto d : lhs_outer_dims) { + output_dims.push_back(output_dimension_number(x_config[d])); + } + for (auto d : rhs_outer_dims) { + output_dims.push_back(output_dimension_number(y_config[d])); + } + + std::vector transpose_dims(output_rank); + for (int64 i = 0; i < output_rank; ++i) { + transpose_dims[output_dims[i]] = i; } PrecisionConfig precision_proto; precision_proto.add_operand_precision(precision); precision_proto.add_operand_precision(precision); + return Transpose(DotGeneral(x, y, dnums, &precision_proto), transpose_dims); + }); +} + +XlaOp BatchDot(XlaOp x, XlaOp y, PrecisionConfig::Precision precision) { + XlaBuilder* builder = x.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape x_shape, builder->GetShape(x)); + TF_ASSIGN_OR_RETURN(Shape y_shape, builder->GetShape(y)); + + // The batch dimensions must be equal and the matrix dimensions must be + // valid. + std::vector batch_dimension_numbers; + const int ndims = x_shape.rank(); + batch_dimension_numbers.reserve(ndims - 2); + for (int i = 0; i < ndims - 2; ++i) { + batch_dimension_numbers.push_back(i); + } + std::vector x_config = batch_dimension_numbers; + x_config.push_back(ndims - 2); + x_config.push_back(ndims); + std::vector y_config = batch_dimension_numbers; + y_config.push_back(ndims); + y_config.push_back(ndims - 1); + std::vector output_config = batch_dimension_numbers; + output_config.push_back(ndims - 2); + output_config.push_back(ndims - 1); + return Einsum(x, x_config, y, y_config, output_config, precision); + }); +} + +StatusOr, 3>> ParseEinsumString( + absl::string_view einsum_config) { + std::array, 3> einsum_config_numeric; + std::vector main_split = + absl::StrSplit(einsum_config, ','); - DotDimensionNumbers dot_dnums; - dot_dnums.add_lhs_contracting_dimensions(x_inner_dim); - dot_dnums.add_rhs_contracting_dimensions(y_inner_dim); - for (auto batch_dimension_number : batch_dimension_numbers) { - dot_dnums.add_lhs_batch_dimensions(batch_dimension_number); - dot_dnums.add_rhs_batch_dimensions(batch_dimension_number); + if (main_split.size() != 2) { + return InvalidArgument("Expected one \",\" in einsum_config."); + } + + auto maybe_invalid_character = [](char d) { + if (absl::ascii_isalpha(d)) { + return Status::OK(); + } + if (d == '.') { + return InvalidArgument("Unsupported \"...\" or \".\" in einsum config."); } + return InvalidArgument("Unexpected character in einsum config."); + }; + + auto& x_config = einsum_config_numeric[0]; + x_config.reserve(main_split[0].size()); + for (auto d : main_split[0]) { + TF_RETURN_IF_ERROR(maybe_invalid_character(d)); + x_config.push_back(static_cast(d)); + } + std::vector y_output_split = + absl::StrSplit(main_split[1], "->"); + if (y_output_split.size() != 2) { + return InvalidArgument("Expected one \"->\" in einsum_config."); + } + auto& y_config = einsum_config_numeric[1]; + y_config.reserve(y_output_split[0].size()); + for (auto d : y_output_split[0]) { + TF_RETURN_IF_ERROR(maybe_invalid_character(d)); + y_config.push_back(static_cast(d)); + } + auto& output_config = einsum_config_numeric[2]; + output_config.reserve(y_output_split[1].size()); + for (auto d : y_output_split[1]) { + TF_RETURN_IF_ERROR(maybe_invalid_character(d)); + output_config.push_back(static_cast(d)); + } + return einsum_config_numeric; +} - return DotGeneral(x, y, dot_dnums, &precision_proto); +XlaOp Einsum(XlaOp x, XlaOp y, absl::string_view einsum_config, + PrecisionConfig::Precision precision) { + XlaBuilder* builder = x.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(auto einsum_config_numeric, + ParseEinsumString(einsum_config)); + return Einsum(x, einsum_config_numeric[0], y, einsum_config_numeric[1], + einsum_config_numeric[2], precision); }); } @@ -181,4 +352,5 @@ XlaOp TransposeInMinorDims(XlaOp x) { XlaOp MaybeTransposeInMinorDims(XlaOp x, bool transpose) { return transpose ? TransposeInMinorDims(x) : x; } + } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/matrix.h b/tensorflow/compiler/xla/client/lib/matrix.h index 916cd83748e7028c474065b86bf02d85166d2c9c..60c41ec45a086726086dac7227fc432a9c62d0c8 100644 --- a/tensorflow/compiler/xla/client/lib/matrix.h +++ b/tensorflow/compiler/xla/client/lib/matrix.h @@ -16,7 +16,11 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_MATRIX_H_ #define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_MATRIX_H_ +#include +#include "absl/strings/string_view.h" +#include "absl/types/span.h" #include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -26,10 +30,15 @@ namespace xla { // else. XlaOp IdentityMatrix(XlaBuilder* builder, PrimitiveType type, int64 m, int64 n); -// Get the diagonals of the last two dimensions. If 'x' has shape -// [..., M, N], then the output has shape [..., min(M, N)], containing the -// diagonal elements (i.e., with indices [..., i, i]). -XlaOp GetMatrixDiagonal(XlaOp x); +// Get the diagonals of the last two dimensions. Use k>0 for diagonals above the +// main diagonal, and k<0 for diagonals below the main diagonal. +// +// If 'x' has shape [..., M, N] +// If k >= 0: then the output has shape [..., min(M, N - k)], containing the +// diagonal elements (i.e., with indices [..., i, i + k]). +// If k < 0: then the output has shape [..., min(M + k, N)], containing the +// diagonal elements (i.e., with indices [..., i - k, i]). +XlaOp GetMatrixDiagonal(XlaOp x, int k = 0); // Returns a lower-triangular mask, i.e., true below the `diagonal`-th diagonal // and false above that diagonal. @@ -65,6 +74,40 @@ xla::XlaOp BatchDot( xla::XlaOp x, xla::XlaOp y, xla::PrecisionConfig::Precision precision = xla::PrecisionConfig::DEFAULT); +// Parse an einsum string into dimension numbers: +// "ab,cb->ac" +// becomes: +// {{0, 1},{2, 1},{0, 2}} +// +// NOTE: This function is meant for testing, there is no need to call it +// directly. + +StatusOr, 3>> ParseEinsumString( + absl::string_view einsum_config); + +// Determine if each dimension label is in at least two inputs. +// +// NOTE: This function is meant for testing, there is no need to call it +// directly. +Status ValidateEinsumNumericDimensions(absl::Span x_config, + absl::Span y_config, + absl::Span output_config); + +// Supports two operand einsum notation like "ab,cb->ac". +xla::XlaOp Einsum( + xla::XlaOp x, xla::XlaOp y, absl::string_view einsum_config, + xla::PrecisionConfig::Precision precision = xla::PrecisionConfig::DEFAULT); + +// Same as above but supporting numeric labels on dimensins. So "ab,cb->ac" +// becomes: +// x_config = {0, 1} +// y_config = {2, 1} +// output_config = {0, 2} +xla::XlaOp Einsum( + xla::XlaOp x, absl::Span x_config, xla::XlaOp y, + absl::Span y_config, absl::Span output_config, + xla::PrecisionConfig::Precision precision = xla::PrecisionConfig::DEFAULT); + // Transposes a stack of matrices `x` by swapping the last two dimensions. xla::XlaOp TransposeInMinorDims(xla::XlaOp x); diff --git a/tensorflow/compiler/xla/client/lib/matrix_test.cc b/tensorflow/compiler/xla/client/lib/matrix_test.cc index 0593a7517ac125ca8dc5395cee76f6bc23232cd3..a93fc2ccb92912a10b9b6c2192b81cd73566f2a0 100644 --- a/tensorflow/compiler/xla/client/lib/matrix_test.cc +++ b/tensorflow/compiler/xla/client/lib/matrix_test.cc @@ -15,13 +15,15 @@ limitations under the License. #include "tensorflow/compiler/xla/client/lib/matrix.h" +#include "absl/strings/string_view.h" #include "tensorflow/compiler/xla/client/lib/slicing.h" #include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/status.h" +#include "tensorflow/compiler/xla/statusor.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" -#include "tensorflow/compiler/xla/xla_data.pb.h" namespace xla { namespace { @@ -51,13 +53,24 @@ void MatrixTest::TestMatrixDiagonal() { XlaBuilder builder("GetMatrixDiagonal"); Array3D input(2, 3, 4); input.FillIota(0); - - XlaOp a; - auto a_data = CreateR3Parameter(input, 0, "a", &builder, &a); - GetMatrixDiagonal(a); - Array2D expected({{0, 5, 10}, {12, 17, 22}}); - - ComputeAndCompareR2(&builder, expected, {a_data.get()}); + std::map> k_and_expected = { + {0, {{0, 5, 10}, {12, 17, 22}}}, + {1, {{1, 6, 11}, {13, 18, 23}}}, + {2, {{2, 7}, {14, 19}}}, + {3, {{3}, {15}}}, + {4, {{}, {}}}, + {-1, {{4, 9}, {16, 21}}}, + {-2, {{8}, {20}}}, + {-3, {{}, {}}}, + {-4, {{}, {}}}, + }; + for (const auto& kv : k_and_expected) { + XlaOp a; + auto a_data = CreateR3Parameter(input, 0, "a", &builder, &a); + GetMatrixDiagonal(a, kv.first); + + ComputeAndCompareR2(&builder, kv.second, {a_data.get()}); + } } XLA_TEST_F(MatrixTest, GetMatrixDiagonal_S32) { TestMatrixDiagonal(); } @@ -101,5 +114,78 @@ XLA_TEST_F(MatrixTest, RowBatchDot) { ComputeAndCompareR3(&builder, {{{33}}, {{292}}}, {a_data.get(), row_data.get(), index_data.get()}); } + +XLA_TEST_F(MatrixTest, Einsum) { + XlaBuilder builder(TestName()); + + int n = 4; + + XlaOp a, row, index; + auto a_data = + CreateR3Parameter(BatchedAValsFull(), 0, "a", &builder, &a); + auto row_data = CreateR3Parameter({{{9, 1, 0, 0}}, {{2, 4, 0, 0}}}, 1, + "row", &builder, &row); + // Select {{3, 6, 0, 1}, {24, 61, 82, 48}} out of BatchedAValsFull(). + auto index_data = CreateR0Parameter(1, 2, "index", &builder, &index); + + auto l_index = DynamicSliceInMinorDims( + a, {index, ConstantR0(&builder, 0)}, {1, n}); + Einsum(l_index, row, "abc,adc->abd"); + + ComputeAndCompareR3(&builder, {{{33}}, {{292}}}, + {a_data.get(), row_data.get(), index_data.get()}); +} + +XLA_TEST_F(MatrixTest, ParseEinsumString) { + auto to_vec = [](absl::string_view s) { + std::vector v; + v.reserve(s.size()); + for (auto c : s) { + v.push_back(int64{c}); + } + return v; + }; + + auto to_string = [&](absl::string_view x, absl::string_view y, + absl::string_view o) { + return absl::StrCat(x, ",", y, "->", o); + }; + + std::vector> good_test_cases = {{"ab", "bc", "ac"}, + {"Bab", "Bbc", "Bac"}, + {"ab", "cd", "dcba"}, + {"abc", "abd", "cbd"}}; + for (auto test_case : good_test_cases) { + auto parse_result_or_status = + ParseEinsumString(to_string(test_case[0], test_case[1], test_case[2])); + EXPECT_TRUE(parse_result_or_status.status().ok()); + auto parse_result = parse_result_or_status.ValueOrDie(); + for (int i = 0; i < 3; ++i) { + EXPECT_EQ(parse_result[i], to_vec(test_case[i])); + } + EXPECT_TRUE(ValidateEinsumNumericDimensions( + parse_result[0], parse_result[1], parse_result[2]) + .ok()); + } + + std::vector einsum_strings_that_fail_parsing = { + "", "a", "ab->ba", "ab,bc,cd->ad", "a...b,bc->a...c"}; + for (auto test_case : einsum_strings_that_fail_parsing) { + auto parse_result_or_status = ParseEinsumString(test_case); + EXPECT_FALSE(parse_result_or_status.status().ok()); + } + + std::vector einsum_strings_that_fail_numeric_validation = { + "a,b->c", "ab,bc->acd", "abz,bc->ac", "ab,bcz->ac"}; + for (auto test_case : einsum_strings_that_fail_numeric_validation) { + auto parse_result_or_status = ParseEinsumString(test_case); + EXPECT_TRUE(parse_result_or_status.status().ok()); + auto parse_result = parse_result_or_status.ValueOrDie(); + EXPECT_FALSE(ValidateEinsumNumericDimensions( + parse_result[0], parse_result[1], parse_result[2]) + .ok()); + } +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/self_adjoint_eigen.cc b/tensorflow/compiler/xla/client/lib/self_adjoint_eigen.cc new file mode 100644 index 0000000000000000000000000000000000000000..c2c8caee6ead40a63915160f074edbf84ec6980e --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/self_adjoint_eigen.cc @@ -0,0 +1,398 @@ +/* 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/self_adjoint_eigen.h" + +#include +#include + +#include "tensorflow/compiler/xla/client/lib/arithmetic.h" +#include "tensorflow/compiler/xla/client/lib/constants.h" +#include "tensorflow/compiler/xla/client/lib/loops.h" +#include "tensorflow/compiler/xla/client/lib/math.h" +#include "tensorflow/compiler/xla/client/lib/matrix.h" +#include "tensorflow/compiler/xla/client/lib/slicing.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/core/lib/core/errors.h" + +namespace xla { + +namespace { + +// Jacobi rotation (also known as Givens rotation): +// G = [[ c, s], +// [-s, c]] +// matmul(G_T, G) = I +struct SymmetricSchurDecomposition { + XlaOp c; // cosine. + XlaOp s; // sine. + XlaOp reduction; // Reduction in the off diagonal after applying G. +}; + +// JacobiUpdate holds the intermediate orthogonal matrix, Jacobi-rotated matrix +// and the off-diagonal norm of the rotated matrix. After each Jacobi iteration, +// off-diagonal norm is reduced. +struct JacobiUpdate { + XlaOp v; + XlaOp w; + XlaOp off_diagonal_norm; +}; + +// Given an n-by-n symmetric A and integers p and q that satisfy 0 <= p < q < n, +// it computes a rotation matrix G = [[c, s], [-s, c]], such that +// G_T * A[[p, q], [p, q]] * G +// is diagonalized. +// +// def sym_schur2x2(A, p, q): +// if np.abs(A[p, q]) > 1e-6: +// tau = (A[q, q] - A[p, p]) / (2 * A[p, q]) +// if tau >= 0: +// t = 1.0 / (tau + np.sqrt(1 + tau ** 2)) +// else: +// t = -1.0 / (-tau + np.sqrt(1 + tau ** 2)) +// c = 1.0 / np.sqrt(1.0 + t ** 2) +// s = t * c +// else: +// c = 1.0 +// s = 0.0 +// return c, s +StatusOr SymmetricShurDecomposition2x2(XlaOp a, + XlaOp p, + XlaOp q, + XlaOp tol) { + XlaBuilder* builder = a.builder(); + TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); + + PrimitiveType type = a_shape.element_type(); + + const int64 num_dims = a_shape.rank(); + + auto zero = ScalarLike(a, 0.0); + auto one = ScalarLike(a, 1.0); + auto two = ScalarLike(a, 2.0); + + auto pqs = DynamicSliceInMinorDims(a, {p, q}, {1, 1}); + + auto ps = DynamicSliceInMinorDims(a, {p, p}, {1, 1}); + auto qs = DynamicSliceInMinorDims(a, {q, q}, {1, 1}); + + auto tau = (qs - ps) / (pqs * two); + auto t_pos = one / (tau + Sqrt(one + Square(tau))); + auto t_neg = -one / (-tau + Sqrt(one + Square(tau))); + auto t = Select(Ge(tau, zero), t_pos, t_neg); + + auto c_temp = Rsqrt(one + Square(t)); + auto s_temp = t * c_temp; + + auto c = Select(Ge(Abs(pqs), tol), c_temp, ZerosLike(c_temp) + one); + auto s = Select(Ge(Abs(pqs), tol), s_temp, ZerosLike(s_temp)); + // Renormalize c and s to compensate for low precision arithmetic, this step + // is redundant if high precision float is used, like float64. + auto rnorm = Rsqrt(Square(c) + Square(s)); + + SymmetricSchurDecomposition schur; + + schur.c = c * rnorm; + schur.s = s * rnorm; + schur.reduction = + Reduce(two * Square(pqs), zero, CreateScalarAddComputation(type, builder), + {num_dims - 2, num_dims - 1}); + return schur; +} + +StatusOr Update(JacobiUpdate jacobi_update, XlaOp p, XlaOp q, + XlaOp tol, int64 n) { + XlaBuilder* builder = jacobi_update.w.builder(); + TF_ASSIGN_OR_RETURN( + SymmetricSchurDecomposition schur, + SymmetricShurDecomposition2x2(jacobi_update.w, p, q, tol)); + + TF_ASSIGN_OR_RETURN(Shape w_shape, builder->GetShape(jacobi_update.w)); + const std::vector batch_dims(w_shape.dimensions().begin(), + w_shape.dimensions().end() - 2); + const int64 num_dims = w_shape.rank(); + + auto zero = ScalarLike(p, 0); + + XlaOp c = schur.c; + XlaOp s = schur.s; + + auto slice_p = DynamicSliceInMinorDims(jacobi_update.w, {p, zero}, {1, n}); + auto slice_q = DynamicSliceInMinorDims(jacobi_update.w, {q, zero}, {1, n}); + + auto slice_p_new = c * slice_p - s * slice_q; + auto slice_q_new = s * slice_p + c * slice_q; + + jacobi_update.w = + DynamicUpdateSliceInMinorDims(jacobi_update.w, slice_p_new, {p, zero}); + jacobi_update.w = + DynamicUpdateSliceInMinorDims(jacobi_update.w, slice_q_new, {q, zero}); + + slice_p = DynamicSliceInMinorDims(jacobi_update.w, {zero, p}, {n, 1}); + slice_q = DynamicSliceInMinorDims(jacobi_update.w, {zero, q}, {n, 1}); + + slice_p_new = c * slice_p - s * slice_q; + slice_q_new = s * slice_p + c * slice_q; + + jacobi_update.w = + DynamicUpdateSliceInMinorDims(jacobi_update.w, slice_p_new, {zero, p}); + jacobi_update.w = + DynamicUpdateSliceInMinorDims(jacobi_update.w, slice_q_new, {zero, q}); + + // Zero out a_{pq} explicitly. + std::vector pq_dims(batch_dims.begin(), batch_dims.end()); + pq_dims.push_back(1); + pq_dims.push_back(1); + auto pq_zero = ScalarLike(jacobi_update.w, 0.0); + auto pq_zeros = Broadcast(pq_zero, pq_dims); + jacobi_update.w = + DynamicUpdateSliceInMinorDims(jacobi_update.w, pq_zeros, {p, q}); + jacobi_update.w = + DynamicUpdateSliceInMinorDims(jacobi_update.w, pq_zeros, {q, p}); + + slice_p = DynamicSliceInMinorDims(jacobi_update.v, {zero, p}, {n, 1}); + slice_q = DynamicSliceInMinorDims(jacobi_update.v, {zero, q}, {n, 1}); + + std::vector broadcast_dims(batch_dims.size()); + std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); + broadcast_dims.push_back(num_dims - 1); + + // Renormalize the p-th and q-th columns. This step is redundant if high + // precision floats are used, like 64-bit float. But for 32-bit float, it + // becomes necessary. This step will not increase the overall complexity. + slice_p_new = c * slice_p - s * slice_q; + slice_p_new = Mul( + slice_p_new, + Rsqrt(Reduce(Square(slice_p_new), pq_zero, + CreateScalarAddComputation(w_shape.element_type(), builder), + {num_dims - 2})), + broadcast_dims); + slice_q_new = s * slice_p + c * slice_q; + slice_q_new = Mul( + slice_q_new, + Rsqrt(Reduce(Square(slice_q_new), pq_zero, + CreateScalarAddComputation(w_shape.element_type(), builder), + {num_dims - 2})), + broadcast_dims); + + jacobi_update.v = + DynamicUpdateSliceInMinorDims(jacobi_update.v, slice_p_new, {zero, p}); + jacobi_update.v = + DynamicUpdateSliceInMinorDims(jacobi_update.v, slice_q_new, {zero, q}); + + jacobi_update.off_diagonal_norm = Sqrt( + Max(Square(jacobi_update.off_diagonal_norm) - schur.reduction, pq_zero)); + + return jacobi_update; +} + +StatusOr> WhileLoopFn( + absl::Span initial_values, // + int matrix_dimension, // + int max_sweep_updates, // + PrimitiveType index_type, // + absl::string_view name, // + XlaBuilder* builder) { + auto while_cond_fn = [&](absl::Span values, + XlaBuilder* cond_builder) -> StatusOr { + auto k = values[0]; + auto off_diagonal_norm = values[5]; + // tol = frobenius_norm * epsilon. + auto tol = values[6] * values[7]; + + auto max_sweeps = ScalarLike(k, max_sweep_updates); + + auto sweep_update_cond = Gt(max_sweeps, k); + + auto tol_cond = ReduceAll(Lt(tol, off_diagonal_norm), + xla::ConstantR0(cond_builder, false), + CreateScalarOrComputation(PRED, cond_builder)); + return And(tol_cond, sweep_update_cond); + }; + + auto while_body_fn = + [&](absl::Span values, + XlaBuilder* body_builder) -> StatusOr> { + auto zero = Zero(body_builder, index_type); + auto one = One(body_builder, index_type); + auto end_index = ScalarLike(one, matrix_dimension); + + // Indexes. + XlaOp k = values[0]; + XlaOp p = values[1]; + XlaOp q = values[2]; + + JacobiUpdate jacobi_update; + jacobi_update.v = values[3]; + jacobi_update.w = values[4]; + jacobi_update.off_diagonal_norm = values[5]; + + XlaOp frobenius_norm = values[6]; + XlaOp tol = values[7]; + + TF_ASSIGN_OR_RETURN(jacobi_update, + Update(jacobi_update, p, q, tol, matrix_dimension)); + + std::vector updated_values; + updated_values.reserve(values.size()); + + q = q + one; + p = Select(Eq(q, end_index), p + one, p); + k = Select(Eq(p, end_index - one), k + one, k); + p = Select(Eq(p, end_index - one), zero, p); + q = Select(Eq(q, end_index), p + one, q); + + updated_values.push_back(k); + updated_values.push_back(p); + updated_values.push_back(q); + + updated_values.push_back(jacobi_update.v); + updated_values.push_back(jacobi_update.w); + updated_values.push_back(jacobi_update.off_diagonal_norm); + + updated_values.push_back(frobenius_norm); + updated_values.push_back(tol); + + return updated_values; + }; + std::vector values; + TF_ASSIGN_OR_RETURN(values, WhileLoopHelper(while_cond_fn, while_body_fn, + initial_values, name, builder)); + + return values; +} + +} // namespace + +// This is the cyclic Jacobi iteration. Please note that the eigenvalues are +// possibly not ordered. +// +// def jacobi(A): +// n, _ = A.shape +// V = np.eye(n) +// nfrob = np.sum(A ** 2) +// ndiag = np.sum(np.diag(A) ** 2) +// off = nfrob - ndiag +// while off > 1e-6 * nfrob: +// for p in range(n - 1): +// for q in range(p + 1, n): +// if off > 1e-6 * nfrob: +// c, s = sym_schur2x2(A, p, q) +// off = off - 2 * A[p, q] ** 2 +// A[[p, q], :] = np.matmul(np.array([[c, -s], [s, c]]), +// A[[p, q], :]) +// A[:, [p, q]] = np.matmul(A[:, [p, q]], +// np.array([[c, s], [-s, c]])) +// V[:, [p, q]] = np.matmul(V[:, [p, q]], +// np.array([[c, s], [-s, c]])) +// +// return A, V +// +// TODO(kuny): Implement parallel order Jacobi. +// +SelfAdjointEigenResult SelfAdjointEigen(XlaOp a, bool lower, int64 max_iter, + float epsilon) { + XlaBuilder* builder = a.builder(); + auto return_error = [&](const Status& status) { + SelfAdjointEigenResult result; + result.v = builder->ReportError(status); + result.w = builder->ReportError(status); + return result; + }; + auto shape_with_status = builder->GetShape(a); + if (!shape_with_status.status().ok()) { + return return_error(shape_with_status.status()); + } + Shape a_shape = shape_with_status.ValueOrDie(); + const int64 num_dims = a_shape.rank(); + if (num_dims < 2) { + return return_error(InvalidArgument( + "Arguments to Eigen decomposition must have rank >= 2: got shape %s.", + a_shape.ToString())); + } + PrimitiveType type = a_shape.element_type(); + if (!primitive_util::IsFloatingPointType(type)) { + return return_error(InvalidArgument( + "Type of the input matrix must be float: got %s.", a_shape.ToString())); + } + + const int64 m = ShapeUtil::GetDimension(a_shape, -2); + const int64 n = ShapeUtil::GetDimension(a_shape, -1); + + if (m != n) { + return return_error(InvalidArgument( + "Arguments to Eigen decomposition must be square matrices: got shape " + "(%d, %d).", + m, n)); + } + + const int64 num_batch_dims = num_dims - 2; + std::vector batch_dims(num_batch_dims); + for (int i = 0; i < num_batch_dims; ++i) { + batch_dims[i] = ShapeUtil::GetDimension(a_shape, i); + } + + auto zero = ScalarLike(a, 0.0); + auto tol = ScalarLike(a, epsilon); + + auto v_init = Broadcast(IdentityMatrix(builder, type, m, m), batch_dims); + auto w_init = Triangle(a, lower); + w_init = w_init + TransposeInMinorDims(w_init) - w_init * v_init; + + auto frobenius_norm = Sqrt(Reduce(Square(w_init), zero, + CreateScalarAddComputation(type, builder), + {num_dims - 2, num_dims - 1})); + auto diag = GetMatrixDiagonal(w_init); + auto diag_square = + Reduce(Square(diag), zero, CreateScalarAddComputation(type, builder), + {num_dims - 2}); + + auto off_diagonal_init = + Sqrt(Max(Square(frobenius_norm) - diag_square, zero)); + + auto output_with_status = WhileLoopFn( + { + Zero(builder, S32), // k + Zero(builder, S32), // p + One(builder, S32), // q + v_init, // + w_init, // + off_diagonal_init, // + frobenius_norm, // + tol, // + }, // + n, // + max_iter, // + S32, // + "CyclicJacobi", // + builder); + if (!output_with_status.status().ok()) { + return return_error(output_with_status.status()); + } + + auto output = output_with_status.ValueOrDie(); + + SelfAdjointEigenResult result; + result.v = output[3]; + result.w = GetMatrixDiagonal(output[4]); + + return result; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/self_adjoint_eigen.h b/tensorflow/compiler/xla/client/lib/self_adjoint_eigen.h new file mode 100644 index 0000000000000000000000000000000000000000..49fc17aa275a8e831e800069290db2dd047416e4 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/self_adjoint_eigen.h @@ -0,0 +1,42 @@ +/* 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_CLIENT_LIB_SELF_ADJOINT_EIGEN_H_ +#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_SELF_ADJOINT_EIGEN_H_ + +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { + +// The eigenvalue decomposition of a symmetric matrix, the original matrix is +// recovered by v * w * v_t. +struct SelfAdjointEigenResult { + // The i-th column is the normalized eigenvector corresponding to the + // eigenvalue w[i]. Will return a matrix object if a is a matrix object. + XlaOp v; + // TODO(kuny): Sort the eigenvalues. + // The eigenvalues in ascending order, each repeated according to its + // multiplicity. + XlaOp w; +}; + +SelfAdjointEigenResult SelfAdjointEigen(XlaOp a, bool lower = true, + int64 max_iter = 100, + float epsilon = 1e-6); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_SELF_ADJOINT_EIGEN_H_ diff --git a/tensorflow/compiler/xla/client/lib/self_adjoint_eigen_test.cc b/tensorflow/compiler/xla/client/lib/self_adjoint_eigen_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..720c49b7716df4c2f1a1fe2d6cda9fe1b135acce --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/self_adjoint_eigen_test.cc @@ -0,0 +1,300 @@ +/* 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/self_adjoint_eigen.h" + +#include "tensorflow/compiler/xla/array2d.h" +#include "tensorflow/compiler/xla/array3d.h" +#include "tensorflow/compiler/xla/client/lib/arithmetic.h" +#include "tensorflow/compiler/xla/client/lib/constants.h" +#include "tensorflow/compiler/xla/client/lib/matrix.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/test.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" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/status_test_util.h" + +namespace xla { + +class SelfAdjointEigenTest : public ClientLibraryTestBase { + protected: + void SetUp() override { + ClientLibraryTestBase::SetUp(); + batch_3d_4x4_ = Array3D{ + { + {4, 6, 8, 10}, + {6, 45, 54, 63}, + {8, 54, 146, 166}, + {10, 63, 166, 310}, + }, + { + {16, 24, 8, 12}, + {24, 61, 82, 48}, + {8, 82, 100, 6}, + {12, 48, 6, 62}, + }, + }; + matrix2d_8x8_ = Array2D{ + {14., 123., 49., 112., 115., 173., 182., 125.}, + {123., 14., 60., 118., 150., 130., 91., 72.}, + {49., 60., 138., 111., 106., 101., 115., 142.}, + {112., 118., 111., 142., 91., 130., 25., 61.}, + {115., 150., 106., 91., 116., 121., 128., 85.}, + {173., 130., 101., 130., 121., 70., 151., 132.}, + {182., 91., 115., 25., 128., 151., 66., 92.}, + {125., 72., 142., 61., 85., 132., 92., 156.}, + }; + low_rank_4x4_ = Array2D{ + // x = [[1, 2, 3, 4], [1, -1, 1, -1]] + // matmul(x.T, x) + {2, 1, 4, 3}, + {1, 5, 5, 9}, + {4, 5, 10, 11}, + {3, 9, 11, 17}, + }; + } + void TearDown() override { ClientLibraryTestBase::TearDown(); } + + Array3D get_unit_matrix_3d(const Array3D& matrix) { + Array3D result(matrix.n1(), matrix.n2(), matrix.n3(), 0.0); + for (int i = 0; i < matrix.n1(); ++i) { + for (int j = 0; j < matrix.n2(); ++j) { + result({i, j, j}) = 1.0; + } + } + return result; + } + + Array3D ExtractTriangularMatrix(const Array3D& matrix, + bool lower) { + Array3D result(matrix); + for (int i = 0; i < result.n1(); ++i) { + for (int j = 0; j < result.n2(); ++j) { + if (lower) { + for (int k = j + 1; k < result.n3(); ++k) { + result({i, j, k}) = 0.0; + } + } else { + for (int k = 0; k < j; ++k) { + result({i, j, k}) = 0.0; + } + } + } + } + return result; + } + + XlaOp ComputeMatmulVWVt(SelfAdjointEigenResult result, XlaBuilder* builder) { + Shape shape = builder->GetShape(result.v).ValueOrDie(); + std::vector out_dims = shape.dimensions(); + std::vector broadcast_dims(shape.rank() - 1); + std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0); + + broadcast_dims[shape.rank() - 2] = shape.rank() - 1; + auto vw = Mul(result.v, BroadcastInDim(result.w, out_dims, broadcast_dims)); + return BatchDot(vw, TransposeInMinorDims(result.v), + PrecisionConfig::HIGHEST); + } + + XlaOp GetAverageAbsoluteError(XlaOp m1, XlaOp m2, XlaBuilder* builder) { + Shape shape = builder->GetShape(m1).ValueOrDie(); + int64 size = 1; + for (auto d : shape.dimensions()) { + size *= d; + } + return ReduceAll(Abs(m1 - m2), ConstantR0WithType(builder, F32, 0), + CreateScalarAddComputation(F32, builder)) / + ConstantR0WithType(builder, F32, size); + } + + Array2D GenerateRandomSymmetricMatrix(int size) { + Array2D result{size, size, 0.0}; + result.FillRandom(10 /* stddev */, 2 /* mean */); + for (int i = 0; i < size; ++i) { + for (int j = 0; j < i; ++j) { + result({j, i}) = result({i, j}); + } + } + return result; + } + + Array3D batch_3d_4x4_; + Array2D matrix2d_8x8_; + Array2D low_rank_4x4_; + Array2D wrong_type_4x4_; +}; + +XLA_TEST_F(SelfAdjointEigenTest, Test_VWVt_EQ_A_2x4x4) { + XlaBuilder builder(TestName()); + + XlaOp a; + auto a_data = CreateR3Parameter(batch_3d_4x4_, 0, "a", &builder, &a); + auto result = SelfAdjointEigen(a); + ComputeMatmulVWVt(result, &builder); + + ComputeAndCompareR3(&builder, batch_3d_4x4_, {a_data.get()}, + ErrorSpec(1e-3, 1e-3)); +} + +XLA_TEST_F(SelfAdjointEigenTest, Test_VWVt_EQ_A_Lower_2x4x4) { + XlaBuilder builder(TestName()); + + XlaOp a; + auto a_data = CreateR3Parameter( + ExtractTriangularMatrix(batch_3d_4x4_, true), 0, "a", &builder, &a); + auto result = SelfAdjointEigen(a); + ComputeMatmulVWVt(result, &builder); + + ComputeAndCompareR3(&builder, batch_3d_4x4_, {a_data.get()}, + ErrorSpec(1e-3, 1e-3)); +} + +XLA_TEST_F(SelfAdjointEigenTest, Test_VWVt_EQ_A_Upper_2x4x4) { + XlaBuilder builder(TestName()); + + XlaOp a; + auto a_data = CreateR3Parameter( + ExtractTriangularMatrix(batch_3d_4x4_, false), 0, "a", &builder, &a); + auto result = SelfAdjointEigen(a, false); + ComputeMatmulVWVt(result, &builder); + + ComputeAndCompareR3(&builder, batch_3d_4x4_, {a_data.get()}, + ErrorSpec(1e-3, 1e-3)); +} + +XLA_TEST_F(SelfAdjointEigenTest, Test_Orthogonality_2x4x4) { + XlaBuilder builder(TestName()); + + XlaOp a; + auto a_data = CreateR3Parameter(batch_3d_4x4_, 0, "a", &builder, &a); + auto result = SelfAdjointEigen(a); + BatchDot(result.v, TransposeInMinorDims(result.v), PrecisionConfig::HIGHEST); + + ComputeAndCompareR3(&builder, get_unit_matrix_3d(batch_3d_4x4_), + {a_data.get()}, ErrorSpec(1e-3, 1e-3)); +} + +XLA_TEST_F(SelfAdjointEigenTest, Test_VtWV_EQ_A_Rank_Deficient_4x4) { + XlaBuilder builder(TestName()); + + XlaOp a; + auto a_data = CreateR2Parameter(low_rank_4x4_, 0, "a", &builder, &a); + auto result = SelfAdjointEigen(a); + ComputeMatmulVWVt(result, &builder); + + ComputeAndCompareR2(&builder, low_rank_4x4_, {a_data.get()}, + ErrorSpec(1e-3, 1e-3)); +} + +XLA_TEST_F(SelfAdjointEigenTest, Test_Eigen_8x8) { + XlaBuilder builder(TestName()); + + // This is computed by numpy.linalg.eigh with float32. + std::vector expected{-182.69205, -116.86245, -105.74489, -9.545369, + 37.81711, 104.732285, 120.29153, 868.00385}; + + XlaOp a; + auto a_data = CreateR2Parameter(matrix2d_8x8_, 0, "a", &builder, &a); + auto result = SelfAdjointEigen(a); + Sort(result.w); + + ComputeAndCompareR1(&builder, expected, {a_data.get()}, + ErrorSpec(1e-3, 1e-3)); +} + +XLA_TEST_F(SelfAdjointEigenTest, Test_Orthogonality_8x8) { + XlaBuilder builder(TestName()); + + float expected_vals = 1e-3; + + XlaOp a; + auto a_data = CreateR2Parameter(matrix2d_8x8_, 0, "a", &builder, &a); + auto result = SelfAdjointEigen(a); + // np.sum(norm(eye(n) - matmul(conj(T(v)), v)) / n**2 + GetAverageAbsoluteError(IdentityMatrix(&builder, F32, 8, 8), + BatchDot(TransposeInMinorDims(result.v), result.v), + &builder); + + ComputeAndCompareR0(&builder, expected_vals, {a_data.get()}, + ErrorSpec(1e-3, 1e-3)); +} + +XLA_TEST_F(SelfAdjointEigenTest, Wrong_Type_Int) { + XlaBuilder builder(TestName()); + + XlaOp a; + auto a_data = CreateR2Parameter(wrong_type_4x4_, 0, "a", &builder, &a); + auto result = SelfAdjointEigen(a); + EXPECT_FALSE(result.v.valid()); + EXPECT_FALSE(result.w.valid()); +} + +XLA_TEST_F(SelfAdjointEigenTest, Various_Size_Random_Matrix_8x8) { + XlaBuilder builder(TestName()); + int size = 8; + Array2D a_val = GenerateRandomSymmetricMatrix(size); + XlaOp a; + auto a_data = CreateR2Parameter(a_val, 0, "a", &builder, &a); + auto result = SelfAdjointEigen(a); + GetAverageAbsoluteError(ComputeMatmulVWVt(result, &builder), a, &builder); + + ComputeAndCompareR0(&builder, 1e-3, {a_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(SelfAdjointEigenTest, Various_Size_Random_Matrix_16x16) { + XlaBuilder builder(TestName()); + int size = 16; + Array2D a_val = GenerateRandomSymmetricMatrix(size); + XlaOp a; + auto a_data = CreateR2Parameter(a_val, 0, "a", &builder, &a); + auto result = SelfAdjointEigen(a); + GetAverageAbsoluteError(ComputeMatmulVWVt(result, &builder), a, &builder); + + ComputeAndCompareR0(&builder, 1e-3, {a_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(SelfAdjointEigenTest, Various_Size_Random_Matrix_32x32) { + XlaBuilder builder(TestName()); + int size = 32; + Array2D a_val = GenerateRandomSymmetricMatrix(size); + XlaOp a; + auto a_data = CreateR2Parameter(a_val, 0, "a", &builder, &a); + auto result = SelfAdjointEigen(a); + GetAverageAbsoluteError(ComputeMatmulVWVt(result, &builder), a, &builder); + + ComputeAndCompareR0(&builder, 1e-3, {a_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(SelfAdjointEigenTest, Various_Size_Random_Matrix_64x64) { + XlaBuilder builder(TestName()); + int size = 64; + Array2D a_val = GenerateRandomSymmetricMatrix(size); + XlaOp a; + auto a_data = CreateR2Parameter(a_val, 0, "a", &builder, &a); + auto result = SelfAdjointEigen(a); + GetAverageAbsoluteError(ComputeMatmulVWVt(result, &builder), a, &builder); + + ComputeAndCompareR0(&builder, 1e-3, {a_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/slicing.cc b/tensorflow/compiler/xla/client/lib/slicing.cc index 611fffba8d0544a011cb641802058c08d95cf5f7..77145ba7d4c72435450d3e33d57b2507eb84d2fc 100644 --- a/tensorflow/compiler/xla/client/lib/slicing.cc +++ b/tensorflow/compiler/xla/client/lib/slicing.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/client/lib/slicing.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" namespace xla { @@ -51,17 +52,17 @@ XlaOp SliceInMinorDims(XlaOp x, absl::Span start, XlaOp UpdateSlice(XlaOp x, XlaOp update, absl::Span start) { XlaBuilder* builder = x.builder(); return builder->ReportErrorOrReturn([&]() -> StatusOr { - // TODO(phawkins): make int64 work on all backends, remove the int32 cast. - std::vector start_as_int32(start.begin(), start.end()); - auto start_constant = ConstantR1(builder, start_as_int32); TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); const int64 n_dims = shape.rank(); - TF_ASSIGN_OR_RETURN(Shape start_constant_shape, - builder->GetShape(start_constant)); - const int64 start_length = - ShapeUtil::GetDimension(start_constant_shape, -1); - TF_RET_CHECK(start_length == n_dims); - return DynamicUpdateSlice(x, update, start_constant); + TF_RET_CHECK(start.size() == n_dims); + + // TODO(phawkins): make int64 work on all backends, remove the int32 cast. + std::vector start_as_int32(start.begin(), start.end()); + std::vector start_ops(start.size()); + for (int i = 0; i < start.size(); ++i) { + start_ops[i] = ConstantR0(builder, start_as_int32[i]); + } + return DynamicUpdateSlice(x, update, start_ops); }); } @@ -90,18 +91,17 @@ std::vector ConcatVectors(absl::Span xs, return output; } -XlaOp PrependZerosInMajorDims(XlaOp x, absl::Span starts) { +StatusOr> PrependZerosInMajorDims( + XlaOp x, absl::Span starts) { XlaBuilder* builder = x.builder(); - return builder->ReportErrorOrReturn([&]() -> StatusOr { - TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); - const int64 n_dims = shape.rank(); - auto zero = Reshape(ConstantR0(builder, 0), {1}); - std::vector padded_starts(n_dims, zero); - for (int i = 0; i < starts.size(); ++i) { - padded_starts[n_dims - starts.size() + i] = Reshape(starts[i], {1}); - } - return ConcatInDim(builder, padded_starts, 0); - }); + TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); + const int64 n_dims = shape.rank(); + auto zero = ConstantR0(builder, 0); + std::vector padded_starts(n_dims, zero); + for (int i = 0; i < starts.size(); ++i) { + padded_starts[n_dims - starts.size() + i] = starts[i]; + } + return padded_starts; } } // namespace @@ -119,7 +119,7 @@ XlaOp DynamicSliceInMinorDims(XlaOp x, absl::Span starts, .subspan( /*pos=*/0, /*len=*/n_dims - sizes.size()); - auto padded_starts = PrependZerosInMajorDims(x, starts); + TF_ASSIGN_OR_RETURN(auto padded_starts, PrependZerosInMajorDims(x, starts)); auto padded_sizes = ConcatVectors(major_dims, sizes); return DynamicSlice(x, padded_starts, padded_sizes); }); @@ -127,8 +127,11 @@ XlaOp DynamicSliceInMinorDims(XlaOp x, absl::Span starts, XlaOp DynamicUpdateSliceInMinorDims(XlaOp x, XlaOp update, absl::Span starts) { - auto padded_starts = PrependZerosInMajorDims(x, starts); - return DynamicUpdateSlice(x, update, padded_starts); + XlaBuilder* builder = x.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(auto padded_starts, PrependZerosInMajorDims(x, starts)); + return DynamicUpdateSlice(x, update, padded_starts); + }); } } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/sorting.cc b/tensorflow/compiler/xla/client/lib/sorting.cc index e8553a08bb014e790822a14e128686b60b8d6b7c..3245f46e6fd6f365f2e4ee90b3c88cf1bd3b5b85 100644 --- a/tensorflow/compiler/xla/client/lib/sorting.cc +++ b/tensorflow/compiler/xla/client/lib/sorting.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/client/lib/sorting.h" +#include "tensorflow/compiler/xla/client/lib/comparators.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/util.h" @@ -30,7 +31,12 @@ XlaOp TopK(XlaOp input, int64 k) { ShapeUtil::MakeShape(S32, AsInt64Slice(input_shape.dimensions())); XlaOp iota_s32 = Iota(builder, iota_shape, last_dim); auto input_dims = input_shape.dimensions(); - XlaOp sort_result = Sort(Neg(input), {iota_s32}); + // TODO(b/122298745): Get rid of Neg() and use CreateScalarGtComputation + // once the TPU backend supports the comparison computations. + XlaOp sort_result = + Sort({Neg(input), iota_s32}, + CreateScalarLtComputation({input_shape.element_type(), S32}, + iota_s32.builder())); std::vector start_indices(input_shape.dimensions_size(), 0); std::vector limit_indices(input_dims.begin(), input_dims.end()); limit_indices[last_dim] = k; diff --git a/tensorflow/compiler/xla/client/lib/sorting_test.cc b/tensorflow/compiler/xla/client/lib/sorting_test.cc index 27ff36c7491ab8397d46f3a49493ff2b904deb2d..ae78910a5b416ceba6da9286b42dde9fb4ebced5 100644 --- a/tensorflow/compiler/xla/client/lib/sorting_test.cc +++ b/tensorflow/compiler/xla/client/lib/sorting_test.cc @@ -77,11 +77,13 @@ XLA_TEST_F(SortingTest, TopKFullSort) { auto x = ConstantR1(&builder, inputs); xla::GetTupleElement(xla::TopK(x, kSize), 0); - std::sort(inputs.begin(), inputs.end(), std::greater()); + absl::c_sort(inputs, std::greater()); ComputeAndCompareR1(&builder, inputs, {}); } -XLA_TEST_F(SortingTest, TopKFullSortWithDuplicates) { +// TODO(b/122298745): Enable this test when the GPU backend supports stable +// sorting. +XLA_TEST_F(SortingTest, DISABLED_ON_GPU(TopKFullSortWithDuplicates)) { XlaBuilder builder(TestName()); XlaOp a; auto a_data = CreateR1Parameter({1, 1, 2, 2, 1}, 0, "a", &builder, &a); diff --git a/tensorflow/compiler/xla/client/lib/triangular_solve.h b/tensorflow/compiler/xla/client/lib/triangular_solve.h deleted file mode 100644 index 50a3b30ebd1c15eb6d2ace4e351cb41f21db7093..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/client/lib/triangular_solve.h +++ /dev/null @@ -1,67 +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. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_TRIANGULAR_SOLVE_H_ -#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_TRIANGULAR_SOLVE_H_ - -#include "tensorflow/compiler/xla/client/xla_builder.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" - -namespace xla { - -// 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))`. -// That is, the innermost matrices in the output satisfy a scalar system -// depending on the value of the value of (left_side, transpose_a, conjugate_a) -// according to: -// (F, F, F) => `output[..., i, k] a[..., k, j] = b[..., i, j]`, -// (F, F, T) => `output[..., i, k] a*[..., k, j] = b[..., i, j]`, -// (F, T, F) => `output[..., i, k] a[..., j, k] = b[..., i, j]`, -// (F, T, T) => `output[..., i, k] a*[..., j, k] = b[..., i, j]`, -// (T, F, F) => ` a[..., i, k] output[..., k, j] = b[..., i, j]`, -// (T, F, T) => `a*[..., i, k] output[..., k, j] = b[..., i, j]`, -// (T, T, F) => ` a[..., i, k] output[..., j, k] = b[..., i, j]`, -// (T, T, T) => `a*[..., i, k] output[..., j, k] = b[..., i, j]`, -// where * denotes complex conjugation and where the index `k` is summed over. -// -// `a` is a tensor of shape `[..., M, M]` whose innermost 2 dimensions form -// square matrices. If lower is true (false), then the strictly upper (lower) -// triangular part of each innermost matrix in `a` is assumed to be zero and is -// not accessed. -// `b` is a tensor of shape `[..., M, K]` if left_side is true, otherwise a -// tensor of shape `[..., K, M]`. -// `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). -// `transpose_a` is a boolean indicating whether the matrix `a` is transposed. -// `conjugate_a` is a boolean indicating whether the entries of `a` are complex -// conjugated (independently of whether they are transposed), so that when both -// transpose_a and conjugate_a are true the effect is a Hermitian adjoint. -// -// Uses a blocked algorithm if `block_size` is > 1; if block_size == 1 then no -// blocking is used. -XlaOp TriangularSolve( - XlaOp a, XlaOp b, bool left_side, bool lower, bool transpose_a, - bool conjugate_a, int64 block_size = 128, - PrecisionConfig::Precision precision = PrecisionConfig::HIGHEST); - -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_TRIANGULAR_SOLVE_H_ diff --git a/tensorflow/compiler/xla/client/local_client.cc b/tensorflow/compiler/xla/client/local_client.cc index 049cd15738a619294b19d5cf74ca514d7b4a00ad..48b5f94538f453785194bc434a91ee0a10c020c2 100644 --- a/tensorflow/compiler/xla/client/local_client.cc +++ b/tensorflow/compiler/xla/client/local_client.cc @@ -164,9 +164,8 @@ StatusOr LocalExecutable::Run( // ExecutableRunOptions.eigen_intra_op_thread_pool. // *) The thread pool used for XLA CPU ops is from // backend_->eigen_intra_op_thread_pool(). - ServiceExecutableRunOptions service_options( - run_options, backend_->StreamBorrower(), - backend_->eigen_intra_op_thread_pool()); + ServiceExecutableRunOptions service_options(run_options, + backend_->StreamBorrower()); if (executable_->dumping_snapshot()) { return ExecuteAndDump(&service_options, arguments); diff --git a/tensorflow/compiler/xla/client/local_client.h b/tensorflow/compiler/xla/client/local_client.h index ddb36680e8b185b053368baffa6f1d5cac50dc07..4f4fc8df31c633749ae9b6dafcdc38d4fd1eba40 100644 --- a/tensorflow/compiler/xla/client/local_client.h +++ b/tensorflow/compiler/xla/client/local_client.h @@ -114,7 +114,7 @@ class LocalClient : public Client { // Build and return a LocalExecutable object. The executable is compiled using // the given XlaComputation, argument layouts and options. // - // The given ExecutableBuildOptions override any values from TF_XLA_FLAGS + // The given ExecutableBuildOptions overrides any values from XLA_FLAGS // environment variable. StatusOr> Compile( const XlaComputation& computation, diff --git a/tensorflow/compiler/xla/client/xla_builder.cc b/tensorflow/compiler/xla/client/xla_builder.cc index 4b752def1ae2decc95b6d791d74386e9e8d2c731..fb9dbe851e7db8da3a496c40a63b39f80cf1ff33 100644 --- a/tensorflow/compiler/xla/client/xla_builder.cc +++ b/tensorflow/compiler/xla/client/xla_builder.cc @@ -22,6 +22,8 @@ limitations under the License. #include #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" #include "absl/memory/memory.h" #include "absl/strings/match.h" #include "absl/strings/str_cat.h" @@ -193,9 +195,9 @@ StatusOr XlaBuilder::GetProgramShape(XlaOp root) const { } void XlaBuilder::IsConstantVisitor(const int64 op_handle, - std::set* visited, + absl::flat_hash_set* visited, bool* is_constant) const { - if (visited->count(op_handle) != 0 || !*is_constant) { + if (visited->contains(op_handle) || !*is_constant) { return; } @@ -209,11 +211,21 @@ void XlaBuilder::IsConstantVisitor(const int64 op_handle, } // TODO(b/32495713): We aren't checking the called computations. break; + case HloOpcode::kGetDimensionSize: { + int64 dimension_number = instr.dimensions(0); + const HloInstructionProto& operand = + *(LookUpInstructionByHandle(instr.operand_ids(0)).ValueOrDie()); + Shape operand_shape(operand.shape()); + if (operand_shape.is_dynamic_dimension(dimension_number)) { + *is_constant = false; + } + break; + } // Non functional ops. case HloOpcode::kRng: case HloOpcode::kAllReduce: - // TODO(b/33009255): Implmement constant folding for cross replica sum. + // TODO(b/33009255): Implement constant folding for cross replica sum. case HloOpcode::kInfeed: case HloOpcode::kOutfeed: case HloOpcode::kCall: @@ -245,6 +257,29 @@ Status XlaBuilder::SetDynamicBinding(int64 dynamic_size_param_num, int64 target_param_num, ShapeIndex target_param_index, int64 target_dim_num) { + bool param_exists = false; + for (HloInstructionProto& instr : instructions_) { + if (instr.opcode() == HloOpcodeString(HloOpcode::kParameter) && + instr.parameter_number() == target_param_num) { + param_exists = true; + Shape param_shape(instr.shape()); + Shape* param_shape_ptr = ¶m_shape; + for (int64 index : target_param_index) { + param_shape_ptr = param_shape_ptr->mutable_tuple_shapes(index); + } + param_shape_ptr->set_dynamic_dimension(target_dim_num, + /*is_dynamic=*/true); + *instr.mutable_shape() = param_shape.ToProto(); + } + } + + if (!param_exists) { + return InvalidArgument( + "Asked to mark parameter %lld as dynamic sized parameter, but the " + "doesn't exists", + target_param_num); + } + TF_RETURN_IF_ERROR(dynamic_parameter_binding_.Bind( DynamicParameterBinding::DynamicParameter{dynamic_size_param_num, dynamic_size_param_index}, @@ -264,29 +299,52 @@ XlaComputation XlaBuilder::BuildAndNoteError() { return build_status.ConsumeValueOrDie(); } -StatusOr XlaBuilder::Build() { +StatusOr XlaBuilder::Build(bool remove_dynamic_dimensions) { if (!first_error_.ok()) { string backtrace; first_error_backtrace_.Dump(tensorflow::DebugWriteToString, &backtrace); return AppendStatus(first_error_, backtrace); } - return Build(instructions_.back().id()); + return Build(instructions_.back().id(), remove_dynamic_dimensions); } -StatusOr XlaBuilder::Build(XlaOp root) { +StatusOr XlaBuilder::Build(XlaOp root, + bool remove_dynamic_dimensions) { if (root.builder_ != this) { return InvalidArgument("Given root operation is not in this computation."); } - return Build(root.handle()); + return Build(root.handle(), remove_dynamic_dimensions); } -StatusOr XlaBuilder::Build(int64 root_id) { +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); } + // TODO(b/121223198): XLA backend cannot handle dynamic dimensions yet, remove + // all dynamic dimensions before building xla program until we have support in + // the backend. + if (remove_dynamic_dimensions) { + std::function remove_dynamic_dimension = + [&](ShapeProto* shape) { + if (shape->tuple_shapes_size() != 0) { + for (int64 i = 0; i < shape->tuple_shapes_size(); ++i) { + remove_dynamic_dimension(shape->mutable_tuple_shapes(i)); + } + } + for (int64 i = 0; i < shape->dimensions_size(); ++i) { + shape->set_is_dynamic_dimension(i, false); + } + }; + + for (auto& instruction : instructions_) { + remove_dynamic_dimension(instruction.mutable_shape()); + } + } + HloComputationProto entry; SetProtoIdAndName(&entry, name_, kNameSeparator, GetNextId()); TF_ASSIGN_OR_RETURN(ProgramShape program_shape, GetProgramShape(root_id)); @@ -348,8 +406,9 @@ StatusOr XlaBuilder::Build(int64 root_id) { alias.param_number, alias.param_index.ToString().c_str()); } - TF_RETURN_IF_ERROR(config.SetUpAlias(alias.output_index, alias.param_number, - alias.param_index)); + TF_RETURN_IF_ERROR(config.SetUpAlias( + alias.output_index, alias.param_number, alias.param_index, + HloInputOutputAliasConfig::AliasKind::kUserAlias)); } *module->mutable_input_output_alias() = config.ToProto(); return Status::OK(); @@ -442,16 +501,19 @@ XlaOp XlaBuilder::BinaryOp(HloOpcode binop, const XlaOp& lhs, const XlaOp& rhs, const Shape& from_shape = should_broadcast_lhs ? lhs_shape : rhs_shape; std::vector to_size; - for (int64 size : shape.dimensions()) { - to_size.push_back(size); + std::vector to_size_is_dynamic; + for (int i = 0; i < shape.rank(); i++) { + to_size.push_back(shape.dimensions(i)); + to_size_is_dynamic.push_back(shape.is_dynamic_dimension(i)); } for (int64 from_dim = 0; from_dim < from_shape.rank(); from_dim++) { int64 to_dim = broadcast_dimensions[from_dim]; to_size[to_dim] = from_shape.dimensions(from_dim); + to_size_is_dynamic[to_dim] = from_shape.is_dynamic_dimension(from_dim); } - const Shape& broadcasted_shape = - ShapeUtil::MakeShape(from_shape.element_type(), to_size); + const Shape& broadcasted_shape = ShapeUtil::MakeShape( + from_shape.element_type(), to_size, to_size_is_dynamic); TF_ASSIGN_OR_RETURN( XlaOp broadcasted_operand, InDimBroadcast(broadcasted_shape, from, broadcast_dimensions)); @@ -511,16 +573,6 @@ XlaOp XlaBuilder::TernaryOp(HloOpcode triop, const XlaOp& lhs, const XlaOp& rhs, }); } -XlaOp XlaBuilder::Add(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kAdd, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::Mul(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kMultiply, lhs, rhs, broadcast_dimensions); -} - XlaOp XlaBuilder::ConstantLiteral(const LiteralSlice& literal) { return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; @@ -610,8 +662,17 @@ XlaOp XlaBuilder::BroadcastInDim( TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); // Output shape, in the case of degenerate broadcast, the out_dim_size is // not necessarily the same as the dimension sizes of the output shape. - const auto& output_shape = + auto output_shape = ShapeUtil::MakeShape(operand_shape.element_type(), out_dim_size); + for (int i = 0; i < broadcast_dimensions.size(); i++) { + if (broadcast_dimensions[i] < 0 || + broadcast_dimensions[i] > out_dim_size.size()) { + return InvalidArgument("Broadcast dimension %lld is out of bound", + broadcast_dimensions[i]); + } + output_shape.set_dynamic_dimension(broadcast_dimensions[i], + operand_shape.is_dynamic_dimension(i)); + } TF_RETURN_IF_ERROR(ShapeInference::InferBroadcastShape( operand_shape, output_shape, broadcast_dimensions) @@ -703,6 +764,34 @@ XlaOp XlaBuilder::DynamicSlice(const XlaOp& operand, const XlaOp& start_indices, }); } +XlaOp XlaBuilder::DynamicSlice(const XlaOp& operand, + absl::Span start_indices, + absl::Span slice_sizes) { + return ReportErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + std::vector start_indices_shape_ptrs; + TF_ASSIGN_OR_RETURN(const auto& start_indices_shapes, + GetOperandShapes(start_indices)); + absl::c_transform(start_indices_shapes, + std::back_inserter(start_indices_shape_ptrs), + [](const Shape& shape) { return &shape; }); + TF_ASSIGN_OR_RETURN(Shape shape, + ShapeInference::InferDynamicSliceShape( + operand_shape, start_indices_shapes, slice_sizes)); + *instr.mutable_shape() = shape.ToProto(); + + for (int64 size : slice_sizes) { + instr.add_dynamic_slice_sizes(size); + } + + std::vector operands = {operand}; + operands.insert(operands.end(), start_indices.begin(), start_indices.end()); + return AddInstruction(std::move(instr), HloOpcode::kDynamicSlice, operands); + }); +} + XlaOp XlaBuilder::DynamicUpdateSlice(const XlaOp& operand, const XlaOp& update, const XlaOp& start_indices) { return ReportErrorOrReturn([&]() -> StatusOr { @@ -722,6 +811,31 @@ XlaOp XlaBuilder::DynamicUpdateSlice(const XlaOp& operand, const XlaOp& update, }); } +XlaOp XlaBuilder::DynamicUpdateSlice(const XlaOp& operand, const XlaOp& update, + absl::Span start_indices) { + return ReportErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + TF_ASSIGN_OR_RETURN(const Shape& update_shape, GetShape(update)); + std::vector start_indices_shape_ptrs; + TF_ASSIGN_OR_RETURN(const auto& start_indices_shapes, + GetOperandShapes(start_indices)); + absl::c_transform(start_indices_shapes, + std::back_inserter(start_indices_shape_ptrs), + [](const Shape& shape) { return &shape; }); + TF_ASSIGN_OR_RETURN(Shape shape, + ShapeInference::InferDynamicUpdateSliceShape( + operand_shape, update_shape, start_indices_shapes)); + *instr.mutable_shape() = shape.ToProto(); + + std::vector operands = {operand, update}; + operands.insert(operands.end(), start_indices.begin(), start_indices.end()); + return AddInstruction(std::move(instr), HloOpcode::kDynamicUpdateSlice, + operands); + }); +} + XlaOp XlaBuilder::ConcatInDim(absl::Span operands, int64 dimension) { return ReportErrorOrReturn([&]() -> StatusOr { @@ -880,36 +994,6 @@ XlaOp XlaBuilder::GetTupleElement(const XlaOp& tuple_data, int64 index) { }); } -XlaOp XlaBuilder::Eq(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kEq, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::Ne(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kNe, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::Ge(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kGe, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::Gt(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kGt, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::Le(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kLe, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::Lt(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kLt, lhs, rhs, broadcast_dimensions); -} - XlaOp XlaBuilder::Dot(const XlaOp& lhs, const XlaOp& rhs, const PrecisionConfig* precision_config) { return ReportErrorOrReturn([&]() -> StatusOr { @@ -930,6 +1014,18 @@ XlaOp XlaBuilder::DotGeneral(const XlaOp& lhs, const XlaOp& rhs, HloInstructionProto instr; TF_ASSIGN_OR_RETURN(const Shape& lhs_shape, GetShape(lhs)); TF_ASSIGN_OR_RETURN(const Shape& rhs_shape, GetShape(rhs)); + // If one operand is a scalar, just multiply the two operands. + if (ShapeUtil::IsScalar(lhs_shape) || ShapeUtil::IsScalar(rhs_shape)) { + if (dimension_numbers.rhs_batch_dimensions_size() != 0 || + dimension_numbers.lhs_batch_dimensions_size() != 0 || + dimension_numbers.rhs_contracting_dimensions_size() != 0 || + dimension_numbers.lhs_contracting_dimensions_size() != 0) { + return InvalidArgument( + "Dots with scalar operands must have no contracting or batch " + "dimensions"); + } + return xla::Mul(lhs, rhs); + } TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferDotOpShape(lhs_shape, rhs_shape, dimension_numbers)); @@ -1425,147 +1521,6 @@ XlaOp XlaBuilder::CustomCall( }); } -XlaOp XlaBuilder::Complex(const XlaOp& real, const XlaOp& imag, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kComplex, real, imag, broadcast_dimensions); -} - -XlaOp XlaBuilder::Conj(const XlaOp& operand) { - return Complex(Real(operand), Neg(Imag(operand))); -} - -XlaOp XlaBuilder::Sub(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kSubtract, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::Div(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kDivide, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::Rem(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kRemainder, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::Max(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kMaximum, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::Min(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kMinimum, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::And(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kAnd, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::Or(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kOr, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::Xor(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kXor, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::Not(const XlaOp& operand) { - return UnaryOp(HloOpcode::kNot, operand); -} - -XlaOp XlaBuilder::ShiftLeft(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kShiftLeft, lhs, rhs, broadcast_dimensions); -} - -XlaOp XlaBuilder::ShiftRightArithmetic( - const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kShiftRightArithmetic, lhs, rhs, - broadcast_dimensions); -} - -XlaOp XlaBuilder::ShiftRightLogical( - const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kShiftRightLogical, lhs, rhs, - broadcast_dimensions); -} - -XlaOp XlaBuilder::Abs(const XlaOp& operand) { - return UnaryOp(HloOpcode::kAbs, operand); -} - -XlaOp XlaBuilder::Atan2(const XlaOp& y, const XlaOp& x, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kAtan2, y, x, broadcast_dimensions); -} - -XlaOp XlaBuilder::Exp(const XlaOp& operand) { - return UnaryOp(HloOpcode::kExp, operand); -} - -XlaOp XlaBuilder::Expm1(const XlaOp& operand) { - return UnaryOp(HloOpcode::kExpm1, operand); -} - -XlaOp XlaBuilder::Floor(const XlaOp& operand) { - return UnaryOp(HloOpcode::kFloor, operand); -} - -XlaOp XlaBuilder::Ceil(const XlaOp& operand) { - return UnaryOp(HloOpcode::kCeil, operand); -} - -XlaOp XlaBuilder::Round(const XlaOp& operand) { - return UnaryOp(HloOpcode::kRoundNearestAfz, operand); -} - -XlaOp XlaBuilder::Log(const XlaOp& operand) { - return UnaryOp(HloOpcode::kLog, operand); -} - -XlaOp XlaBuilder::Log1p(const XlaOp& operand) { - return UnaryOp(HloOpcode::kLog1p, operand); -} - -XlaOp XlaBuilder::Sign(const XlaOp& operand) { - return UnaryOp(HloOpcode::kSign, operand); -} - -XlaOp XlaBuilder::Clz(const XlaOp& operand) { - return UnaryOp(HloOpcode::kClz, operand); -} - -XlaOp XlaBuilder::Cos(const XlaOp& operand) { - return UnaryOp(HloOpcode::kCos, operand); -} - -XlaOp XlaBuilder::Sin(const XlaOp& operand) { - return UnaryOp(HloOpcode::kSin, operand); -} - -XlaOp XlaBuilder::Tanh(const XlaOp& operand) { - return UnaryOp(HloOpcode::kTanh, operand); -} - -XlaOp XlaBuilder::Real(const XlaOp& operand) { - return UnaryOp(HloOpcode::kReal, operand); -} - -XlaOp XlaBuilder::Imag(const XlaOp& operand) { - return UnaryOp(HloOpcode::kImag, operand); -} - -XlaOp XlaBuilder::IsFinite(const XlaOp& operand) { - return UnaryOp(HloOpcode::kIsFinite, operand); -} - XlaOp XlaBuilder::Transpose(const XlaOp& operand, absl::Span permutation) { return ReportErrorOrReturn([&]() -> StatusOr { @@ -1596,36 +1551,144 @@ XlaOp XlaBuilder::Rev(const XlaOp& operand, }); } +namespace { +// Switch from a floating point value to a integer value in such a way that when +// using the integer value to compare, we get the same result for normal values, +// and -Nan is treated as the smallest value, and Nan is treated as the largest +// value. +// If f is a float, and +// x = bit_cast(f); +// y = x < 0 ? numeric_limits::max() - x : x; +// then y is ordered as an int32 such that finite values have the obvious order, +// -0 is ordered before 0, and -NaN and NaN appear at the beginning and end of +// the ordering. +// Note that in order to avoid -x to overflow, we calculate +// numeric_limits::max() - x as unsigned, and then convert back to +// signed. +XlaOp BitcastConvertFloatingPointToIntegral(const XlaOp& value, + int64 bit_width) { + PrimitiveType signed_type; + PrimitiveType unsigned_type; + XlaOp max_value; + switch (bit_width) { + case 16: + max_value = + ConstantR0(value.builder(), + static_cast(std::numeric_limits::max())); + signed_type = S16; + unsigned_type = U16; + break; + case 32: + max_value = + ConstantR0(value.builder(), + static_cast(std::numeric_limits::max())); + signed_type = S32; + unsigned_type = U32; + break; + case 64: + max_value = + ConstantR0(value.builder(), + static_cast(std::numeric_limits::max())); + signed_type = S64; + unsigned_type = U64; + break; + default: + return value.builder()->ReportError( + InvalidArgument("Invalid bit width %lld for Comparator floating " + "point parameter.", + bit_width)); + } + auto signed_value = BitcastConvertType(value, signed_type); + auto unsigned_value = BitcastConvertType(value, unsigned_type); + auto flipped_value = + BitcastConvertType(Sub(max_value, unsigned_value), signed_type); + auto is_negative = + Lt(signed_value, + ConstantLiteral(value.builder(), LiteralUtil::Zero(signed_type))); + return Select(is_negative, flipped_value, signed_value); +} +} // namespace + XlaOp XlaBuilder::Sort(const XlaOp& keys, absl::Span values, int64 dimension) { + return ReportErrorOrReturn([&]() -> StatusOr { + std::vector operands{keys}; + for (const XlaOp& value : values) { + operands.push_back(value); + } + // Build the default less-than comparator (copied from lib/comparators.cc). + // TODO(b/122298745): Remove the deprecated API method so that this code + // duplication can be deleted. + auto b = this->CreateSubBuilder("comparator"); + std::vector operand_types; + for (const XlaOp& operand : operands) { + TF_ASSIGN_OR_RETURN(auto operand_shape, GetShape(operand)); + operand_types.push_back(operand_shape.element_type()); + } + + int64 parameter_count = 0; + XlaOp first_lhs_param; + XlaOp first_rhs_param; + + for (auto operand_type : operand_types) { + auto scalar_shape = ShapeUtil::MakeShape(operand_type, {}); + auto lhs_param = + b->Parameter(parameter_count * 2, scalar_shape, + absl::StrCat("p.", parameter_count, ".lhs")); + auto rhs_param = + b->Parameter(parameter_count * 2 + 1, scalar_shape, + absl::StrCat("p.", parameter_count, ".rhs")); + if (parameter_count == 0) { + first_lhs_param = lhs_param; + first_rhs_param = rhs_param; + } + ++parameter_count; + } + if (primitive_util::IsFloatingPointType(operand_types[0])) { + PrimitiveType compare_type = operand_types[0]; + // Special-case handling for BF16. We currently do not support direct + // comparisons with BF16, so we convert to F32 and then use the F32 + // comparison logic. + if (compare_type == BF16) { + compare_type = F32; + first_lhs_param = b->ConvertElementType(first_lhs_param, F32); + first_rhs_param = b->ConvertElementType(first_rhs_param, F32); + } + int64 bit_width = primitive_util::BitWidth(compare_type); + first_lhs_param = + BitcastConvertFloatingPointToIntegral(first_lhs_param, bit_width); + first_rhs_param = + BitcastConvertFloatingPointToIntegral(first_rhs_param, bit_width); + } + Lt(first_lhs_param, first_rhs_param); + + TF_ASSIGN_OR_RETURN(auto comparator, b->Build()); + return Sort(operands, comparator, dimension); + }); +} + +XlaOp XlaBuilder::Sort(absl::Span operands, + const XlaComputation& comparator, int64 dimension) { return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; std::vector operand_shape_ptrs; - TF_ASSIGN_OR_RETURN(const Shape& keys_shape, GetShape(keys)); - operand_shape_ptrs.push_back(&keys_shape); - TF_ASSIGN_OR_RETURN(std::vector values_shapes, - GetOperandShapes(values)); - absl::c_transform(values_shapes, std::back_inserter(operand_shape_ptrs), + TF_ASSIGN_OR_RETURN(std::vector operand_shapes, + GetOperandShapes(operands)); + absl::c_transform(operand_shapes, std::back_inserter(operand_shape_ptrs), [](const Shape& shape) { return &shape; }); TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferVariadicOpShape( HloOpcode::kSort, operand_shape_ptrs)); *instr.mutable_shape() = shape.ToProto(); if (dimension == -1) { - TF_ASSIGN_OR_RETURN(const Shape& keys_shape, GetShape(keys)); + TF_ASSIGN_OR_RETURN(const Shape& keys_shape, GetShape(operands[0])); dimension = keys_shape.rank() - 1; } instr.add_dimensions(dimension); - std::vector operands{keys}; - operands.insert(operands.end(), values.begin(), values.end()); + AddCalledComputation(comparator, &instr); return AddInstruction(std::move(instr), HloOpcode::kSort, operands); }); } -XlaOp XlaBuilder::Pow(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions) { - return BinaryOp(HloOpcode::kPower, lhs, rhs, broadcast_dimensions); -} - XlaOp XlaBuilder::ConvertElementType(const XlaOp& operand, PrimitiveType new_element_type) { return ReportErrorOrReturn([&]() -> StatusOr { @@ -1651,10 +1714,6 @@ XlaOp XlaBuilder::BitcastConvertType(const XlaOp& operand, }); } -XlaOp XlaBuilder::Neg(const XlaOp& operand) { - return UnaryOp(HloOpcode::kNegate, operand); -} - XlaOp XlaBuilder::Clamp(const XlaOp& min, const XlaOp& operand, const XlaOp& max) { return TernaryOp(HloOpcode::kClamp, min, operand, max); @@ -2035,8 +2094,8 @@ XlaOp XlaBuilder::CrossReplicaSum( TF_ASSIGN_OR_RETURN(const Shape& shape, GetShape(operand)); const Shape& scalar_shape = ShapeUtil::MakeShape(shape.element_type(), {}); auto b = CreateSubBuilder("sum"); - b->Add(b->Parameter(/*parameter_number=*/0, scalar_shape, "x"), - b->Parameter(/*parameter_number=*/1, scalar_shape, "y")); + Add(b->Parameter(/*parameter_number=*/0, scalar_shape, "x"), + b->Parameter(/*parameter_number=*/1, scalar_shape, "y")); TF_ASSIGN_OR_RETURN(auto computation, b->Build()); return CrossReplicaSum(operand, computation, replica_groups, /*channel_id=*/absl::nullopt); @@ -2145,6 +2204,14 @@ XlaOp XlaBuilder::CollectivePermute( }); } +XlaOp XlaBuilder::ReplicaId() { + return ReportErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + *instr.mutable_shape() = ShapeUtil::MakeShape(U32, {}).ToProto(); + return AddInstruction(std::move(instr), HloOpcode::kReplicaId, {}); + }); +} + XlaOp XlaBuilder::SelectAndScatter(const XlaOp& operand, const XlaComputation& select, absl::Span window_dimensions, @@ -2415,7 +2482,7 @@ StatusOr XlaBuilder::IsConstant(const XlaOp& operand) const { TF_RETURN_IF_ERROR(LookUpInstruction(operand).status()); bool is_constant = true; - std::set visited; + absl::flat_hash_set visited; IsConstantVisitor(operand.handle(), &visited, &is_constant); return is_constant; } @@ -2462,21 +2529,58 @@ StatusOr XlaBuilder::BuildConstantSubGraph( worklist.pop(); TF_ASSIGN_OR_RETURN(const HloInstructionProto* instr_proto, LookUpInstructionByHandle(handle)); - for (int64 id : instr_proto->operand_ids()) { - if (related_ops.insert(id).second) { - worklist.push(id); + + if (instr_proto->opcode() == + HloOpcodeString(HloOpcode::kGetDimensionSize)) { + // At this point, BuildConstantSubGraph should never encounter a + // GetDimensionSize with a dynamic dimension. IsConstant check would have + // failed at the beginning of this function. + // + // Replace GetDimensionSize with a Constant representing the static bound + // of the shape. + int64 dimension = instr_proto->dimensions(0); + int64 operand_handle = instr_proto->operand_ids(0); + TF_ASSIGN_OR_RETURN(const HloInstructionProto* operand_proto, + LookUpInstructionByHandle(operand_handle)); + + TF_RET_CHECK(!operand_proto->shape().is_dynamic_dimension(dimension)); + auto constant_dimension_size = + static_cast(operand_proto->shape().dimensions(dimension)); + + Literal literal = LiteralUtil::CreateR0(constant_dimension_size); + + HloInstructionProto const_instr; + *const_instr.mutable_shape() = literal.shape().ToProto(); + *const_instr.mutable_literal() = literal.ToProto(); + *const_instr.mutable_opcode() = HloOpcodeString(HloOpcode::kConstant); + + const_instr.set_id(handle); + *const_instr.mutable_name() = + GetFullName(const_instr.opcode(), kNameSeparator, const_instr.id()); + *entry.add_instructions() = + const_instr; // Add to the result constant graph. + } else { + for (int64 id : instr_proto->operand_ids()) { + if (related_ops.insert(id).second) { + worklist.push(id); + } + } + for (int64 called_id : instr_proto->called_computation_ids()) { + related_calls.insert(called_id); } - } - for (int64 called_id : instr_proto->called_computation_ids()) { - related_calls.insert(called_id); } } // Add related ops to the computation. for (int64 id : related_ops) { - auto* instr = entry.add_instructions(); TF_ASSIGN_OR_RETURN(const HloInstructionProto* instr_src, LookUpInstructionByHandle(id)); + + if (instr_src->opcode() == HloOpcodeString(HloOpcode::kGetDimensionSize)) { + continue; + } + auto* instr = entry.add_instructions(); + *instr = *instr_src; // Ensures that the instruction names are unique among the graph. const string& new_name = @@ -2749,12 +2853,21 @@ XlaOp DynamicSlice(const XlaOp& operand, const XlaOp& start_indices, absl::Span slice_sizes) { return operand.builder()->DynamicSlice(operand, start_indices, slice_sizes); } +XlaOp DynamicSlice(const XlaOp& operand, absl::Span start_indices, + absl::Span slice_sizes) { + return operand.builder()->DynamicSlice(operand, start_indices, slice_sizes); +} XlaOp DynamicUpdateSlice(const XlaOp& operand, const XlaOp& update, const XlaOp& start_indices) { return operand.builder()->DynamicUpdateSlice(operand, update, start_indices); } +XlaOp DynamicUpdateSlice(const XlaOp& operand, const XlaOp& update, + absl::Span start_indices) { + return operand.builder()->DynamicUpdateSlice(operand, update, start_indices); +} + XlaOp ConcatInDim(XlaBuilder* builder, absl::Span operands, int64 dimension) { return builder->ConcatInDim(operands, dimension); @@ -2778,32 +2891,38 @@ XlaOp GetTupleElement(const XlaOp& tuple_data, int64 index) { XlaOp Eq(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Eq(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kEq, lhs, rhs, + broadcast_dimensions); } XlaOp Ne(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Ne(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kNe, lhs, rhs, + broadcast_dimensions); } XlaOp Ge(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Ge(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kGe, lhs, rhs, + broadcast_dimensions); } XlaOp Gt(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Gt(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kGt, lhs, rhs, + broadcast_dimensions); } -XlaOp Lt(const XlaOp& lhs, const XlaOp& rhs, +XlaOp Le(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Lt(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kLe, lhs, rhs, + broadcast_dimensions); } -XlaOp Le(const XlaOp& lhs, const XlaOp& rhs, +XlaOp Lt(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Le(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kLt, lhs, rhs, + broadcast_dimensions); } XlaOp Dot(const XlaOp& lhs, const XlaOp& rhs, @@ -2877,6 +2996,29 @@ XlaOp Fft(const XlaOp& operand, FftType fft_type, return operand.builder()->Fft(operand, fft_type, fft_length); } +XlaOp TriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower, + bool unit_diagonal, + TriangularSolveOptions::Transpose transpose_a) { + XlaBuilder* builder = a.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + TF_ASSIGN_OR_RETURN(const Shape& a_shape, builder->GetShape(a)); + TF_ASSIGN_OR_RETURN(const Shape& b_shape, builder->GetShape(b)); + xla::TriangularSolveOptions& options = + *instr.mutable_triangular_solve_options(); + options.set_left_side(left_side); + options.set_lower(lower); + options.set_unit_diagonal(unit_diagonal); + options.set_transpose_a(transpose_a); + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferTriangularSolveShape( + a_shape, b_shape, options)); + *instr.mutable_shape() = shape.ToProto(); + + return builder->AddInstruction(std::move(instr), + HloOpcode::kTriangularSolve, {a, b}); + }); +} + XlaOp Infeed(XlaBuilder* builder, const Shape& shape, const string& config) { return builder->Infeed(shape, config); } @@ -2906,78 +3048,96 @@ XlaOp CustomCallWithLayout(XlaBuilder* builder, const string& call_target_name, operand_shapes_with_layout); } -XlaOp Complex(const XlaOp& real, const XlaOp& imag, +XlaOp Complex(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return real.builder()->Complex(real, imag, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kComplex, lhs, rhs, + broadcast_dimensions); } -XlaOp Conj(const XlaOp& operand) { return operand.builder()->Conj(operand); } +XlaOp Conj(const XlaOp& operand) { + return Complex(Real(operand), Neg(Imag(operand))); +} XlaOp Add(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Add(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kAdd, lhs, rhs, + broadcast_dimensions); } XlaOp Sub(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Sub(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kSubtract, lhs, rhs, + broadcast_dimensions); } XlaOp Mul(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Mul(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kMultiply, lhs, rhs, + broadcast_dimensions); } XlaOp Div(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Div(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kDivide, lhs, rhs, + broadcast_dimensions); } XlaOp Rem(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Rem(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kRemainder, lhs, rhs, + broadcast_dimensions); } XlaOp Max(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Max(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kMaximum, lhs, rhs, + broadcast_dimensions); } XlaOp Min(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Min(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kMinimum, lhs, rhs, + broadcast_dimensions); } XlaOp And(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->And(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kAnd, lhs, rhs, + broadcast_dimensions); } XlaOp Or(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Or(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kOr, lhs, rhs, + broadcast_dimensions); } XlaOp Xor(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Xor(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kXor, lhs, rhs, + broadcast_dimensions); } -XlaOp Not(const XlaOp& operand) { return operand.builder()->Not(operand); } +XlaOp Not(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kNot, operand); +} XlaOp ShiftLeft(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->ShiftLeft(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kShiftLeft, lhs, rhs, + broadcast_dimensions); } XlaOp ShiftRightArithmetic(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->ShiftRightArithmetic(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kShiftRightArithmetic, lhs, rhs, + broadcast_dimensions); } XlaOp ShiftRightLogical(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->ShiftRightLogical(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kShiftRightLogical, lhs, rhs, + broadcast_dimensions); } XlaOp Reduce(const XlaOp& operand, const XlaOp& init_value, @@ -3049,6 +3209,8 @@ XlaOp CollectivePermute( return operand.builder()->CollectivePermute(operand, source_target_pairs); } +XlaOp ReplicaId(XlaBuilder* builder) { return builder->ReplicaId(); } + XlaOp SelectAndScatter(const XlaOp& operand, const XlaComputation& select, absl::Span window_dimensions, absl::Span window_strides, Padding padding, @@ -3070,48 +3232,67 @@ XlaOp SelectAndScatterWithGeneralPadding( init_value, scatter); } -XlaOp Abs(const XlaOp& operand) { return operand.builder()->Abs(operand); } +XlaOp Abs(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kAbs, operand); +} -XlaOp Atan2(const XlaOp& y, const XlaOp& x, +XlaOp Atan2(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return y.builder()->Atan2(y, x, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kAtan2, lhs, rhs, + broadcast_dimensions); } -XlaOp Exp(const XlaOp& operand) { return operand.builder()->Exp(operand); } - -XlaOp Expm1(const XlaOp& operand) { return operand.builder()->Expm1(operand); } - -XlaOp Floor(const XlaOp& operand) { return operand.builder()->Floor(operand); } - -XlaOp Ceil(const XlaOp& operand) { return operand.builder()->Ceil(operand); } - -XlaOp Round(const XlaOp& operand) { return operand.builder()->Round(operand); } - -XlaOp Log(const XlaOp& operand) { return operand.builder()->Log(operand); } - -XlaOp Log1p(const XlaOp& operand) { return operand.builder()->Log1p(operand); } - -XlaOp Sign(const XlaOp& operand) { return operand.builder()->Sign(operand); } - -XlaOp Clz(const XlaOp& operand) { return operand.builder()->Clz(operand); } - -XlaOp Cos(const XlaOp& operand) { return operand.builder()->Cos(operand); } - -XlaOp Sin(const XlaOp& operand) { return operand.builder()->Sin(operand); } - -XlaOp Tanh(const XlaOp& operand) { return operand.builder()->Tanh(operand); } - -XlaOp Real(const XlaOp& operand) { return operand.builder()->Real(operand); } - -XlaOp Imag(const XlaOp& operand) { return operand.builder()->Imag(operand); } +XlaOp Exp(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kExp, operand); +} +XlaOp Expm1(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kExpm1, operand); +} +XlaOp Floor(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kFloor, operand); +} +XlaOp Ceil(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kCeil, operand); +} +XlaOp Round(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kRoundNearestAfz, operand); +} +XlaOp Log(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kLog, operand); +} +XlaOp Log1p(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kLog1p, operand); +} +XlaOp Sign(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kSign, operand); +} +XlaOp Clz(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kClz, operand); +} +XlaOp Cos(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kCos, operand); +} +XlaOp Sin(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kSin, operand); +} +XlaOp Tanh(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kTanh, operand); +} +XlaOp Real(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kReal, operand); +} +XlaOp Imag(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kImag, operand); +} XlaOp Pow(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions) { - return lhs.builder()->Pow(lhs, rhs, broadcast_dimensions); + return lhs.builder()->BinaryOp(HloOpcode::kPower, lhs, rhs, + broadcast_dimensions); } XlaOp IsFinite(const XlaOp& operand) { - return operand.builder()->IsFinite(operand); + return operand.builder()->UnaryOp(HloOpcode::kIsFinite, operand); } XlaOp ConvertElementType(const XlaOp& operand, PrimitiveType new_element_type) { @@ -3122,7 +3303,9 @@ XlaOp BitcastConvertType(const XlaOp& operand, PrimitiveType new_element_type) { return operand.builder()->BitcastConvertType(operand, new_element_type); } -XlaOp Neg(const XlaOp& operand) { return operand.builder()->Neg(operand); } +XlaOp Neg(const XlaOp& operand) { + return operand.builder()->UnaryOp(HloOpcode::kNegate, operand); +} XlaOp Transpose(const XlaOp& operand, absl::Span permutation) { return operand.builder()->Transpose(operand, permutation); @@ -3136,6 +3319,11 @@ XlaOp Sort(const XlaOp& keys, absl::Span values, int64 dimension) { return keys.builder()->Sort(keys, values, dimension); } +XlaOp Sort(absl::Span operands, const XlaComputation& comparator, + int64 dimension) { + return operands[0].builder()->Sort(operands, comparator, dimension); +} + XlaOp Clamp(const XlaOp& min, const XlaOp& operand, const XlaOp& max) { return min.builder()->Clamp(min, operand, max); } diff --git a/tensorflow/compiler/xla/client/xla_builder.h b/tensorflow/compiler/xla/client/xla_builder.h index 68ddf2cdd8992156f3792097d53f8d57323cfb9a..1e39c8766f318f7f31778265dddae6b2b32e111d 100644 --- a/tensorflow/compiler/xla/client/xla_builder.h +++ b/tensorflow/compiler/xla/client/xla_builder.h @@ -56,6 +56,9 @@ class XlaOp { } ~XlaOp() = default; + XlaOp(const XlaOp& other) = default; + XlaOp& operator=(const XlaOp& other) = default; + // Precondition: !IsUninitialized(). // // It's very common to do foo.builder()->bar(). Without this precondition, if @@ -197,11 +200,19 @@ class XlaBuilder { // status. Note that all ops that have been enqueued will be moved to the // computation being returned. The root of the computation will be the last // added operation. - StatusOr Build(); + // + // `remove_dynamic_dimensions` tells the builder whether to remove the + // dyanmic dimensions information in all ops. + // + // TODO(b/121223198): Delete `remove_dynamic_dimensions` and keeps the + // dynamic dimensions information when XLA backend can handle dynamic + // dimensions. + StatusOr Build(bool remove_dynamic_dimensions = true); // Overload of Build which specifies a particular root instruction for the // computation. - StatusOr Build(XlaOp root); + StatusOr Build(XlaOp root, + bool remove_dynamic_dimensions = true); // Builds the computation with the requested operations, or notes an error in // the parent XlaBuilder and returns an empty computation if building failed. @@ -269,6 +280,10 @@ class XlaBuilder { // and its real dynamic size is represented by `dynamic_param_index` in // parameter `dynamic_param_num`. // + // Note that this should be called before the dynamic parameters are used to + // create other operations, otherwise created operations won't have the + // dynamic dimensions information. + // // TODO(b/119520625): Remove this API once we have more dynamic shape infra // ready. Status SetDynamicBinding(int64 dynamic_size_param_num, @@ -293,7 +308,7 @@ class XlaBuilder { }; // Build helper which takes the id of the root operation.. - StatusOr Build(int64 root_id); + StatusOr Build(int64 root_id, bool remove_dynamic_dimensions); // Description for the methods below can be found in the corresponding public // functions section in this file. @@ -303,38 +318,6 @@ class XlaBuilder { XlaOp ConstantLiteral(const LiteralSlice& literal); - template - XlaOp ConstantR0(NativeT value); - template - XlaOp ConstantR1(absl::Span values); - XlaOp ConstantR1(const tensorflow::core::Bitmap& values); - template - XlaOp ConstantR2( - std::initializer_list> values); - template - XlaOp ConstantFromArrayWithLayout(const Array& values, - const Layout& layout); - template - XlaOp ConstantFromArray(const Array& values); - template - XlaOp ConstantR2FromArray2DWithLayout(const Array2D& values, - const Layout& layout); - template - XlaOp ConstantR2FromArray2D(const Array2D& values); - template - XlaOp ConstantR3FromArray3DWithLayout(const Array3D& values, - const Layout& layout); - template - XlaOp ConstantR3FromArray3D(const Array3D& values); - template - XlaOp ConstantR4FromArray4DWithLayout(const Array4D& values, - const Layout& layout); - template - XlaOp ConstantR4FromArray4D(const Array4D& values); - - template - XlaOp ConstantR1(int64 length, NativeT value); - XlaOp Broadcast(const XlaOp& operand, absl::Span broadcast_sizes); @@ -359,11 +342,18 @@ class XlaBuilder { XlaOp SliceInDim(const XlaOp& operand, int64 start_index, int64 limit_index, int64 stride, int64 dimno); + ABSL_DEPRECATED("Use span-of-indices form instead") XlaOp DynamicSlice(const XlaOp& operand, const XlaOp& start_indices, absl::Span slice_sizes); + XlaOp DynamicSlice(const XlaOp& operand, + absl::Span start_indices, + absl::Span slice_sizes); + ABSL_DEPRECATED("Use span-of-indices form instead") XlaOp DynamicUpdateSlice(const XlaOp& operand, const XlaOp& update, const XlaOp& start_indices); + XlaOp DynamicUpdateSlice(const XlaOp& operand, const XlaOp& update, + absl::Span start_indices); XlaOp ConcatInDim(absl::Span operands, int64 dimension); @@ -375,24 +365,6 @@ class XlaBuilder { XlaOp GetTupleElement(const XlaOp& tuple_data, int64 index); - XlaOp Eq(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Ne(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Ge(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Gt(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Lt(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Le(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - XlaOp Dot(const XlaOp& lhs, const XlaOp& rhs, const PrecisionConfig* precision_config = nullptr); @@ -457,50 +429,6 @@ class XlaBuilder { const Shape& shape_with_layout, const string& opaque, absl::optional> operand_shapes_with_layout); - XlaOp Complex(const XlaOp& real, const XlaOp& imag, - absl::Span broadcast_dimensions = {}); - - XlaOp Conj(const XlaOp& operand); - - XlaOp Add(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Sub(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Mul(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Div(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Rem(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Max(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Min(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp And(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Or(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Xor(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp Not(const XlaOp& operand); - - XlaOp ShiftLeft(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - XlaOp ShiftRightArithmetic(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - XlaOp ShiftRightLogical(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - XlaOp Reduce(const XlaOp& operand, const XlaOp& init_value, const XlaComputation& computation, absl::Span dimensions_to_reduce); @@ -543,6 +471,8 @@ class XlaBuilder { const XlaOp& operand, const std::vector>& source_target_pairs); + XlaOp ReplicaId(); + XlaOp SelectAndScatter(const XlaOp& operand, const XlaComputation& select, absl::Span window_dimensions, absl::Span window_strides, @@ -557,44 +487,6 @@ class XlaBuilder { absl::Span> padding, const XlaOp& source, const XlaOp& init_value, const XlaComputation& scatter); - XlaOp Abs(const XlaOp& operand); - - XlaOp Atan2(const XlaOp& y, const XlaOp& x, - absl::Span broadcast_dimensions = {}); - - XlaOp Exp(const XlaOp& operand); - - XlaOp Expm1(const XlaOp& operand); - - XlaOp Floor(const XlaOp& operand); - - XlaOp Ceil(const XlaOp& operand); - - XlaOp Round(const XlaOp& operand); - - XlaOp Log(const XlaOp& operand); - - XlaOp Log1p(const XlaOp& operand); - - XlaOp Sign(const XlaOp& operand); - - XlaOp Clz(const XlaOp& operand); - - XlaOp Cos(const XlaOp& operand); - - XlaOp Sin(const XlaOp& operand); - - XlaOp Tanh(const XlaOp& operand); - - XlaOp Real(const XlaOp& operand); - - XlaOp Imag(const XlaOp& operand); - - XlaOp Pow(const XlaOp& lhs, const XlaOp& rhs, - absl::Span broadcast_dimensions = {}); - - XlaOp IsFinite(const XlaOp& operand); - XlaOp Iota(const Shape& shape, int64 iota_dimension); XlaOp Iota(PrimitiveType type, int64 size); @@ -605,14 +497,15 @@ class XlaBuilder { XlaOp BitcastConvertType(const XlaOp& operand, PrimitiveType new_element_type); - XlaOp Neg(const XlaOp& operand); - XlaOp Transpose(const XlaOp& operand, absl::Span permutation); XlaOp Rev(const XlaOp& operand, absl::Span dimensions); + ABSL_DEPRECATED("Use form with comparator computation instead") XlaOp Sort(const XlaOp& keys, absl::Span values = {}, int64 dimension = -1); + XlaOp Sort(absl::Span operands, const XlaComputation& comparator, + int64 dimension = -1); XlaOp Clamp(const XlaOp& min, const XlaOp& operand, const XlaOp& max); @@ -727,7 +620,8 @@ class XlaBuilder { // operation such as `RngNormal` or `Infeed`. The visitor walks the // computation starting at a given operation and sets is_constant to false iff // a parameter or stateful operation is encountered. - void IsConstantVisitor(const int64 op_handle, std::set* visited, + void IsConstantVisitor(const int64 op_handle, + absl::flat_hash_set* visited, bool* is_constant) const; // Checks bounds for convolution parameters. @@ -803,48 +697,6 @@ class XlaBuilder { const Shape& shape, const string& name); friend XlaOp ConstantLiteral(XlaBuilder* builder, const LiteralSlice& literal); - template - friend XlaOp ConstantR0(XlaBuilder* builder, NativeT value); - template - friend XlaOp ConstantR1(XlaBuilder* builder, - absl::Span values); - friend XlaOp ConstantR1(XlaBuilder* builder, - const tensorflow::core::Bitmap& values); - template - friend XlaOp ConstantR2( - XlaBuilder* builder, - std::initializer_list> values); - template - friend XlaOp ConstantFromArrayWithLayout(XlaBuilder* builder, - const Array& values, - const Layout& layout); - template - friend XlaOp ConstantFromArray(XlaBuilder* builder, - const Array& values); - template - friend XlaOp ConstantR2FromArray2DWithLayout(XlaBuilder* builder, - const Array2D& values, - const Layout& layout); - template - friend XlaOp ConstantR2FromArray2D(XlaBuilder* builder, - const Array2D& values); - template - friend XlaOp ConstantR3FromArray3DWithLayout(XlaBuilder* builder, - const Array3D& values, - const Layout& layout); - template - friend XlaOp ConstantR3FromArray3D(XlaBuilder* builder, - const Array3D& values); - template - friend XlaOp ConstantR4FromArray4DWithLayout(XlaBuilder* builder, - const Array4D& values, - const Layout& layout); - template - friend XlaOp ConstantR4FromArray4D(XlaBuilder* builder, - const Array4D& values); - - template - friend XlaOp ConstantR1(XlaBuilder* builder, int64 length, NativeT value); friend XlaOp Broadcast(const XlaOp& operand, absl::Span broadcast_sizes); @@ -874,9 +726,14 @@ class XlaBuilder { friend XlaOp DynamicSlice(const XlaOp& operand, const XlaOp& start_indices, absl::Span slice_sizes); + friend XlaOp DynamicSlice(const XlaOp& operand, + absl::Span start_indices, + absl::Span slice_sizes); friend XlaOp DynamicUpdateSlice(const XlaOp& operand, const XlaOp& update, const XlaOp& start_indices); + friend XlaOp DynamicUpdateSlice(const XlaOp& operand, const XlaOp& update, + absl::Span start_indices); friend XlaOp ConcatInDim(XlaBuilder* builder, absl::Span operands, int64 dimension); @@ -937,6 +794,9 @@ class XlaBuilder { const PrecisionConfig* precision_config); friend XlaOp Fft(const XlaOp& operand, FftType fft_type, absl::Span fft_length); + friend XlaOp TriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower, + bool unit_diagonal, + TriangularSolveOptions::Transpose transpose_a); friend XlaOp Infeed(XlaBuilder* builder, const Shape& shape, const string& config); friend void Outfeed(const XlaOp& operand, const Shape& shape_with_layout, @@ -1015,6 +875,7 @@ class XlaBuilder { friend XlaOp CollectivePermute( const XlaOp& operand, const std::vector>& source_target_pairs); + friend XlaOp ReplicaId(XlaBuilder* builder); friend XlaOp SelectAndScatter(const XlaOp& operand, const XlaComputation& select, absl::Span window_dimensions, @@ -1061,6 +922,8 @@ class XlaBuilder { friend XlaOp Rev(const XlaOp& operand, absl::Span dimensions); friend XlaOp Sort(const XlaOp& keys, absl::Span values, int64 dimension); + friend XlaOp Sort(absl::Span operands, + const XlaComputation& comparator, int64 dimension); friend XlaOp Clamp(const XlaOp& min, const XlaOp& operand, const XlaOp& max); friend XlaOp Map(XlaBuilder* builder, absl::Span operands, const XlaComputation& computation, @@ -1318,10 +1181,15 @@ XlaOp SliceInDim(const XlaOp& operand, int64 start_index, int64 limit_index, // The size of the slice in each dimension is passed in 'slice_sizes', // which specify the end point of exclusive slice intervals in each // dimension [start, start + size). -// The shape of 'start_indices' must be rank == 1, with dimension size -// equal to the rank of the 'operand'. +// The shape of each element of 'start_indices' must be scalar, with the span +// size equal to the rank of the 'operand'. All elements of 'start_indices' must +// have the same shape. // Slice index calculations are computed modulo input dimension sizes to // prevent dynamic start indices from generating out-of-bound array accesses. +XlaOp DynamicSlice(const XlaOp& operand, absl::Span start_indices, + absl::Span slice_sizes); + +ABSL_DEPRECATED("Use span-of-indices form instead") XlaOp DynamicSlice(const XlaOp& operand, const XlaOp& start_indices, absl::Span slice_sizes); @@ -1337,10 +1205,15 @@ XlaOp DynamicSlice(const XlaOp& operand, const XlaOp& start_indices, // [4 5 6] => DynamicUpdateslice(data, update, start) => [4 10 11] // [7 8 9] [7 8 9 ] // -// The shape of 'start_indices' must be rank == 1, with dimension size -// equal to the rank of the 'operand'. +// The shape of each element of 'start_indices' must be scalar, with the span +// size equal to the rank of the 'operand'. All elements of 'start_indices' must +// have the same shape. // Slice index calculations are computed modulo update dimension sizes to // prevent dynamic start indices from generating out-of-bound array accesses. +XlaOp DynamicUpdateSlice(const XlaOp& operand, const XlaOp& update, + absl::Span start_indices); + +ABSL_DEPRECATED("Use span-of-indices form instead") XlaOp DynamicUpdateSlice(const XlaOp& operand, const XlaOp& update, const XlaOp& start_indices); @@ -1446,6 +1319,32 @@ XlaOp ConvGeneralDilated(const XlaOp& lhs, const XlaOp& rhs, XlaOp Fft(const XlaOp& operand, FftType fft_type, absl::Span fft_length); +// Solves systems of linear equations with lower or upper triangular coefficient +// matrices by forward- or back-substitution. Broadcasting along leading +// dimensions, this routine solves for x in 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))`. +// +// * `a` is a tensor of shape `[..., M, M]` whose innermost 2 dimensions form +// square matrices. If `lower` is true (false), then the strictly upper +// (lower) triangular part of each innermost matrix in `a` is assumed to be +// zero and is not accessed. +// * `b` is a tensor of shape `[..., M, K]` if `left_side` is true, otherwise a +// tensor of shape `[..., K, M]`. +// * `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). +// * 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 +// `a`: the identity function, transpose(a), or conjugate(transpose(a)) +XlaOp TriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower, + bool unit_diagonal, + TriangularSolveOptions::Transpose transpose_a); + // 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, @@ -1545,9 +1444,33 @@ XlaOp Min(const XlaOp& lhs, const XlaOp& rhs, XlaOp And(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); +// Overload to call And with 3 or more operands. We need the following somewhat +// convoluted overload set to disambiguate with the overload that takes the +// `broadcast_dimensions` optional param. +inline XlaOp And(const XlaOp& op1, const XlaOp& op2, const XlaOp& op3) { + return And(op1, And(op2, op3)); +} +template +XlaOp And(const XlaOp& op1, const XlaOp& op2, const XlaOp& op3, + const XlaOpTs&... operands) { + return And(op1, And(op2, And(op3, operands...))); +} + XlaOp Or(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); +// Overload to call Or with 3 or more operands. As with `And`, we need the +// following complicated overload set to handle the default arg in the `Or` +// overload above. +inline XlaOp Or(const XlaOp& op1, const XlaOp& op2, const XlaOp& op3) { + return Or(op1, Or(op2, op3)); +} +template +XlaOp Or(const XlaOp& op1, const XlaOp& op2, const XlaOp& op3, + const XlaOpTs&... operands) { + return Or(op1, Or(op2, Or(op3, operands...))); +} + XlaOp Xor(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); @@ -1640,6 +1563,9 @@ XlaOp CollectivePermute( const XlaOp& operand, const std::vector>& source_target_pairs); +// Enqueues an operation that returns the replica ID. +XlaOp ReplicaId(XlaBuilder* builder); + // Enqueues an operation that scatters the `source` array to the selected // indices of each window. XlaOp SelectAndScatter(const XlaOp& operand, const XlaComputation& select, @@ -1711,10 +1637,14 @@ XlaOp Imag(const XlaOp& operand); XlaOp Pow(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); -// Enqueues an operator that tests if the operand's values are finite, i.e., -// not Inf or NaN. Defined only for floating-point types. Returns an array of -// booleans with the same shape where entries are true iff the corresponding -// entry was NaN. +// Enqueues an operator that tests if the operand's values are finite, i.e., not +// +/-Inf or NaN. Returns an array of booleans with the same shape where +// entries are true iff the corresponding entry was not infinite or NaN. +// +// Defined only for real-valued (i.e. not complex) floating-point types; raises +// an error for other types. +// +// See also IsInf, IsPosInf, IsNegInf, and IsNan in lib/math.h. XlaOp IsFinite(const XlaOp& operand); // Enqueues an iota operation onto the computation. @@ -1750,7 +1680,7 @@ XlaOp Rev(const XlaOp& operand, absl::Span dimensions); // of keys, in ascending order. // * If the keys have higher rank, the keys are sorted along the provided // dimension. For example, for a rank-2 tensor (a matrix) of keys, a dimension -// value of 0 will indepenently sort every column, and a dimension value of 1 +// value of 0 will independently sort every column, and a dimension value of 1 // will independently sort each row. If no dimension number is provided, then // the last dimension is chosen by default. // @@ -1760,9 +1690,36 @@ XlaOp Rev(const XlaOp& operand, absl::Span dimensions); // * The result is a tuple that consists of a sorted tensor of keys (along the // provided dimension, as above) as the first element, and tensors with their // corresponding values as the other elements. +ABSL_DEPRECATED("Use form with comparator computation instead") XlaOp Sort(const XlaOp& keys, absl::Span values = {}, int64 dimension = -1); +// Enqueues a sort instruction onto the computation, using 'comparator' for +// comparisons. 'comparator' needs to define a strict weak order. +// If only one operand is provided: +// * If the operand is a rank-1 tensor (an array), the result is a sorted array. +// The resulting sorting order has the property that for all index positions +// i, j with i < j, either +// comparator(value[i], value[j]) = comparator(value[j], value[i]) = false or +// comparator(value[i], value[j]) = true. +// * If the operand has higher rank, the operand is sorted along the provided +// dimension. For example, for a rank-2 tensor (a matrix), a dimension value +// of 0 will independently sort every column, and a dimension value of 1 will +// independently sort each row. If no dimension number is provided, then the +// last dimension is chosen by default. For the dimension which is sorted, the +// same sorting order applies as in the rank-1 case. +// +// If more than one operand is provided: +// * 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. +// Default comparator computations can be found in lib/comparators.h +XlaOp Sort(absl::Span operands, const XlaComputation& comparator, + int64 dimension = -1); + // Enqueues a clamp instruction onto the computation. XlaOp Clamp(const XlaOp& min, const XlaOp& operand, const XlaOp& max); @@ -1901,81 +1858,6 @@ XlaOp GetDimensionSize(const XlaOp& operand, int64 dimension); // Implementation details below this point. // -template -XlaOp XlaBuilder::ConstantR0(NativeT value) { - return ConstantLiteral(LiteralUtil::CreateR0(value)); -} - -template -XlaOp XlaBuilder::ConstantR1(absl::Span values) { - return ConstantLiteral(LiteralUtil::CreateR1(values)); -} - -template -XlaOp XlaBuilder::ConstantR1(int64 length, NativeT value) { - Literal literal(ShapeUtil::MakeShape( - primitive_util::NativeToPrimitiveType(), {length})); - literal.PopulateWithValue(value); - return ConstantLiteral(literal); -} - -inline XlaOp XlaBuilder::ConstantR1(const tensorflow::core::Bitmap& values) { - return ConstantLiteral(LiteralUtil::CreateR1(values)); -} - -template -XlaOp XlaBuilder::ConstantR2( - std::initializer_list> values) { - return ConstantLiteral(LiteralUtil::CreateR2(values)); -} - -template -XlaOp XlaBuilder::ConstantFromArrayWithLayout(const Array& values, - const Layout& layout) { - return ConstantLiteral( - LiteralUtil::CreateFromArrayWithLayout(values, layout)); -} - -template -XlaOp XlaBuilder::ConstantFromArray(const Array& values) { - return ConstantLiteral(LiteralUtil::CreateFromArray(values)); -} - -template -XlaOp XlaBuilder::ConstantR2FromArray2DWithLayout( - const Array2D& values, const Layout& layout) { - return ConstantLiteral( - LiteralUtil::CreateFromArrayWithLayout(values, layout)); -} - -template -XlaOp XlaBuilder::ConstantR2FromArray2D(const Array2D& values) { - return ConstantLiteral(LiteralUtil::CreateR2FromArray2D(values)); -} - -template -XlaOp XlaBuilder::ConstantR3FromArray3DWithLayout( - const Array3D& values, const Layout& layout) { - return ConstantLiteral( - LiteralUtil::CreateR3FromArray3DWithLayout(values, layout)); -} - -template -XlaOp XlaBuilder::ConstantR3FromArray3D(const Array3D& values) { - return ConstantFromArray(values); -} - -template -XlaOp XlaBuilder::ConstantR4FromArray4DWithLayout( - const Array4D& values, const Layout& layout) { - return ConstantFromArrayWithLayout(values, layout); -} - -template -XlaOp XlaBuilder::ConstantR4FromArray4D(const Array4D& values) { - return ConstantFromArray(values); -} - // Free function template implementations. template diff --git a/tensorflow/compiler/xla/client/xla_builder_test.cc b/tensorflow/compiler/xla/client/xla_builder_test.cc index ba929a1200915a9dfc3ca1f67dd89bb13dce7c61..c9fa738a19d0928d56ac4b98beb5fc0ed195518b 100644 --- a/tensorflow/compiler/xla/client/xla_builder_test.cc +++ b/tensorflow/compiler/xla/client/xla_builder_test.cc @@ -25,6 +25,7 @@ limitations under the License. #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" namespace xla { @@ -39,7 +40,8 @@ using ::testing::HasSubstr; class XlaBuilderTest : public ::testing::Test { protected: StatusOr> BuildHloModule(XlaBuilder* b) { - TF_ASSIGN_OR_RETURN(XlaComputation computation, b->Build()); + TF_ASSIGN_OR_RETURN(XlaComputation computation, + b->Build(/*remove_dynamic_dimensions=*/false)); const HloModuleProto& proto = computation.proto(); TF_ASSIGN_OR_RETURN(const auto& config, HloModule::CreateModuleConfigFromProto( @@ -50,7 +52,8 @@ class XlaBuilderTest : public ::testing::Test { // Overload which explicitly specifies the root instruction. StatusOr> BuildHloModule(XlaBuilder* b, XlaOp root) { - TF_ASSIGN_OR_RETURN(XlaComputation computation, b->Build(root)); + TF_ASSIGN_OR_RETURN(XlaComputation computation, + b->Build(root, /*remove_dynamic_dimensions=*/false)); const HloModuleProto& proto = computation.proto(); TF_ASSIGN_OR_RETURN(const auto& config, HloModule::CreateModuleConfigFromProto( @@ -132,6 +135,38 @@ TEST_F(XlaBuilderTest, BinaryOperatorsBuildExpectedHLO) { op::ShiftRightLogical(op::Constant(), op::Constant())); } +TEST_F(XlaBuilderTest, VariadicAnd) { + XlaBuilder b(TestName()); + Shape s = ShapeUtil::MakeShape(PRED, {}); + And(Parameter(&b, 0, s, "p0"), Parameter(&b, 1, s, "p1"), + Parameter(&b, 2, s, "p2")); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + // Don't specify in the test whether And(x, y, z) is right- or + // left-associative; accept either one. + EXPECT_THAT( + module->entry_computation()->root_instruction(), + ::testing::AnyOf(op::And(op::Parameter(0), + op::And(op::Parameter(1), op::Parameter(2))), + op::And(op::And(op::Parameter(0), op::Parameter(1)), + op::Parameter(2)))); +} + +TEST_F(XlaBuilderTest, VariadicOr) { + XlaBuilder b(TestName()); + Shape s = ShapeUtil::MakeShape(PRED, {}); + Or(Parameter(&b, 0, s, "p0"), Parameter(&b, 1, s, "p1"), + Parameter(&b, 2, s, "p2")); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + // Don't specify in the test whether Or(x, y, z) is right- or + // left-associative; accept either one. + EXPECT_THAT( + module->entry_computation()->root_instruction(), + ::testing::AnyOf( + op::Or(op::Parameter(0), op::Or(op::Parameter(1), op::Parameter(2))), + op::Or(op::Or(op::Parameter(0), op::Parameter(1)), + op::Parameter(2)))); +} + TEST_F(XlaBuilderTest, ShiftRightOperatorOnNonIntegerProducesError) { XlaBuilder b(TestName()); ConstantR0(&b, 1) >> ConstantR0(&b, 2); @@ -446,6 +481,461 @@ TEST_F(XlaBuilderTest, ProtoMatches) { EXPECT_EQ(c0_string, c1_string); } +TEST_F(XlaBuilderTest, DynamicParameter) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {5}), ShapeUtil::MakeShape(F32, {6})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + Parameter(&b, 1, ShapeUtil::MakeShape(U32, {}), "p1"); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/1, + /*dynamic_size_param_index=*/{}, + /*target_param_num=*/0, + /*target_param_index=*/{1}, + /*target_dim_num=*/0)); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b, /*root=*/p0)); + const Shape& param_shape = module->entry_computation() + ->parameter_instruction(0) + ->shape() + .tuple_shapes(1); + EXPECT_TRUE(param_shape.is_dynamic_dimension(0)); +} + +TEST_F(XlaBuilderTest, DynamicUnary) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {5}), ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{1}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/0)); + auto gte = GetTupleElement(p0, 0); + Neg(gte); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE(result_shape.is_dynamic_dimension(0)); +} + +TEST_F(XlaBuilderTest, DynamicBinary) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {5}), ShapeUtil::MakeShape(F32, {5}), + ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{2}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/0)); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{2}, + /*target_param_num=*/0, + /*target_param_index=*/{1}, + /*target_dim_num=*/0)); + auto gte0 = GetTupleElement(p0, 0); + auto gte1 = GetTupleElement(p0, 1); + Add(gte0, gte1); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE(result_shape.is_dynamic_dimension(0)); +} + +TEST_F(XlaBuilderTest, DynamicBinaryHasBroadcast) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {5, 4}), ShapeUtil::MakeShape(F32, {5}), + ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{2}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/0)); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{2}, + /*target_param_num=*/0, + /*target_param_index=*/{1}, + /*target_dim_num=*/0)); + auto gte0 = GetTupleElement(p0, 0); + auto gte1 = GetTupleElement(p0, 1); + Add(gte0, gte1, {0}); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE(ContainersEqual(result_shape.dynamic_dimensions(), {true, false})) + << result_shape; +} + +TEST_F(XlaBuilderTest, DynamicBroadcast) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {5, 4}), ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{1}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/0)); + auto gte = GetTupleElement(p0, 0); + BroadcastInDim(gte, /*out_dim_size=*/{3, 5, 4}, + /*broadcast_dimensions=*/{1, 2}); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE( + ContainersEqual(result_shape.dynamic_dimensions(), {false, true, false})) + << result_shape; +} + +TEST_F(XlaBuilderTest, DynamicBinaryHasDegenerateBroadcast) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {10}), ShapeUtil::MakeShape(F32, {1, 15}), + ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{1}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/0)); + auto gte0 = GetTupleElement(p0, 0); + auto gte1 = GetTupleElement(p0, 1); + Add(gte0, gte1, /*broadcast_dimensions=*/{0}); // f32[<=10, 15] + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE(ContainersEqual(result_shape.dynamic_dimensions(), {true, false})) + << result_shape; +} + +TEST_F(XlaBuilderTest, DynamicSelectOnlyPredDynamic) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(PRED, {10}), ShapeUtil::MakeShape(F32, {10}), + ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{1}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/0)); + auto gte0 = GetTupleElement(p0, 0); + auto gte1 = GetTupleElement(p0, 1); + + Select(gte0, gte1, gte1); + + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE(ContainersEqual(result_shape.dynamic_dimensions(), {true})) + << result_shape; +} + +TEST_F(XlaBuilderTest, DynamicPad) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {5, 4}), ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + auto pad_val = ConstantR0(&b, -1); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{1}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/0)); + auto gte = GetTupleElement(p0, 0); + PaddingConfig padding_config; + for (int i = 0; i < 2; i++) { + auto dimension = padding_config.add_dimensions(); + dimension->set_edge_padding_low(0); + dimension->set_edge_padding_high(0); + dimension->set_interior_padding(0); + } + Pad(gte, pad_val, padding_config); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE(ContainersEqual(result_shape.dynamic_dimensions(), {true, false})) + << result_shape; +} + +TEST_F(XlaBuilderTest, DynamicConvolution) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {1, 2, 2, 128}), + ShapeUtil::MakeShape(F32, {2, 2, 128, 8}), ShapeUtil::MakeShape(U32, {}), + ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{2}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/0)); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{3}, + /*target_param_num=*/0, + /*target_param_index=*/{1}, + /*target_dim_num=*/2)); + auto input = GetTupleElement(p0, 0); + auto filter = GetTupleElement(p0, 1); + ConvolutionDimensionNumbers dnums; + dnums.set_input_batch_dimension(0); + dnums.set_output_batch_dimension(0); + dnums.add_input_spatial_dimensions(1); + dnums.add_output_spatial_dimensions(1); + dnums.add_input_spatial_dimensions(2); + dnums.add_output_spatial_dimensions(2); + dnums.set_input_feature_dimension(3); + dnums.set_output_feature_dimension(3); + dnums.add_kernel_spatial_dimensions(0); + dnums.add_kernel_spatial_dimensions(1); + dnums.set_kernel_input_feature_dimension(2); + dnums.set_kernel_output_feature_dimension(3); + ConvWithGeneralDimensions(input, filter, {1, 1}, Padding::kValid, dnums, + /*feature_group_count=*/1); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE(ContainersEqual(result_shape.dynamic_dimensions(), + {true, false, false, false})) + << result_shape; +} + +TEST_F(XlaBuilderTest, DynamicDot) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {2, 3, 4}), + ShapeUtil::MakeShape(F32, {2, 4, 5}), ShapeUtil::MakeShape(U32, {}), + ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{2}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/0)); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{2}, + /*target_param_num=*/0, + /*target_param_index=*/{1}, + /*target_dim_num=*/0)); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{3}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/1)); + + auto lhs = GetTupleElement(p0, 0); + auto rhs = GetTupleElement(p0, 1); + DotDimensionNumbers dnums; + dnums.add_lhs_contracting_dimensions(2); + dnums.add_rhs_contracting_dimensions(1); + dnums.add_lhs_batch_dimensions(0); + dnums.add_rhs_batch_dimensions(0); + DotGeneral(lhs, rhs, dnums); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE( + ContainersEqual(result_shape.dynamic_dimensions(), {true, true, false})) + << result_shape; +} + +TEST_F(XlaBuilderTest, DynamicReduce) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {5, 4, 3}), ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + auto init = ConstantR0(&b, 0); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{1}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/1)); + auto gte = GetTupleElement(p0, 0); + XlaBuilder bsum(TestName()); + Add(Parameter(&bsum, 0, ShapeUtil::MakeShape(F32, {}), "x"), + Parameter(&bsum, 1, ShapeUtil::MakeShape(F32, {}), "y")); + TF_ASSERT_OK_AND_ASSIGN(auto sum, bsum.Build()); + Reduce(gte, init, sum, {0}); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE(ContainersEqual(result_shape.dynamic_dimensions(), {true, false})) + << result_shape; +} + +TEST_F(XlaBuilderTest, DynamicReduceWindow) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {2, 4, 8}), ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + auto init = ConstantR0(&b, 0.f); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{1}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/0)); + auto gte = GetTupleElement(p0, 0); + XlaBuilder bsum(TestName()); + Add(Parameter(&bsum, 0, ShapeUtil::MakeShape(F32, {}), "x"), + Parameter(&bsum, 1, ShapeUtil::MakeShape(F32, {}), "y")); + TF_ASSERT_OK_AND_ASSIGN(auto sum, bsum.Build()); + ReduceWindow(gte, init, sum, /*window_dimensions=*/{1, 2, 4}, + /*window_strides=*/{1, 1, 1}, Padding::kValid); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE( + ContainersEqual(result_shape.dynamic_dimensions(), {true, false, false})) + << result_shape; +} + +TEST_F(XlaBuilderTest, DynamicSelectAndScatter) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {2, 4, 8}), + ShapeUtil::MakeShape(F32, {2, 2, 2}), ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + auto init = ConstantR0(&b, 0.f); + XlaBuilder bsum(TestName()); + Add(Parameter(&bsum, 0, ShapeUtil::MakeShape(F32, {}), "x"), + Parameter(&bsum, 1, ShapeUtil::MakeShape(F32, {}), "y")); + TF_ASSERT_OK_AND_ASSIGN(auto sum, bsum.Build()); + XlaBuilder bge(TestName()); + Ge(Parameter(&bge, 0, ShapeUtil::MakeShape(F32, {}), "x"), + Parameter(&bge, 1, ShapeUtil::MakeShape(F32, {}), "y")); + TF_ASSERT_OK_AND_ASSIGN(auto ge, bge.Build()); + + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{2}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/0)); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{2}, + /*target_param_num=*/0, + /*target_param_index=*/{1}, + /*target_dim_num=*/0)); + auto gte0 = GetTupleElement(p0, 0); + auto source = GetTupleElement(p0, 1); + SelectAndScatter(gte0, ge, {1, 2, 4}, {1, 2, 4}, Padding::kValid, source, + init, sum); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE( + ContainersEqual(result_shape.dynamic_dimensions(), {true, false, false})) + << result_shape; +} + +TEST_F(XlaBuilderTest, DynamicReshape) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {2, 3, 4, 5, 6}), + ShapeUtil::MakeShape(U32, {}), ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{1}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/2)); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{2}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/3)); + auto gte = GetTupleElement(p0, 0); // f32[2, 3, <=4, <=5, 6] + Reshape(gte, /*new_sizes=*/{6, 4, 1, 5, 2, 3}); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE(result_shape.is_dynamic_dimension(1)); + EXPECT_TRUE(result_shape.is_dynamic_dimension(3)); + EXPECT_TRUE(ContainersEqual(result_shape.dynamic_dimensions(), + {false, true, false, true, false, false})) + << result_shape; +} + +TEST_F(XlaBuilderTest, DynamicSelect) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {4, 5, 6}), + ShapeUtil::MakeShape(F32, {4, 5, 6}), ShapeUtil::MakeShape(U32, {}), + ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + auto pred = Parameter(&b, 1, ShapeUtil::MakeShape(PRED, {}), "pred"); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{2}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/1)); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{3}, + /*target_param_num=*/0, + /*target_param_index=*/{1}, + /*target_dim_num=*/1)); + auto gte0 = GetTupleElement(p0, 0); + auto gte1 = GetTupleElement(p0, 1); + Select(pred, gte0, gte1); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE(result_shape.is_dynamic_dimension(1)); + EXPECT_FALSE(result_shape.is_dynamic_dimension(2)); + EXPECT_TRUE( + ContainersEqual(result_shape.dynamic_dimensions(), {false, true, false})) + << result_shape; +} + +TEST_F(XlaBuilderTest, DynamicSelectNotCompatible) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {4, 5, 6}), + ShapeUtil::MakeShape(F32, {4, 5, 6}), ShapeUtil::MakeShape(U32, {}), + ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + auto pred = Parameter(&b, 1, ShapeUtil::MakeShape(PRED, {}), "pred"); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{2}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/1)); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{3}, + /*target_param_num=*/0, + /*target_param_index=*/{1}, + /*target_dim_num=*/2)); + auto gte0 = GetTupleElement(p0, 0); // f32[4,<=5,6] + auto gte1 = GetTupleElement(p0, 1); // f32[4,5,<=6] + Select(pred, gte0, gte1); + Status status = BuildHloModule(&b).status(); + ASSERT_IS_NOT_OK(status); + EXPECT_THAT(status.error_message(), + ::testing::HasSubstr("Operands to select must be the same shape; " + "got f32[4,<=5,6] and f32[4,5,<=6]")); +} + +TEST_F(XlaBuilderTest, DynamicTranspose) { + XlaBuilder b(TestName()); + Shape tuple_param_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {3, 5}), ShapeUtil::MakeShape(U32, {})}); + auto p0 = Parameter(&b, 0, tuple_param_shape, "p0"); + ASSERT_IS_OK(b.SetDynamicBinding(/*dynamic_size_param_num=*/0, + /*dynamic_size_param_index=*/{1}, + /*target_param_num=*/0, + /*target_param_index=*/{0}, + /*target_dim_num=*/0)); + auto gte = GetTupleElement(p0, 0); + Transpose(gte, /*permutation=*/{1, 0}); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + const Shape& result_shape = + module->entry_computation()->root_instruction()->shape(); + EXPECT_TRUE(ContainersEqual(result_shape.dynamic_dimensions(), {false, true})) + << result_shape; +} + TEST_F(XlaBuilderTest, AfterAllWithNonTokenOperands) { XlaBuilder b(TestName()); AfterAll(&b, {CreateToken(&b), ConstantR0(&b, 1.0)}); diff --git a/tensorflow/compiler/xla/error_spec.h b/tensorflow/compiler/xla/error_spec.h index a1463aa15941b9c265db94e2eb3cc176fab6695b..4359f3b7deb8e585494cb2a9c7115eac6a312c8e 100644 --- a/tensorflow/compiler/xla/error_spec.h +++ b/tensorflow/compiler/xla/error_spec.h @@ -30,6 +30,19 @@ struct ErrorSpec { // In effect, this allows the tested operation to produce incorrect results // for inputs outside its mathematical domain. bool relaxed_nans; + + // If this is true, then we treat each +/-inf in the actual result as + // equivalent to our choice of either +/-inf or the min/max floating-point + // value. + // + // If the expected result is +/-inf, the actual result must still be +/-inf. + // + // In effect, this allows the tested operation to overflow, so long as it's + // overflowing on "large" values. + // + // (We could have a symmetric more_infs_ok flag if necessary; right now it + // appears not to be.) + bool fewer_infs_ok = false; }; } // namespace xla diff --git a/tensorflow/compiler/xla/executable_run_options.h b/tensorflow/compiler/xla/executable_run_options.h index ba3217f31b55bd1428f67da6154a46c8bc304053..6f36d11dfb34eb27e79ea4ff797d35f80fb44b27 100644 --- a/tensorflow/compiler/xla/executable_run_options.h +++ b/tensorflow/compiler/xla/executable_run_options.h @@ -16,9 +16,6 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_EXECUTABLE_RUN_OPTIONS_H_ #define TENSORFLOW_COMPILER_XLA_EXECUTABLE_RUN_OPTIONS_H_ -// Pulls in the ::stream_executor -> ::xla::se namespace alias. -#include "tensorflow/compiler/xla/types.h" - // These classes are forward declared so that ExecutableRunOptions can be linked // into an XLA-compiled binary without having to link all of the pointed-to // objects (e.g., for an ahead-of-time compiled CPU binary, the gpu tools don't @@ -28,12 +25,6 @@ class Stream; class Platform; } // namespace stream_executor -namespace tensorflow { -namespace thread { -class ThreadPool; -} // namespace thread -} // namespace tensorflow - namespace Eigen { struct ThreadPoolDevice; } // namespace Eigen diff --git a/tensorflow/compiler/xla/g3doc/_book.yaml b/tensorflow/compiler/xla/g3doc/_book.yaml index 267701e9c0e42a21d2cda6238520f6a9692e7e76..d756cd74c98b98a6fda099690d966562bd694e2c 100644 --- a/tensorflow/compiler/xla/g3doc/_book.yaml +++ b/tensorflow/compiler/xla/g3doc/_book.yaml @@ -25,6 +25,8 @@ upper_tabs: path: /xla/operation_semantics - title: Shapes and layout path: /xla/shapes + - title: Tiled layout + path: /xla/tiled_layout - title: Using AOT compilation path: /xla/tfcompile - heading: Tutorials diff --git a/tensorflow/compiler/xla/g3doc/broadcasting.md b/tensorflow/compiler/xla/g3doc/broadcasting.md index 2870869a2cef13a9105b9dc9fa4d657834288f86..5c0525c1e9adf9f37d945170d05e7c18fa3d8852 100644 --- a/tensorflow/compiler/xla/g3doc/broadcasting.md +++ b/tensorflow/compiler/xla/g3doc/broadcasting.md @@ -168,7 +168,7 @@ consult the Broadcasting of a lower-rank array to a higher-rank array **and** broadcasting using degenerate dimensions can both be performed in the same binary operation. -For example, a vector of size 4 and an matrix of size 1x2 can be added together +For example, a vector of size 4 and a matrix of size 1x2 can be added together using broadcast dimensions value of (0): |1 2 3 4| + [5 6] // [5 6] is a 1x2 matrix, not a vector. @@ -176,7 +176,7 @@ using broadcast dimensions value of (0): First the vector is broadcast up to rank 2 (matrix) using the broadcast dimensions. The single value (0) in the broadcast dimensions indicates that dimension zero of the vector matches to dimension zero of the matrix. This -produces an matrix of size 4xM where the value M is chosen to match the +produces a matrix of size 4xM where the value M is chosen to match the corresponding dimension size in the 1x2 array. Therefore, a 4x2 matrix is produced: diff --git a/tensorflow/compiler/xla/g3doc/operation_semantics.md b/tensorflow/compiler/xla/g3doc/operation_semantics.md index 78d461701750c4477bc0f7a44d96958d501a93fa..db90d184b5218614ac49363ebf2a7e25fffe44de 100644 --- a/tensorflow/compiler/xla/g3doc/operation_semantics.md +++ b/tensorflow/compiler/xla/g3doc/operation_semantics.md @@ -636,11 +636,15 @@ details, see `tf.nn.depthwise_conv2d`. The `batch_group_count` (default value 1) argument can be used for depthwise filters during backpropagation. `batch_group_count` needs to be a divisor of the -size of the `lhs` batch dimension. If `batch_group_count` is greater than 1, it -means that conceptually the output batch dimension is split evenely in -`batch_group_count` groups, such that each group consists of a consecutive -subsequence of batches. Each output batch element is the reduced value of the -batch group size. +size of the `lhs` (input) batch dimension. If `batch_group_count` is greater +than 1, it means that the output batch dimension should be of size +`batch_group_size` where `batch_group_size = input batch / batch_group_count`. +For convolutions with `batch_group_count` greater than 1, the input batch size +must evenly divide into batch_group_size and output feature size, which implies +that the output feature size must be equal to batch_group_count. Conceptually, +this can be achieved by performing the usual convolution, and then scraping +`batch_group_size` number of elements on the diagonal of the matrix formed by +output batch and output feature. The output shape has these dimensions, in this order: @@ -871,9 +875,7 @@ DotGeneral performs the sum of products over contracting dimensions specified in 'dimension_numbers'. Associated contracting dimension numbers from the 'lhs' and 'rhs' do not need -to be the same, but must be listed in the same order in both -'lhs/rhs_contracting_dimensions' arrays and have the same dimension sizes. -There must be exactly one contracting dimension on both 'lhs' and 'rhs'. +to be the same and but must have the same dimension sizes. Example with contracting dimension numbers: @@ -892,10 +894,8 @@ DotGeneral(lhs, rhs, dnums) -> { {6.0, 12.0}, {15.0, 30.0} } ``` -Associated batch dimension numbers from the 'lhs' and 'rhs' must have the same -dimension number, must be listed in the same order in both arrays, must -have the same dimension sizes, and must be ordered before contracting and -non-contracting/non-batch dimension numbers. +Associated batch dimension numbers from the 'lhs' and 'rhs' must +have the same dimension sizes. Example with batch dimension numbers (batch size 2, 2x2 matrices): @@ -944,21 +944,21 @@ dimension: [start, start + size). The shape of `start_indices` must be rank == `DynamicSlice(operand, start_indices, size_indices)` -| Arguments | Type | Semantics | -| --------------- | ------------------- | ----------------------------------- | -| `operand` | `XlaOp` | N dimensional array of type T | -| `start_indices` | `XlaOp` | Rank 1 array of N integers | -: : : containing the starting indices of : -: : : the slice for each dimension. Value : -: : : must be greater than or equal to : -: : : zero. : -| `size_indices` | `ArraySlice` | List of N integers containing the | -: : : slice size for each dimension. Each : -: : : value must be strictly greater than : -: : : zero, and start + size must be less : -: : : than or equal to the size of the : -: : : dimension to avoid wrapping modulo : -: : : dimension size. : +| Arguments | Type | Semantics | +| --------------- | --------------------- | ---------------------------------- | +| `operand` | `XlaOp` | N dimensional array of type T | +| `start_indices` | sequence of N `XlaOp` | List of N scalar integers | +: : : containing the starting indices of : +: : : the slice for each dimension. : +: : : Value must be greater than or : +: : : equal to zero. : +| `size_indices` | `ArraySlice` | List of N integers containing the | +: : : slice size for each dimension. : +: : : Each value must be strictly : +: : : greater than zero, and start + : +: : : size must be less than or equal to : +: : : the size of the dimension to avoid : +: : : wrapping modulo dimension size. : The effective slice indices are computed by applying the following transformation for each index `i` in `[1, N)` before performing the slice: @@ -1009,19 +1009,22 @@ the rank of `operand`. `DynamicUpdateSlice(operand, update, start_indices)` -| Arguments | Type | Semantics | -| --------------- | ------- | ------------------------------------------------ | -| `operand` | `XlaOp` | N dimensional array of type T | -| `update` | `XlaOp` | N dimensional array of type T containing the | -: : : slice update. Each dimension of update shape : -: : : must be strictly greater than zero, and start + : -: : : update must be less than or equal to the operand : -: : : size for each dimension to avoid generating : -: : : out-of-bounds update indices. : -| `start_indices` | `XlaOp` | Rank 1 array of N integers containing the | -: : : starting indices of the slice for each : -: : : dimension. Value must be greater than or equal : -: : : to zero. : +| Arguments | Type | Semantics | +| --------------- | --------------------- | ---------------------------------- | +| `operand` | `XlaOp` | N dimensional array of type T | +| `update` | `XlaOp` | N dimensional array of type T | +: : : containing the slice update. Each : +: : : dimension of update shape must be : +: : : strictly greater than zero, and : +: : : start + update must be less than : +: : : or equal to the operand size for : +: : : each dimension to avoid generating : +: : : out-of-bounds update indices. : +| `start_indices` | sequence of N `XlaOp` | List of N scalar integers | +: : : containing the starting indices of : +: : : the slice for each dimension. : +: : : Value must be greater than or : +: : : equal to zero. : The effective slice indices are computed by applying the following transformation for each index `i` in `[1, N)` before performing the slice: @@ -1183,7 +1186,7 @@ if and only if the corresponding input element is finite. `Sign(operand)` Element-wise sign operation `x -> sgn(x)` where -$$\text{sgn}(x) = \begin{cases} -1 & x < 0\\ 0 & x = 0\\ 1 & x > 0 \end{cases}$$ +$$\text{sgn}(x) = \begin{cases} -1 & x < 0\\ -0 & x = -0\\ NaN & x = NaN\\ +0 & x = +0\\ 1 & x > 0 \end{cases}$$ using the comparison operator of the element type of `operand`. @@ -1870,6 +1873,20 @@ non-deterministic. Therefore, the reduction function should not be overly sensitive to reassociation. See the discussion about associativity in the context of [`Reduce`](#reduce) for more details. +## ReplicaId + +See also +[`XlaBuilder::ReplicaId`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Returns the unique ID (U32 scalar) of the replica. + + `ReplicaId()` + +The unique ID of each replica is an unsigned integer in the interval `[0, N)`, +where `N` is the number of replicas. Since all the replicas are running the same +program, a `ReplicaId()` call in the program will return a different value on +each replica. + ## Reshape See also diff --git a/tensorflow/compiler/xla/g3doc/layout_with_tiling.md b/tensorflow/compiler/xla/g3doc/tiled_layout.md similarity index 96% rename from tensorflow/compiler/xla/g3doc/layout_with_tiling.md rename to tensorflow/compiler/xla/g3doc/tiled_layout.md index 5e990851af7495ebd4417e44f1d955fcc14dadf1..21e88ceab6208cdf940826d769fd93713044d5a0 100644 --- a/tensorflow/compiler/xla/g3doc/layout_with_tiling.md +++ b/tensorflow/compiler/xla/g3doc/tiled_layout.md @@ -1,9 +1,7 @@ # Tiled layout -*Note: This doc describes how tiled layout is intended to work. Tiling is being -implemented, but this is an early effort and it is currently not even guaranteed -to get an Unimplemented error if one tries to use tiling - it may be just -silently ignored.* +Caution: Tiled layout is *pre-release* and this describes how it's intended to +work. Errors may be silently ignored.
![](images/xla_array_layout_figure1.png) diff --git a/tensorflow/compiler/xla/index_util.cc b/tensorflow/compiler/xla/index_util.cc index 7e22a32e545e4155545ffcfb9582187eadec3a82..eebd8245abe759b71b3fe732943761325ea04b81 100644 --- a/tensorflow/compiler/xla/index_util.cc +++ b/tensorflow/compiler/xla/index_util.cc @@ -21,7 +21,6 @@ limitations under the License. #include "absl/strings/str_join.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/logging.h" namespace xla { diff --git a/tensorflow/compiler/xla/layout.cc b/tensorflow/compiler/xla/layout.cc index e3b5fcd5274881cec31ecf906e3461685f82a1f4..d2f7985aab7123a80ce626b27aa612edda87f761 100644 --- a/tensorflow/compiler/xla/layout.cc +++ b/tensorflow/compiler/xla/layout.cc @@ -30,7 +30,19 @@ TileProto Tile::ToProto() const { } string Tile::ToString() const { - return absl::StrCat("(", absl::StrJoin(dimensions(), ","), ")"); + std::vector elements; + for (auto dim : dimensions()) { + if (dim >= 0) { + elements.push_back(std::to_string(dim)); + } else { + CHECK_EQ(dim, kCombineDimension) + << "Tile dimension size needs to be mininum int64 value if it's " + "negative. Value is " + << dim; + elements.push_back("*"); + } + } + return absl::StrCat("(", absl::StrJoin(elements, ","), ")"); } /* static */ Layout Layout::CreateFromProto(const LayoutProto& proto) { @@ -64,11 +76,19 @@ LayoutProto Layout::ToProto() const { } string Layout::ToString() const { - // TODO(b/119839262): Emit tiles in string. if (format() == SPARSE) { + CHECK_EQ(tiles_size(), 0) << "Sparse layout should not be tiled."; return absl::StrCat("sparse{", max_sparse_elements(), "}"); } else if (format() == DENSE) { - return absl::StrCat("{", absl::StrJoin(minor_to_major(), ","), "}"); + string colon_string = tiles().empty() ? "" : "T"; + for (Tile tile : tiles()) { + absl::StrAppend(&colon_string, tile.ToString()); + } + if (element_size_in_bits() != 0) { + absl::StrAppend(&colon_string, "E(", element_size_in_bits(), ")"); + } + return absl::StrCat("{", absl::StrJoin(minor_to_major(), ","), + colon_string.empty() ? "" : ":", colon_string, "}"); } else { CHECK_EQ(format(), INVALID_FORMAT); return "invalid{}"; diff --git a/tensorflow/compiler/xla/layout.h b/tensorflow/compiler/xla/layout.h index 313368c39e4c976fc481941eb17325101f2ba69a..1faa1629980f5a8d954b70a4d860dbf708de2624 100644 --- a/tensorflow/compiler/xla/layout.h +++ b/tensorflow/compiler/xla/layout.h @@ -55,6 +55,20 @@ class Tile { // Returns the dimensions of the tile. const std::vector& dimensions() const { return dimensions_; } + Tile& add_dimensions(int64 value) { + dimensions_.push_back(value); + return *this; + } + + Tile& clear_dimensions() { + dimensions_.clear(); + return *this; + } + + // This dimension size means the corresponding dimension in the shape is + // combined with the next minor dimension before tiling is applied. + static constexpr int64 kCombineDimension = std::numeric_limits::min(); + private: // The bounds of the tile. std::vector dimensions_; diff --git a/tensorflow/compiler/xla/layout_test.cc b/tensorflow/compiler/xla/layout_test.cc index fb6abd3f6523b978e72b21ec082ae06973e86243..7d43b0b87c8eeabf1d30187625a67967a11c3eb4 100644 --- a/tensorflow/compiler/xla/layout_test.cc +++ b/tensorflow/compiler/xla/layout_test.cc @@ -38,10 +38,10 @@ TEST_F(LayoutTest, ToString) { "sparse{123}"); EXPECT_EQ(Layout({4, 5, 6}).ToString(), "{4,5,6}"); EXPECT_EQ(Layout({3, 2, 1, 0}, {Tile({42, 123}), Tile({4, 5})}).ToString(), - "{3,2,1,0}"); + "{3,2,1,0:T(42,123)(4,5)}"); EXPECT_EQ( Layout({1, 0}, {Tile({2, 55})}).set_element_size_in_bits(42).ToString(), - "{1,0}"); + "{1,0:T(2,55)E(42)}"); } TEST_F(LayoutTest, StreamOut) { diff --git a/tensorflow/compiler/xla/layout_util.cc b/tensorflow/compiler/xla/layout_util.cc index 42649d669222e6a871c81a01c54148c93f12a140..62314118ca9713a04cb4e3cf6ad261b966d85f15 100644 --- a/tensorflow/compiler/xla/layout_util.cc +++ b/tensorflow/compiler/xla/layout_util.cc @@ -54,12 +54,24 @@ void SetDefaultLayoutToContainer(std::vector* minor_to_major) { } // namespace /* static */ Layout LayoutUtil::MakeLayout( - absl::Span minor_to_major) { + absl::Span minor_to_major, absl::Span tiles, + int64 element_size_in_bits) { Layout layout; layout.set_format(DENSE); for (int64 dimension_number : minor_to_major) { layout.add_minor_to_major(dimension_number); } + for (Tile tile : tiles) { + for (int64 dim : tile.dimensions()) { + if (dim < 0 && dim != Tile::kCombineDimension) { + LOG(FATAL) << "Tile dimension size needs to be mininum int64 value if " + "it's negative. Value is " + << dim; + } + } + *layout.add_tiles() = tile; + } + layout.set_element_size_in_bits(element_size_in_bits); return layout; } @@ -235,6 +247,10 @@ Layout CreateDefaultLayoutForRank(int64 rank) { } dimensions_in_layout[dim] = true; } + } else { + if (layout.tiles_size() != 0) { + return InvalidArgument("Only dense layouts can be tiled."); + } } return Status::OK(); @@ -290,8 +306,8 @@ Layout CreateDefaultLayoutForRank(int64 rank) { /* static */ bool LayoutUtil::HasLayout(const Shape& shape) { if (shape.IsTuple()) { // Tuple shape: all subshapes must have a layout. - return std::all_of(shape.tuple_shapes().begin(), shape.tuple_shapes().end(), - [](const Shape& s) { return HasLayout(s); }); + return absl::c_all_of(shape.tuple_shapes(), + [](const Shape& s) { return HasLayout(s); }); } else if (!shape.IsArray()) { // Opaque, token types etc. ignore layout. return true; @@ -424,7 +440,7 @@ Status LayoutUtil::CopyLayoutBetweenShapes(const Shape& src, Shape* dst) { positions_in_layout.push_back( PositionInContainer(layout.minor_to_major(), dim)); } - std::sort(positions_in_layout.begin(), positions_in_layout.end()); + absl::c_sort(positions_in_layout); for (size_t i = 1; i < positions_in_layout.size(); ++i) { if (1 != positions_in_layout[i] - positions_in_layout[i - 1]) { return false; diff --git a/tensorflow/compiler/xla/layout_util.h b/tensorflow/compiler/xla/layout_util.h index 609dba67bcdbcb11be0906b7d87a52a17ba0dfbd..9997aef465daa48ee77050e03d97cde0ea2425cc 100644 --- a/tensorflow/compiler/xla/layout_util.h +++ b/tensorflow/compiler/xla/layout_util.h @@ -36,7 +36,9 @@ class LayoutUtil { public: // Creates a layout with the given minor-to-major dimension order. (This is a // convenience function for protobuf construction.) - static Layout MakeLayout(absl::Span minor_to_major); + static Layout MakeLayout(absl::Span minor_to_major, + absl::Span tiles = {}, + int64 element_size_in_bits = 0); // Similar to MakeLayout, but take indices in reverse order. static Layout MakeLayoutFromMajorToMinor( diff --git a/tensorflow/compiler/xla/layout_util_test.cc b/tensorflow/compiler/xla/layout_util_test.cc index 4cc94c270cd64eb19761cc1044861c7d185b7888..12da214063676717aa075e66aa54974f4cc2b31b 100644 --- a/tensorflow/compiler/xla/layout_util_test.cc +++ b/tensorflow/compiler/xla/layout_util_test.cc @@ -317,6 +317,81 @@ TEST_F(LayoutUtilTest, DefaultLayoutGettersMajorToMinor) { ShapeUtil::MakeShape(F32, {10, 20, 30, 15, 25})))); } +TEST_F(LayoutUtilTest, HumanStringWithTiling) { + Shape shape = ShapeUtil::MakeShapeWithLayout(F32, {2, 3, 4}, {0, 1, 2}); + Tile* tile; + + // No tiling. + EXPECT_EQ(ShapeUtil::HumanStringWithLayout(shape), "f32[2,3,4]{0,1,2}"); + + // 2D tile. + tile = shape.mutable_layout()->add_tiles(); + tile->add_dimensions(512); + tile->add_dimensions(1024); + EXPECT_EQ(ShapeUtil::HumanStringWithLayout(shape), + "f32[2,3,4]{0,1,2:T(512,1024)}"); + + // 1D tile. + shape.mutable_layout()->clear_tiles(); + tile = shape.mutable_layout()->add_tiles(); + tile->add_dimensions(512); + EXPECT_EQ(ShapeUtil::HumanStringWithLayout(shape), + "f32[2,3,4]{0,1,2:T(512)}"); + + // 2 tiles. + shape = ShapeUtil::MakeShapeWithLayout(BF16, {2, 3, 4}, {1, 2, 0}); + tile = shape.mutable_layout()->add_tiles(); + tile->add_dimensions(16); + tile->add_dimensions(256); + tile = shape.mutable_layout()->add_tiles(); + tile->add_dimensions(2); + tile->add_dimensions(1); + EXPECT_EQ(ShapeUtil::HumanStringWithLayout(shape), + "bf16[2,3,4]{1,2,0:T(16,256)(2,1)}"); + + // PRED with element size of 8 bits. + shape = ShapeUtil::MakeShapeWithLayout(PRED, {8, 8, 8}, {0, 2, 1}); + tile = shape.mutable_layout()->add_tiles(); + tile->add_dimensions(8); + tile->add_dimensions(128); + EXPECT_EQ(ShapeUtil::HumanStringWithLayout(shape), + "pred[8,8,8]{0,2,1:T(8,128)}"); + + // PRED with element size of 32 bits. + shape.mutable_layout()->clear_tiles(); + tile = shape.mutable_layout()->add_tiles(); + tile->add_dimensions(8); + tile->add_dimensions(128); + shape.mutable_layout()->set_element_size_in_bits(32); + EXPECT_EQ(ShapeUtil::HumanStringWithLayout(shape), + "pred[8,8,8]{0,2,1:T(8,128)E(32)}"); + + // No tile. PRED with element size of 32 bits. + shape.mutable_layout()->clear_tiles(); + shape.mutable_layout()->set_element_size_in_bits(32); + EXPECT_EQ(ShapeUtil::HumanStringWithLayout(shape), + "pred[8,8,8]{0,2,1:E(32)}"); + + // Tile with negative dimension size for combining dimensions. + shape = ShapeUtil::MakeShapeWithLayout(BF16, {2, 3, 1004}, {2, 1, 0}); + tile = shape.mutable_layout()->add_tiles(); + tile->add_dimensions(2); + tile->add_dimensions(Tile::kCombineDimension); + tile->add_dimensions(128); + EXPECT_EQ(ShapeUtil::HumanStringWithLayout(shape), + "bf16[2,3,1004]{2,1,0:T(2,*,128)}"); + + // Tile with two negative dimensions. + shape = ShapeUtil::MakeShapeWithLayout(BF16, {8, 2, 3, 1004}, {3, 2, 1, 0}); + tile = shape.mutable_layout()->add_tiles(); + tile->add_dimensions(2); + tile->add_dimensions(Tile::kCombineDimension); + tile->add_dimensions(Tile::kCombineDimension); + tile->add_dimensions(128); + EXPECT_EQ(ShapeUtil::HumanStringWithLayout(shape), + "bf16[8,2,3,1004]{3,2,1,0:T(2,*,*,128)}"); +} + TEST_F(LayoutUtilTest, ValidateLayout_ValidArrayLayout) { Shape shape = ShapeUtil::MakeShapeWithLayout(F32, {2, 3}, {0, 1}); auto status = diff --git a/tensorflow/compiler/xla/literal.cc b/tensorflow/compiler/xla/literal.cc index 6c441899300b31d9cc544f1de8bad7454f063fa7..5cd738d0f7769ceac7eb3bdbc5abd3196d9cf99c 100644 --- a/tensorflow/compiler/xla/literal.cc +++ b/tensorflow/compiler/xla/literal.cc @@ -29,10 +29,12 @@ limitations under the License. #include "absl/strings/str_join.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/index_util.h" +#include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/platform/logging.h" @@ -42,7 +44,6 @@ namespace xla { namespace { using absl::StrCat; -using absl::StrFormat; constexpr bool kLittleEndian = __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__; @@ -411,6 +412,7 @@ Status LiteralBase::Piece::CopyFrom(const LiteralBase::Piece& src) { COPY_ELEMENTS(F32, float); COPY_ELEMENTS(F64, double); COPY_ELEMENTS(C64, complex64); + COPY_ELEMENTS(C128, complex128); COPY_ELEMENTS(PRED, bool); #undef COPY_ELEMENTS default: @@ -548,6 +550,9 @@ Status MutableLiteralBase::CopySliceFrom(const LiteralSlice& src_literal, case C64: return CopySliceFromInternal(src_literal, src_base, dest_base, copy_size); + case C128: + return CopySliceFromInternal(src_literal, src_base, dest_base, + copy_size); case PRED: return CopySliceFromInternal(src_literal, src_base, dest_base, copy_size); @@ -766,6 +771,8 @@ Literal LiteralBase::Slice(absl::Span start_indices, return SliceInternal(result_shape, start_indices); case C64: return SliceInternal(result_shape, start_indices); + case C128: + return SliceInternal(result_shape, start_indices); default: LOG(FATAL) << "not yet implemented: " << PrimitiveType_Name(result_shape.element_type()); @@ -814,6 +821,10 @@ string LiteralBase::GetAsString(absl::Span multi_index, complex64 c = Get(multi_index, shape_index); return StrCat("(", c.real(), ", ", c.imag(), ")"); } + case C128: { + complex128 c = Get(multi_index, shape_index); + return StrCat("(", c.real(), ", ", c.imag(), ")"); + } default: LOG(FATAL) << PrimitiveType_Name(subshape.element_type()); } @@ -868,6 +879,11 @@ string LiteralBase::GetSparseElementAsString( GetSparseElement(sparse_element_number, shape_index); return StrCat("(", c.real(), ", ", c.imag(), ")"); } + case C128: { + complex128 c = + GetSparseElement(sparse_element_number, shape_index); + return StrCat("(", c.real(), ", ", c.imag(), ")"); + } default: LOG(FATAL) << "Invalid element type for sparse arrays: " << PrimitiveType_Name(subshape.element_type()); @@ -996,6 +1012,9 @@ void LiteralBase::Piece::SortSparseElements() { case C64: SortSparseElementsInternal(); break; + case C128: + SortSparseElementsInternal(); + break; case F16: SortSparseElementsInternal(); break; @@ -1230,7 +1249,24 @@ Literal ConvertBetweenNativeTypesWithConverter(const LiteralBase& src_literal, } template -Literal ConvertBetweenNativeTypes(const LiteralBase& src_literal) { +typename std::enable_if<(std::is_same::value) && + (std::is_same::value || + std::is_same::value), + Literal>::type +ConvertBetweenNativeTypes(const LiteralBase& src_literal) { + auto converter = [](NativeSrcT src) { + return NativeDestT(static_cast(src)); + }; + return ConvertBetweenNativeTypesWithConverter( + src_literal, converter); +} + +template +typename std::enable_if<(!std::is_same::value) || + (!std::is_same::value && + !std::is_same::value), + Literal>::type +ConvertBetweenNativeTypes(const LiteralBase& src_literal) { auto converter = [](NativeSrcT src) { return static_cast(src); }; return ConvertBetweenNativeTypesWithConverter( src_literal, converter); @@ -1274,22 +1310,6 @@ BitcastBetweenNativeTypes(const LiteralBase& src_literal) { LOG(FATAL) << "Invalid bitcast between types of different sizes."; } -template -Literal ConvertToC64(const LiteralBase& src_literal) { - CHECK(src_literal.shape().IsArray()); - Literal result_literal( - ShapeUtil::ChangeElementType(src_literal.shape(), C64)); - using NativeSrcT = - typename primitive_util::PrimitiveTypeToNative::type; - absl::Span src_data = src_literal.data(); - absl::Span dest_data = result_literal.data(); - int64 num_elements = src_literal.element_count(); - for (int64 i = 0; i < num_elements; ++i) { - dest_data[i] = complex64(static_cast(src_data[i]), 0); - } - return result_literal; -} - template Literal ConvertIfTypesMatch(const LiteralBase& src_literal, bool bitcast) { CHECK_EQ(primitive_src_type, src_literal.shape().element_type()); @@ -1319,9 +1339,11 @@ StatusOr ConvertIfDestTypeMatches(const LiteralBase& src_literal, bitcast); CONVERT_IF_TYPES_MATCH(PRED) CONVERT_IF_TYPES_MATCH(S8) + CONVERT_IF_TYPES_MATCH(S16) CONVERT_IF_TYPES_MATCH(S32) CONVERT_IF_TYPES_MATCH(S64) CONVERT_IF_TYPES_MATCH(U8) + CONVERT_IF_TYPES_MATCH(U16) CONVERT_IF_TYPES_MATCH(U32) CONVERT_IF_TYPES_MATCH(U64) CONVERT_IF_TYPES_MATCH(F16) @@ -1330,10 +1352,15 @@ StatusOr ConvertIfDestTypeMatches(const LiteralBase& src_literal, CONVERT_IF_TYPES_MATCH(BF16) #undef CONVERT_IF_TYPES_MATCH case C64: - if (!bitcast) { - return ConvertToC64(src_literal); + if (bitcast) { + break; } - break; + return ConvertIfTypesMatch(src_literal, false); + case C128: + if (bitcast) { + break; + } + return ConvertIfTypesMatch(src_literal, false); // Other types are not yet supported. default: break; @@ -1357,9 +1384,11 @@ StatusOr ConvertSwitch(const LiteralBase& literal, bitcast); CONVERT_IF_DEST_TYPE_MATCHES(PRED) CONVERT_IF_DEST_TYPE_MATCHES(S8) + CONVERT_IF_DEST_TYPE_MATCHES(S16) CONVERT_IF_DEST_TYPE_MATCHES(S32) CONVERT_IF_DEST_TYPE_MATCHES(S64) CONVERT_IF_DEST_TYPE_MATCHES(U8) + CONVERT_IF_DEST_TYPE_MATCHES(U16) CONVERT_IF_DEST_TYPE_MATCHES(U32) CONVERT_IF_DEST_TYPE_MATCHES(U64) CONVERT_IF_DEST_TYPE_MATCHES(F16) @@ -1481,6 +1510,8 @@ bool LiteralBase::Piece::EqualElements(const LiteralBase::Piece& other) const { return EqualElementsInternal(other, &multi_index); case C64: return EqualElementsInternal(other, &multi_index); + case C128: + return EqualElementsInternal(other, &multi_index); default: LOG(FATAL) << "Unimplemented: LiteralBase::Piece::EqualElements for type " << PrimitiveType_Name(subshape().element_type()); @@ -1596,26 +1627,20 @@ bool LiteralBase::IsAllFloat(float value) const { return true; } - auto piece_is_all = [&]() { - switch (shape().element_type()) { - case F32: - return AllElementsEqualValue(piece.data(), value); - case F64: - return AllElementsEqualValue(piece.data(), value); - case F16: - return AllElementsEqualValue(piece.data(), - static_cast(value)); - case BF16: - return AllElementsEqualValue( - piece.data(), static_cast(value)); - default: - return false; - } - }; - if (!piece_is_all()) { - return false; + switch (shape().element_type()) { + case F32: + return AllElementsEqualValue(piece.data(), value); + case F64: + return AllElementsEqualValue(piece.data(), value); + case F16: + return AllElementsEqualValue(piece.data(), + static_cast(value)); + case BF16: + return AllElementsEqualValue( + piece.data(), static_cast(value)); + default: + return false; } - return true; }); } @@ -1624,6 +1649,9 @@ bool LiteralBase::IsAllComplex(complex64 value) const { case C64: return AllElementsEqualValue(root_piece().data(), value); + case C128: + return AllElementsEqualValue(root_piece().data(), + value); default: return false; } @@ -1703,6 +1731,11 @@ bool LiteralBase::IsAllFirst() const { auto data = piece.data(); return AllElementsEqualValue(data, data[0]); } + + case C128: { + auto data = piece.data(); + return AllElementsEqualValue(data, data[0]); + } default: return false; } @@ -1752,6 +1785,8 @@ bool LiteralBase::IsR1Iota() const { return Get({idx}) == static_cast(idx); case C64: return Get({idx}) == complex64(idx, 0.0f); + case C128: + return Get({idx}) == complex128(idx, 0.0f); case PRED: return Get({idx}) == idx; // token, opaque, tuple, etc. are all not iota. @@ -1795,6 +1830,8 @@ bool LiteralBase::IsZero(absl::Span indices) const { return Get(indices) == 0.0; case C64: return Get(indices) == complex64(0.0f, 0.0f); + case C128: + return Get(indices) == complex128(0.0f, 0.0f); case F16: return Get(indices) == static_cast(0.0f); case BF16: @@ -1882,6 +1919,12 @@ void LiteralBase::Piece::WriteToProto(LiteralProto* proto) const { proto->add_c64s(value.imag()); } break; + case C128: + for (complex128 value : data()) { + proto->add_c128s(value.real()); + proto->add_c128s(value.imag()); + } + break; case TUPLE: case TOKEN: // Nothing to do but assign the shape which is done above. @@ -2014,7 +2057,17 @@ Status LiteralBase::Piece::CopyFromProto(const LiteralProto& proto) { for (int64 i = 0; i < complex_data.size(); ++i) { complex_data[i] = complex64{proto.c64s(i * 2), proto.c64s(i * 2 + 1)}; } - } break; + break; + } + case C128: { + auto complex_data = data(); + TF_RET_CHECK(proto.c128s_size() == complex_data.size() * 2); + for (int64 i = 0; i < complex_data.size(); ++i) { + complex_data[i] = + complex128{proto.c128s(i * 2), proto.c128s(i * 2 + 1)}; + } + break; + } case TUPLE: return InvalidArgument("Should not be called on tuple shapes: %s", ShapeUtil::HumanString(subshape())); diff --git a/tensorflow/compiler/xla/literal.h b/tensorflow/compiler/xla/literal.h index 041151fda1280d6ae7b35d5857ca79788d4f7203..c418be895d6c3faa6a85ca2c73c6f42b0a021104 100644 --- a/tensorflow/compiler/xla/literal.h +++ b/tensorflow/compiler/xla/literal.h @@ -963,6 +963,10 @@ void MutableLiteralBase::AppendSparseElement( CHECK(LayoutUtil::IsSparseArray(subshape)); int64 rank = subshape.rank(); CHECK_EQ(multi_index.size(), rank); + for (int64 i = 0; i < rank; ++i) { + CHECK_GE(multi_index[i], 0); + CHECK_LT(multi_index[i], subshape.dimensions(i)); + } int64 last_element = p.sparse_indices()->index_count(); CHECK_LT(last_element, LayoutUtil::MaxSparseElements(subshape.layout())); p.sparse_indices()->Append(multi_index); diff --git a/tensorflow/compiler/xla/literal_comparison.cc b/tensorflow/compiler/xla/literal_comparison.cc index 258bc966b1a2ee3faa75e3319185489a107b3461..9b3de75dd4e9d495778af86fb8fc07909ab4ba81 100644 --- a/tensorflow/compiler/xla/literal_comparison.cc +++ b/tensorflow/compiler/xla/literal_comparison.cc @@ -90,6 +90,12 @@ bool CompareEqual(complex64 lhs, complex64 rhs, return CompareEqual(lhs.real(), rhs.real(), multi_index) && CompareEqual(lhs.imag(), rhs.imag(), multi_index); } +template <> +bool CompareEqual(complex128 lhs, complex128 rhs, + absl::Span multi_index) { + return CompareEqual(lhs.real(), rhs.real(), multi_index) && + CompareEqual(lhs.imag(), rhs.imag(), multi_index); +} template Status MakeBitwiseErrorStatus(NativeT lhs, NativeT rhs, @@ -143,6 +149,14 @@ Status MakeErrorStatus(complex64 lhs, complex64 rhs, } return MakeErrorStatus(lhs.imag(), rhs.imag(), multi_index); } +template <> +Status MakeErrorStatus(complex128 lhs, complex128 rhs, + absl::Span multi_index) { + if (!CompareEqual(lhs.real(), rhs.real(), multi_index)) { + return MakeErrorStatus(lhs.real(), rhs.real(), multi_index); + } + return MakeErrorStatus(lhs.imag(), rhs.imag(), multi_index); +} // A recursive function which iterates through every index of expected and // actual literal and compares their values elementwise. Returns true if all @@ -186,39 +200,26 @@ int64 RecursiveElementCount(const Shape& shape) { } } -// Returns whether the actual and expected values are mismatched with respect to -// nans. 'relaxed_nans' is interpreted as in xla::ErrorSpec. +// Returns whether the given value is infinity. template -bool NanMismatch(NativeT expected, NativeT actual, bool relaxed_nans) { - if (relaxed_nans) { - return !std::isnan(expected) && std::isnan(actual); - } else { - return std::isnan(expected) != std::isnan(actual); - } -} - -template <> -bool NanMismatch(complex64 expected, complex64 actual, - bool relaxed_nans) { - return NanMismatch(expected.real(), actual.real(), relaxed_nans) || - NanMismatch(expected.imag(), actual.imag(), relaxed_nans); +bool IsInf(NativeT val) { + return std::isinf(val); } template <> -bool NanMismatch(half expected, half actual, bool relaxed_nans) { - return NanMismatch(static_cast(expected), - static_cast(actual), relaxed_nans); +bool IsInf(half val) { + return std::isinf(static_cast(val)); } -// Returns whether the given value is infinity. +// Returns whether the given value is nan. template -bool IsInf(NativeT val) { - return std::isinf(val); +float IsNan(NativeT value) { + return std::isnan(value); } template <> -bool IsInf(half val) { - return std::isinf(static_cast(val)); +float IsNan(half value) { + return IsNan(static_cast(value)); } // Converts the given floating-point value to a string. @@ -232,6 +233,11 @@ string FpValueToString(complex64 value) { return absl::StrFormat("%8.4g + %8.4fi", value.real(), value.imag()); } +template <> +string FpValueToString(complex128 value) { + return absl::StrFormat("%8.4g + %8.4fi", value.real(), value.imag()); +} + // Returns the absolute value of the given floating point value. This function // is used instead of std::abs directly in order to allow type-dependent // implementations for NearComparator. @@ -364,21 +370,39 @@ class NearComparator { // the given literal_index and keeps track of various mismatch statistics. template void CompareValues(T expected, T actual, int64 linear_index) { - const bool is_nan_mismatch = - NanMismatch(expected, actual, error_.relaxed_nans); float abs_error; float rel_error; if (CompareEqual(expected, actual, {linear_index})) { abs_error = 0; rel_error = 0; - } else if (is_nan_mismatch) { - num_nan_mismatches_++; - // A nan mismatch is considered to have infinite error. rel_error is used - // for sorting a std::set of the top mismatchs, and a nan value here will - // result in undefined behavior because nan's do not satisfy the strict - // weak ordering requirement of std containers. - abs_error = std::numeric_limits::infinity(); - rel_error = std::numeric_limits::infinity(); + } else if (IsNan(expected) || IsNan(actual)) { + if ((!error_.relaxed_nans && IsNan(expected) != IsNan(actual)) || + (error_.relaxed_nans && !IsNan(expected) && IsNan(actual))) { + num_nan_mismatches_++; + // A nan mismatch is considered to have infinite error. rel_error is + // used for sorting a std::set of the top mismatchs, and a nan value + // here will result in undefined behavior because nan's do not satisfy + // the strict weak ordering requirement of std containers. + abs_error = std::numeric_limits::infinity(); + rel_error = std::numeric_limits::infinity(); + } else { + abs_error = 0; + rel_error = 0; + } + } else if (IsInf(actual) && !IsInf(expected) && error_.fewer_infs_ok) { + // `fewer_infs_ok` gives us the option of comparing as though `actual` + // were float_max/min rather than inf. + T actual_finite = actual > T{0} ? std::numeric_limits::max() + : std::numeric_limits::lowest(); + abs_error = FpAbsoluteValue(actual_finite - expected); + + // Avoid division by 0 even though it's well-defined because ubsan can be + // configured to treat this as a fatal error. + if (expected != T{0}) { + rel_error = abs_error / FpAbsoluteValue(expected); + } else { + rel_error = std::numeric_limits::infinity(); + } } else if (IsInf(expected) || IsInf(actual)) { // If either the expected or actual value is infinity but not both, // then both absolute and relative error are regarded as inifity. @@ -398,8 +422,7 @@ class NearComparator { } const bool is_abs_mismatch = abs_error > error_.abs; const bool is_rel_mismatch = rel_error > error_.rel; - const bool is_mismatch = - is_nan_mismatch || (is_abs_mismatch && is_rel_mismatch); + const bool is_mismatch = is_abs_mismatch && is_rel_mismatch; // Update the error of the relative bucket only if the *absolute* error // bound is exceeded and vice versa. @@ -434,7 +457,7 @@ class NearComparator { mismatches_.data()[linear_index] = true; } - // For complex64 types, we compare real and imaginary parts individually. + // For complex types, we compare real and imaginary parts individually. void CompareValues(complex64 expected, complex64 actual, int64 linear_index) { bool mismatch = false; CompareValues(expected.real(), actual.real(), linear_index); @@ -457,6 +480,29 @@ class NearComparator { mismatches_.data()[linear_index] = mismatch; } + void CompareValues(complex128 expected, complex128 actual, + int64 linear_index) { + bool mismatch = false; + CompareValues(expected.real(), actual.real(), linear_index); + if (mismatches_.data()[linear_index] == true) { + mismatch = true; + // Delay the mismatch count increase for real part, instead increase + // mismatch by 1 for the entire complex number. + num_mismatches_--; + } + CompareValues(expected.imag(), actual.imag(), linear_index); + if (mismatches_.data()[linear_index] == true) { + mismatch = true; + // Delay the mismatch count increase for imag part, instead increase + // mismatch by 1 for the entire complex number. + num_mismatches_--; + } + if (mismatch == true) { + num_mismatches_++; + } + mismatches_.data()[linear_index] = mismatch; + } + // Compares the two literals elementwise. void CompareLiterals() { // Fast path optimization for the case were layouts match. @@ -665,6 +711,9 @@ Status EqualHelper(const LiteralSlice& expected, const LiteralSlice& actual) { case C64: result = Equal(expected, actual, index, 0); break; + case C128: + result = Equal(expected, actual, index, 0); + break; case TUPLE: { for (int i = 0; i < ShapeUtil::TupleElementCount(expected.shape()); ++i) { result.Update(EqualHelper(LiteralSlice(expected, {i}), @@ -687,7 +736,7 @@ Status EqualHelper(const LiteralSlice& expected, const LiteralSlice& actual) { // via recursion. shape_index is the ShapeIndex of expected (or actual) // currently being compared. Status NearHelper(const LiteralSlice& expected, const LiteralSlice& actual, - const ErrorSpec& error, bool detailed_message, + const ErrorSpec& error, absl::optional detailed_message, const MiscompareCallback& miscompare_callback, const ShapeIndex& shape_index) { TF_RETURN_IF_ERROR(EqualShapes(expected.shape(), actual.shape())); @@ -728,26 +777,32 @@ Status NearHelper(const LiteralSlice& expected, const LiteralSlice& actual, if (ShapeUtil::ElementIsFloating(expected.shape()) || ShapeUtil::ElementIsComplex(expected.shape())) { + bool use_detailed_message = detailed_message.value_or( + ShapeUtil::ElementsIn(expected.shape()) >= 64); switch (expected.shape().element_type()) { case BF16: return NearComparator::Compare( - expected, actual, error, detailed_message, miscompare_callback); + expected, actual, error, use_detailed_message, miscompare_callback); break; case F16: return NearComparator::Compare( - expected, actual, error, detailed_message, miscompare_callback); + expected, actual, error, use_detailed_message, miscompare_callback); break; case F32: return NearComparator::Compare( - expected, actual, error, detailed_message, miscompare_callback); + expected, actual, error, use_detailed_message, miscompare_callback); break; case F64: return NearComparator::Compare( - expected, actual, error, detailed_message, miscompare_callback); + expected, actual, error, use_detailed_message, miscompare_callback); break; case C64: return NearComparator::Compare( - expected, actual, error, detailed_message, miscompare_callback); + expected, actual, error, use_detailed_message, miscompare_callback); + break; + case C128: + return NearComparator::Compare( + expected, actual, error, use_detailed_message, miscompare_callback); break; default: LOG(FATAL) << "Unsupported primitive type in near comparator: " @@ -838,7 +893,7 @@ Status Equal(const LiteralSlice& expected, const LiteralSlice& actual) { } Status Near(const LiteralSlice& expected, const LiteralSlice& actual, - const ErrorSpec& error, bool detailed_message, + const ErrorSpec& error, absl::optional detailed_message, const MiscompareCallback& miscompare_callback) { VLOG(1) << "Expected literal:"; XLA_VLOG_LINES(1, expected.ToString()); diff --git a/tensorflow/compiler/xla/literal_comparison.h b/tensorflow/compiler/xla/literal_comparison.h index 9e5bf7c1d062ef0f25d07a80d6ded8106df5dacc..23fff3fa348f1652eaec344da4c40ccf3ad1079a 100644 --- a/tensorflow/compiler/xla/literal_comparison.h +++ b/tensorflow/compiler/xla/literal_comparison.h @@ -55,9 +55,10 @@ using MiscompareCallback = // being compared. // // If detailed_message is true, then the error message in the assertion result -// will contain a more detailed breakdown of mismatches. +// will contain a more detailed breakdown of mismatches. By default, we display +// a detailed message only for "large" inputs. Status Near(const LiteralSlice& expected, const LiteralSlice& actual, - const ErrorSpec& error, bool detailed_message, + const ErrorSpec& error, absl::optional detailed_message, const MiscompareCallback& miscompare_callback); // Calling ToString on a literal with over 100 million elements takes around diff --git a/tensorflow/compiler/xla/literal_test.cc b/tensorflow/compiler/xla/literal_test.cc index 0a71cfd5abf8511af9680f5537c05f5eb833397c..b54a71ae68218ef578535a913f5867d843236e32 100644 --- a/tensorflow/compiler/xla/literal_test.cc +++ b/tensorflow/compiler/xla/literal_test.cc @@ -118,6 +118,9 @@ TEST_F(LiteralUtilTest, LiteralScalarToString) { auto c64_lit = LiteralUtil::CreateR0({3.14f, 2.78f}); EXPECT_EQ("c64[] (3.14, 2.78)", c64_lit.ToString()); + auto c128_lit = LiteralUtil::CreateR0({3.14f, 2.78f}); + EXPECT_EQ("c128[] (3.14, 2.78)", c128_lit.ToString()); + auto bf16_lit = LiteralUtil::CreateR0(static_cast(0.5f)); EXPECT_EQ("bf16[] 0.5", bf16_lit.ToString()); @@ -469,6 +472,21 @@ TEST_F(LiteralUtilTest, C64Equality) { EXPECT_NE(vector, vector_reversed); } +TEST_F(LiteralUtilTest, C128Equality) { + // Test equality with tuples. + auto vector = LiteralUtil::CreateR1({{1.0, 2.0}, {3.0, 4.0}}); + + // Tuple with the same elements. One element is shared with the original + // tuple, the other is a clone of the element in the original tuple. + auto vector_clone = + LiteralUtil::CreateR1({{1.0, 2.0}, {3.0, 4.0}}); + EXPECT_EQ(vector, vector_clone); + + auto vector_reversed = + LiteralUtil::CreateR1({{3.0, 4.0}, {1.0, 2.0}}); + EXPECT_NE(vector, vector_reversed); +} + TEST_F(LiteralUtilTest, IsAllTuple) { auto element1 = LiteralUtil::CreateR0(0.0); auto element2 = LiteralUtil::CreateR2({{0.0, 0.0}, {0.0, 0.0}}); @@ -623,7 +641,7 @@ template class LiteralUtilTestTemplated : public ::testing::Test {}; using TestedTypes = ::testing::Types; -TYPED_TEST_CASE(LiteralUtilTestTemplated, TestedTypes); +TYPED_TEST_SUITE(LiteralUtilTestTemplated, TestedTypes); TYPED_TEST(LiteralUtilTestTemplated, Relayout2x2) { // Make a non-integer for floating point types. @@ -836,6 +854,13 @@ TEST_F(LiteralUtilTest, PopulateR1C64) { EXPECT_EQ(output, expected); } +TEST_F(LiteralUtilTest, PopulateR1C128) { + Literal output(ShapeUtil::MakeShape(C128, {1})); + output.PopulateR1({{77, 88}}); + auto expected = LiteralUtil::CreateR1({{77, 88}}); + EXPECT_EQ(output, expected); +} + TEST_F(LiteralUtilTest, PopulateR2C64) { Literal output(ShapeUtil::MakeShape(C64, {2, 2})); output.PopulateR2({{{7, 8}, {9, 10}}, {{1, 2}, {3, 4}}}); @@ -897,6 +922,14 @@ TEST_F(LiteralUtilTest, PopulateWithValueR2C64) { EXPECT_EQ(output, expected); } +TEST_F(LiteralUtilTest, PopulateWithValueR2C128) { + Literal output(ShapeUtil::MakeShape(C128, {2, 2})); + output.PopulateWithValue({4, 2}); + auto expected = + LiteralUtil::CreateR2({{{4, 2}, {4, 2}}, {{4, 2}, {4, 2}}}); + EXPECT_EQ(output, expected); +} + TEST_F(LiteralUtilTest, PopulateWithValueR0F16) { Literal output(ShapeUtil::MakeShape(F16, {})); half h(0.25f); @@ -1237,11 +1270,21 @@ TEST_F(LiteralUtilTest, ConvertIfTypesMatch) { {{0, 19, 0, 21}, {22, 0, 24, 0}}, {{26, 0, 28, 0}, {0, 31, 0, 33}}, }}, layout_r4_dim0major_); + auto s16 = LiteralUtil::CreateR4WithLayout({{ + {{10, 0, 12, 0}, {0, 15, 0, 17}}, + {{0, 19, 0, 21}, {22, 0, 24, 0}}, + {{26, 0, 28, 0}, {0, 31, 0, 33}}, + }}, layout_r4_dim0major_); auto s32 = LiteralUtil::CreateR4WithLayout({{ {{10, 0, 12, 0}, {0, 15, 0, 17}}, {{0, 19, 0, 21}, {22, 0, 24, 0}}, {{26, 0, 28, 0}, {0, 31, 0, 33}}, }}, layout_r4_dim0major_); + auto u16 = LiteralUtil::CreateR4WithLayout({{ + {{10, 0, 12, 0}, {0, 15, 0, 17}}, + {{0, 19, 0, 21}, {22, 0, 24, 0}}, + {{26, 0, 28, 0}, {0, 31, 0, 33}}, + }}, layout_r4_dim0major_); auto u32 = LiteralUtil::CreateR4WithLayout({{ {{10, 0, 12, 0}, {0, 15, 0, 17}}, {{0, 19, 0, 21}, {22, 0, 24, 0}}, @@ -1298,9 +1341,19 @@ TEST_F(LiteralUtilTest, ConvertIfTypesMatch) { {{0.0f, 19.0f, 0.0f, 21.0f}, {22.0f, 0.0f, 24.0f, 0.0f}}, {{26.0f, 0.0f, 28.0f, 0.0f}, {0.0f, 31.0f, 0.0f, 33.0f}}, }}, layout_r4_dim0major_); - // clang-format on + auto c128 = LiteralUtil::CreateR4WithLayout({{ + {{10.0, 0.0, 12.0, 0.0}, {0.0, 15.0, 0.0, 17.0}}, + {{0.0, 19.0, 0.0, 21.0}, {22.0, 0.0, 24.0, 0.0}}, + {{26.0, 0.0, 28.0, 0.0}, {0.0, 31.0, 0.0, 33.0}}, + }}, layout_r4_dim0major_); // clang-format on Literal conv; + conv = s8.Convert(U16).ConsumeValueOrDie(); + EXPECT_EQ(conv, u16); + + conv = s8.Convert(S16).ConsumeValueOrDie(); + EXPECT_EQ(conv, s16); + conv = s8.Convert(U32).ConsumeValueOrDie(); EXPECT_EQ(conv, u32); @@ -1352,12 +1405,26 @@ TEST_F(LiteralUtilTest, ConvertIfTypesMatch) { conv = f16.Convert(C64).ConsumeValueOrDie(); EXPECT_EQ(conv, c64); + conv = s32.Convert(S16).ConsumeValueOrDie(); + EXPECT_EQ(conv, s16); + + conv = s32.Convert(U16).ConsumeValueOrDie(); + EXPECT_EQ(conv, u16); + + conv = s32.Convert(C128).ConsumeValueOrDie(); + EXPECT_EQ(conv, c128); + + conv = f16.Convert(C128).ConsumeValueOrDie(); + EXPECT_EQ(conv, c128); + EXPECT_EQ(s32.Convert(TUPLE).status().code(), tensorflow::error::UNIMPLEMENTED); - EXPECT_EQ(s32.Convert(S16).status().code(), tensorflow::error::UNIMPLEMENTED); - EXPECT_EQ(s32.Convert(U16).status().code(), tensorflow::error::UNIMPLEMENTED); EXPECT_EQ(c64.Convert(F32).status().code(), tensorflow::error::UNIMPLEMENTED); EXPECT_EQ(c64.Convert(S32).status().code(), tensorflow::error::UNIMPLEMENTED); + EXPECT_EQ(c128.Convert(F32).status().code(), + tensorflow::error::UNIMPLEMENTED); + EXPECT_EQ(c128.Convert(S32).status().code(), + tensorflow::error::UNIMPLEMENTED); } TEST_F(LiteralUtilTest, BitcastConvert) { @@ -1719,7 +1786,8 @@ TEST_F(LiteralUtilTest, CreateFromShapeZeroInitialized) { Literal tuple = Literal::CreateFromShape(ShapeUtil::MakeTupleShape( {ShapeUtil::MakeShape(F64, {}), ShapeUtil::MakeShape(PRED, {2}), - ShapeUtil::MakeShape(U64, {2, 1}), ShapeUtil::MakeShape(C64, {})})); + ShapeUtil::MakeShape(U64, {2, 1}), ShapeUtil::MakeShape(C64, {}), + ShapeUtil::MakeShape(C128, {})})); EXPECT_EQ(tuple.Get({}, {0}), 0.0); EXPECT_EQ(tuple.Get({0}, {1}), false); @@ -1727,6 +1795,7 @@ TEST_F(LiteralUtilTest, CreateFromShapeZeroInitialized) { EXPECT_EQ(tuple.Get({0, 0}, {2}), 0); EXPECT_EQ(tuple.Get({1, 0}, {2}), 0); EXPECT_EQ(tuple.Get({}, {3}), complex64(0.0f, 0.0f)); + EXPECT_EQ(tuple.Get({}, {4}), complex128(0.0, 0.0)); } TEST_F(LiteralUtilTest, ProtoRoundTrip) { @@ -1736,6 +1805,8 @@ TEST_F(LiteralUtilTest, ProtoRoundTrip) { auto vector_int8 = LiteralUtil::CreateR1({-128, 0, 2, 4, 7, 56, 127}); auto vector_uint8 = LiteralUtil::CreateR1({128, 0, 2, 56, 127, 255}); auto vector_c64 = LiteralUtil::CreateR1({{1.0, 2.0}, {3.0, 4.0}}); + auto vector_c128 = + LiteralUtil::CreateR1({{1.0, 2.0}, {3.0, 4.0}}); auto vector_bfloat16 = LiteralUtil::CreateR1( {bfloat16{-1.0}, bfloat16{2.0}, bfloat16{-3.0}}); auto vector_half = @@ -1756,6 +1827,7 @@ TEST_F(LiteralUtilTest, ProtoRoundTrip) { EXPECT_EQ(vector_int8, to_from_proto(vector_int8)); EXPECT_EQ(vector_uint8, to_from_proto(vector_uint8)); EXPECT_EQ(vector_c64, to_from_proto(vector_c64)); + EXPECT_EQ(vector_c128, to_from_proto(vector_c128)); EXPECT_EQ(vector_bfloat16, to_from_proto(vector_bfloat16)); EXPECT_EQ(matrix_pred, to_from_proto(matrix_pred)); EXPECT_EQ(tuple, to_from_proto(tuple)); diff --git a/tensorflow/compiler/xla/literal_util.cc b/tensorflow/compiler/xla/literal_util.cc index 67eab1bff8a773419ad1bf038295e2bfbd710145..26b029c8d0c52e38510f9279def7c4af2904931d 100644 --- a/tensorflow/compiler/xla/literal_util.cc +++ b/tensorflow/compiler/xla/literal_util.cc @@ -106,12 +106,16 @@ Literal ConvertType(LiteralSlice literal) { switch (primitive_type) { case U8: return LiteralUtil::CreateR0(0); + case U16: + return LiteralUtil::CreateR0(0); case U32: return LiteralUtil::CreateR0(0); case U64: return LiteralUtil::CreateR0(0); case S8: return LiteralUtil::CreateR0(0); + case S16: + return LiteralUtil::CreateR0(0); case S32: return LiteralUtil::CreateR0(0); case S64: @@ -126,11 +130,10 @@ Literal ConvertType(LiteralSlice literal) { return LiteralUtil::CreateR0(0); case C64: return LiteralUtil::CreateR0(0); + case C128: + return LiteralUtil::CreateR0(0); case PRED: return LiteralUtil::CreateR0(false); - case S16: - case U16: - LOG(FATAL) << "u16/s16 literals not yet implemented"; case TUPLE: LOG(FATAL) << "tuple element type cannot take on value of 0"; case OPAQUE: @@ -164,6 +167,8 @@ Literal ConvertType(LiteralSlice literal) { return LiteralUtil::CreateR0(1); case C64: return LiteralUtil::CreateR0(1); + case C128: + return LiteralUtil::CreateR0(1); case PRED: return LiteralUtil::CreateR0(true); case S16: @@ -200,6 +205,8 @@ Literal ConvertType(LiteralSlice literal) { -std::numeric_limits::infinity()); case C64: LOG(FATAL) << "C64 element type has no minimum value"; + case C128: + LOG(FATAL) << "C128 element type has no minimum value"; case PRED: return LiteralUtil::CreateR0(false); case S16: @@ -344,6 +351,10 @@ Literal ConvertType(LiteralSlice literal) { new_literal.Set(to_multi_index, literal.Get(from_multi_index)); break; + case C128: + new_literal.Set(to_multi_index, + literal.Get(from_multi_index)); + break; default: LOG(FATAL) << "Unhandled primitive element type: " << PrimitiveType_Name(literal.shape().element_type()); @@ -392,6 +403,10 @@ Literal ConvertType(LiteralSlice literal) { return LiteralUtil::CreateR0(literal.GetFirstElement()); case U64: return LiteralUtil::CreateR0(literal.GetFirstElement()); + + case C128: + return LiteralUtil::CreateR0( + literal.GetFirstElement()); default: LOG(FATAL) << "Unhandled primitive type " << literal.shape().element_type(); diff --git a/tensorflow/compiler/xla/metric_table_report.cc b/tensorflow/compiler/xla/metric_table_report.cc index 4eab4fa4290c270697c00be20840cf4e85459183..bad65ac32018fafcc7634b989f1b4b0867aa5c0d 100644 --- a/tensorflow/compiler/xla/metric_table_report.cc +++ b/tensorflow/compiler/xla/metric_table_report.cc @@ -15,9 +15,9 @@ limitations under the License. #include "tensorflow/compiler/xla/metric_table_report.h" -#include #include +#include "absl/strings/ascii.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "tensorflow/core/platform/logging.h" @@ -55,7 +55,7 @@ string MetricTableReport::MakeReport(double expected_metric_sum) { const auto metric_greater = [](const Entry& a, const Entry& b) { return a.metric > b.metric; }; - std::sort(entries_.begin(), entries_.end(), metric_greater); + absl::c_sort(entries_, metric_greater); // Create the report AppendLine(); @@ -117,7 +117,7 @@ std::vector MetricTableReport::MakeCategories( auto metric_sum_greater = [](const Category& a, const Category& b) { return a.metric_sum > b.metric_sum; }; - std::sort(categories.begin(), categories.end(), metric_sum_greater); + absl::c_sort(categories, metric_sum_greater); return categories; } @@ -249,7 +249,7 @@ string MetricTableReport::MetricString(double metric) { string output; // Copy leading non-digit characters unconditionally. // This picks up the leading sign. - while (!sp1.empty() && !isdigit(sp1[0])) { + while (!sp1.empty() && !absl::ascii_isdigit(sp1[0])) { output.push_back(sp1[0]); sp1.remove_prefix(1); } diff --git a/tensorflow/compiler/xla/parse_flags_from_env.cc b/tensorflow/compiler/xla/parse_flags_from_env.cc index 5b568888d14f21c1330556d017eafba6c8dd2228..e1e22f784172b5f3850f0bc510322dfad9e7f1bb 100644 --- a/tensorflow/compiler/xla/parse_flags_from_env.cc +++ b/tensorflow/compiler/xla/parse_flags_from_env.cc @@ -24,6 +24,7 @@ limitations under the License. #include #include +#include "absl/strings/ascii.h" #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" #include "absl/types/span.h" @@ -37,7 +38,7 @@ limitations under the License. namespace xla { -static const char kWS[] = " \t\r\n"; // whitespace +static const char kWS[] = " \t\r\n"; // whitespace // The following struct represents an argv[]-style array, parsed // from data gleaned from the environment. @@ -104,7 +105,8 @@ static void ParseArgvFromString(const string& flag_str, EnvArgv* a) { // Set e to the index just past the end of the flag. size_t e = b; while (e != flag_str.size() && isascii(flag_str[e]) && - (strchr("-_", flag_str[e]) != nullptr || isalnum(flag_str[e]))) { + (strchr("-_", flag_str[e]) != nullptr || + absl::ascii_isalnum(flag_str[e]))) { e++; } if (e != flag_str.size() && flag_str[e] == '=' && @@ -184,6 +186,14 @@ bool ParseFlagsFromEnvAndDieIfUnknown( tensorflow::mutex_lock lock(env_argv_mu); auto* env_argv = &EnvArgvs()[string(envvar)]; SetArgvFromEnv(envvar, env_argv); // a no-op if already initialized + + if (VLOG_IS_ON(1)) { + VLOG(1) << "For env var " << envvar << " found arguments:"; + for (int i = 0; i < env_argv->argc; i++) { + VLOG(1) << " argv[" << i << "] = " << env_argv->argv[i]; + } + } + bool result = tensorflow::Flags::Parse(&env_argv->argc, &env_argv->argv[0], flag_list); diff --git a/tensorflow/compiler/xla/primitive_util.cc b/tensorflow/compiler/xla/primitive_util.cc index 00ad01fc407017624a9183d69e61cb0d382e3f11..1eedddf72c1d393cb1b88e589881e24de02ad802 100644 --- a/tensorflow/compiler/xla/primitive_util.cc +++ b/tensorflow/compiler/xla/primitive_util.cc @@ -18,16 +18,32 @@ limitations under the License. #include "absl/strings/ascii.h" #include "absl/strings/numbers.h" #include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/logging.h" namespace xla { namespace primitive_util { +int SignificandWidth(PrimitiveType type) { + switch (type) { + case F32: + return std::numeric_limits::digits; + case F64: + return std::numeric_limits::digits; + case BF16: + return kBFloat16MantissaBits + 1; + case F16: + return 11; + default: + LOG(FATAL) << "Not a floating data type " << type; + } +} + bool IsFloatingPointType(PrimitiveType type) { return type == F16 || type == F32 || type == F64 || type == BF16; } -bool IsComplexType(PrimitiveType type) { return type == C64; } +bool IsComplexType(PrimitiveType type) { return type == C64 || type == C128; } bool IsSignedIntegralType(PrimitiveType type) { return type == S8 || type == S16 || type == S32 || type == S64; @@ -67,6 +83,9 @@ int BitWidth(PrimitiveType type) { case C64: return 64; + case C128: + return 128; + case TUPLE: LOG(FATAL) << "TUPLE is an invalid type for BitWidth"; @@ -78,10 +97,27 @@ int BitWidth(PrimitiveType type) { } } +xla::PrimitiveType UnsignedIntegralTypeForBitWidth(int64 src_bitwidth) { + switch (src_bitwidth) { + case 8: + return xla::U8; + case 16: + return xla::U16; + case 32: + return xla::U32; + case 64: + return xla::U64; + default: + return xla::PRIMITIVE_TYPE_INVALID; + } +} + PrimitiveType ComplexComponentType(PrimitiveType complex_type) { switch (complex_type) { case C64: return F32; + case C128: + return F64; default: LOG(FATAL) << "Primitive type is not complex: " << PrimitiveType_Name(complex_type); diff --git a/tensorflow/compiler/xla/primitive_util.h b/tensorflow/compiler/xla/primitive_util.h index 70603b6fed1be50c427799e6dce7b8bf9631a6f4..295d353003276b4c1731f7d6a378fd1ae0288d3c 100644 --- a/tensorflow/compiler/xla/primitive_util.h +++ b/tensorflow/compiler/xla/primitive_util.h @@ -29,6 +29,10 @@ limitations under the License. namespace xla { namespace primitive_util { +// Returns the count of significand (mantissa) bits for float datatypes. +// For non-float datatypes, results in a LOG(FATAL). +int SignificandWidth(PrimitiveType type); + // The number of exponent bits in a BF16 value. const int kBFloat16ExponentBits = 8; @@ -126,6 +130,11 @@ inline PrimitiveType NativeToPrimitiveType() { return C64; } +template <> +inline PrimitiveType NativeToPrimitiveType() { + return C128; +} + bool IsFloatingPointType(PrimitiveType type); bool IsComplexType(PrimitiveType type); @@ -142,6 +151,8 @@ bool IsArrayType(PrimitiveType primitive_type); // Returns the number of bits in the representation for a given type. int BitWidth(PrimitiveType type); +PrimitiveType UnsignedIntegralTypeForBitWidth(int64 src_bitwidth); + // Returns the real, imag component type underlying the given complex type. // LOG(FATAL)'s if complex_type is not complex. PrimitiveType ComplexComponentType(PrimitiveType complex_type); @@ -225,6 +236,11 @@ struct PrimitiveTypeToNative { using type = complex64; }; +template <> +struct PrimitiveTypeToNative { + using type = complex128; +}; + // Returns the lower-case name of the given primitive type. const string& LowercasePrimitiveTypeName(PrimitiveType s); diff --git a/tensorflow/compiler/xla/protobuf_util.h b/tensorflow/compiler/xla/protobuf_util.h index f22fc8b8499dd4a5329276040331a2ed9e89bea9..4a88a48f2857a327aba3600ca72191e5c7b28585 100644 --- a/tensorflow/compiler/xla/protobuf_util.h +++ b/tensorflow/compiler/xla/protobuf_util.h @@ -16,6 +16,8 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_PROTOBUF_UTIL_H_ #define TENSORFLOW_COMPILER_XLA_PROTOBUF_UTIL_H_ +#include "google/protobuf/duration.pb.h" +#include "absl/time/time.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/platform/protobuf.h" @@ -43,6 +45,20 @@ Status DumpProtoToDirectory(const tensorflow::protobuf::Message& message, // dirpath along as-is. void RegisterDirectoryExpander(const std::function& expander); +// Converts an absl::Duration to a google::protobuf::Duration. +inline google::protobuf::Duration ToDurationProto(absl::Duration duration) { + google::protobuf::Duration proto; + proto.set_seconds(absl::IDivDuration(duration, absl::Seconds(1), &duration)); + proto.set_nanos( + absl::IDivDuration(duration, absl::Nanoseconds(1), &duration)); + return proto; +} + +// Converts a google::protobuf::Duration to an absl::Duration. +inline absl::Duration FromDurationProto(google::protobuf::Duration proto) { + return absl::Seconds(proto.seconds()) + absl::Nanoseconds(proto.nanos()); +} + } // namespace protobuf_util } // namespace xla diff --git a/tensorflow/compiler/xla/python/BUILD b/tensorflow/compiler/xla/python/BUILD index ddffafa9017a565f01c3214360a958e6840e9148..f7e2d26b7aa7f23f5e0a2e7623863402ef549789 100644 --- a/tensorflow/compiler/xla/python/BUILD +++ b/tensorflow/compiler/xla/python/BUILD @@ -3,8 +3,8 @@ licenses(["notice"]) # Apache 2.0 package(default_visibility = ["//tensorflow:internal"]) load("//tensorflow:tensorflow.bzl", "tf_py_wrap_cc") -load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured") load("//tensorflow/core:platform/default/build_config.bzl", "pyx_library") +load("//tensorflow/compiler/xla:xla.bzl", "xla_python_default_plugins") py_library( name = "xla_client", @@ -77,7 +77,6 @@ cc_library( "//tensorflow/compiler/xla/client/lib:cholesky", "//tensorflow/compiler/xla/client/lib:math", "//tensorflow/compiler/xla/client/lib:qr", - "//tensorflow/compiler/xla/client/lib:triangular_solve", "//tensorflow/compiler/xla/service:platform_util", "//tensorflow/compiler/xla/service:shaped_buffer", "//tensorflow/compiler/xla/service/cpu:custom_call_target_registry", @@ -98,6 +97,11 @@ tf_py_wrap_cc( "local_computation_builder.i", "//tensorflow/python:platform/base.i", ], + version_script = select({ + "//tensorflow:darwin": "pywrap_xla_exported_symbols.lds", + "//tensorflow:windows": None, + "//conditions:default": "pywrap_xla_version_script.lds", + }), deps = [ ":local_computation_builder", ":numpy_bridge", @@ -105,7 +109,5 @@ tf_py_wrap_cc( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service:cpu_plugin", - ] + if_cuda_is_configured([ - "//tensorflow/compiler/xla/service:gpu_plugin", - ]), + ] + xla_python_default_plugins(), ) diff --git a/tensorflow/compiler/xla/python/local_computation_builder.cc b/tensorflow/compiler/xla/python/local_computation_builder.cc index b0f0e9e57060f43e464ea2b2e20b9000679d18b6..24138a173de787122ce7b372c03c8fc3703dc118 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.cc +++ b/tensorflow/compiler/xla/python/local_computation_builder.cc @@ -27,7 +27,6 @@ limitations under the License. #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/lib/triangular_solve.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/client/xla_computation.h" #include "tensorflow/compiler/xla/executable_run_options.h" @@ -92,7 +91,12 @@ Status InitializePlatformName(const string& platform_name) { "previously created with a platform name of %s.", platform_name, *g_platform_name); } - TF_RETURN_IF_ERROR(PlatformUtil::GetPlatform(platform_name).status()); + TF_ASSIGN_OR_RETURN(se::Platform * platform, + PlatformUtil::GetPlatform(platform_name)); + if (platform->VisibleDeviceCount() <= 0) { + return InvalidArgument("Platform %s has no visible devices.", + platform_name); + } *g_platform_name = platform_name; return Status::OK(); } @@ -102,7 +106,7 @@ int GetReplicaCount() { return g_replica_count; } -LocalClient* GetOrCreateLocalClient() { +StatusOr GetOrCreateLocalClient() { string* platform_name = GetPlatformNameString(); tensorflow::mutex_lock lock(g_local_client_mutex); if (g_local_client != nullptr) { @@ -111,7 +115,8 @@ LocalClient* GetOrCreateLocalClient() { LocalClientOptions options; options.set_platform(PlatformUtil::GetPlatform(*platform_name).ValueOrDie()); options.set_number_of_replicas(g_replica_count); - g_local_client = ClientLibrary::GetOrCreateLocalClient(options).ValueOrDie(); + TF_ASSIGN_OR_RETURN(g_local_client, + ClientLibrary::GetOrCreateLocalClient(options)); CHECK(g_local_client != nullptr); return g_local_client; } @@ -133,7 +138,7 @@ Status RegisterCpuCustomCallTarget(const string& fn_name, PyObject* capsule) { Status TransferToInfeedLocal(const Literal& literal) { VLOG(1) << "Infeeding literal without replica number; shape: " << literal.shape(); - LocalClient* client = GetOrCreateLocalClient(); + TF_ASSIGN_OR_RETURN(LocalClient * client, GetOrCreateLocalClient()); return client->TransferToInfeedLocal(literal, /*device_ordinal=*/0); } @@ -141,7 +146,7 @@ Status TransferToInfeedLocalReplica(const Literal& literal, int replica_number) { VLOG(1) << "Infeeding shape " << literal.shape() << " to replica number: " << replica_number; - LocalClient* client = GetOrCreateLocalClient(); + TF_ASSIGN_OR_RETURN(LocalClient * client, GetOrCreateLocalClient()); TF_ASSIGN_OR_RETURN(int device_ordinal, client->ReplicaNumberToDeviceOrdinal(replica_number)); return client->TransferToInfeedLocal(literal, device_ordinal); @@ -151,7 +156,7 @@ StatusOr TransferFromOutfeedLocalReplica(const Shape& shape, int replica_number) { VLOG(1) << "Outfeeding literal from replica number: " << replica_number << " shape: " << shape; - LocalClient* client = GetOrCreateLocalClient(); + TF_ASSIGN_OR_RETURN(LocalClient * client, GetOrCreateLocalClient()); TF_ASSIGN_OR_RETURN(int device_ordinal, client->ReplicaNumberToDeviceOrdinal(replica_number)); return client->TransferFromOutfeedLocal(shape, device_ordinal); @@ -168,7 +173,7 @@ static StatusOr ToBuffer(LocalClient* client, StatusOr LocalShapedBuffer::FromLiteral( const Literal& argument, const absl::optional& shape_with_layout, int replica_number) { - LocalClient* client = GetOrCreateLocalClient(); + TF_ASSIGN_OR_RETURN(LocalClient * client, GetOrCreateLocalClient()); TF_ASSIGN_OR_RETURN(int device_ordinal, client->ReplicaNumberToDeviceOrdinal(replica_number)); VLOG(1) << "Creating shaped buffer from literal on replica/ordinal: " @@ -198,7 +203,7 @@ const Shape& LocalShapedBuffer::shape() const { } StatusOr LocalShapedBuffer::ToLiteral() const { - LocalClient* client = GetOrCreateLocalClient(); + TF_ASSIGN_OR_RETURN(LocalClient * client, GetOrCreateLocalClient()); return client->ShapedBufferToLiteral(*shaped_buffer()); } @@ -333,37 +338,34 @@ CompiledLocalComputation::CompiledLocalComputation( StatusOr CompiledLocalComputation::Execute( absl::Span argument_handles) { - LocalClient* client = GetOrCreateLocalClient(); - StatusOr device_ordinal_status = client->ReplicaNumberToDeviceOrdinal(0); + if (num_replicas() != 1) { + return InvalidArgument( + "Attempted to execute computation with %d replicas using Execute()", + num_replicas()); + } + TF_ASSIGN_OR_RETURN(LocalClient * client, GetOrCreateLocalClient()); + TF_ASSIGN_OR_RETURN(DeviceAssignment device_assignment, + client->backend().computation_placer()->AssignDevices( + 1, /*computation_count=*/1)); StatusOr result_buffer_status; - if (!device_ordinal_status.ok()) { - result_buffer_status = device_ordinal_status.status(); - } else { - const int device_ordinal = device_ordinal_status.ValueOrDie(); - VLOG(3) << "Replica 0 mapped to device ordinal for execution: " - << device_ordinal; - - std::vector argument_buffers; - argument_buffers.reserve(argument_handles.size()); - for (auto& handle : argument_handles) { - argument_buffers.push_back(handle->shaped_buffer()); - } - - DeviceAssignment device_assignment = - client->backend() - .computation_placer() - ->AssignDevices(1, /*computation_count=*/1) - .ConsumeValueOrDie(); + const int device_ordinal = device_assignment(0, 0); + VLOG(3) << "Replica 0 mapped to device ordinal for execution: " + << device_ordinal; + + std::vector argument_buffers; + argument_buffers.reserve(argument_handles.size()); + for (auto& handle : argument_handles) { + argument_buffers.push_back(handle->shaped_buffer()); + } - ExecutableRunOptions options; - options.set_device_ordinal(device_ordinal); - options.set_allocator(client->backend().memory_allocator()); - options.set_intra_op_thread_pool( - client->backend().eigen_intra_op_thread_pool_device()); - options.set_device_assignment(&device_assignment); + ExecutableRunOptions options; + options.set_device_ordinal(device_ordinal); + options.set_allocator(client->backend().memory_allocator()); + options.set_intra_op_thread_pool( + client->backend().eigen_intra_op_thread_pool_device()); + options.set_device_assignment(&device_assignment); - result_buffer_status = executable_->Run(argument_buffers, options); - } + result_buffer_status = executable_->Run(argument_buffers, options); if (!result_buffer_status.ok()) { return InternalError( @@ -376,29 +378,30 @@ StatusOr CompiledLocalComputation::Execute( StatusOr CompiledLocalComputation::ExecutePerReplica( absl::Span> argument_handles) { - LocalClient* client = GetOrCreateLocalClient(); - const int num_replicas = GetReplicaCount(); + TF_ASSIGN_OR_RETURN(LocalClient * client, GetOrCreateLocalClient()); + const int num_devices = client->device_count(); - if (argument_handles.size() != num_replicas) { + if (argument_handles.size() != num_replicas()) { return InvalidArgument( "Attempted to execute with %d replicas when replica count is %d", - argument_handles.size(), num_replicas); + argument_handles.size(), num_devices); + } + if (argument_handles.size() > num_devices) { + return InvalidArgument( + "Attempted to execute with %d replicas when device count is %d", + argument_handles.size(), num_devices); } - VLOG(1) << "Executing with " << num_replicas << " replicas."; + VLOG(1) << "Executing with " << num_replicas() << " replicas."; - // Each replica populates a StatusOr result, but only the output value of - // replica zero is returned. - std::vector> results(num_replicas); - auto execute = [this, client, num_replicas, &argument_handles, + TF_ASSIGN_OR_RETURN(DeviceAssignment device_assignment, + client->backend().computation_placer()->AssignDevices( + num_replicas(), /*computation_count=*/1)); + + std::vector> results(num_replicas()); + auto execute = [this, client, &device_assignment, &argument_handles, &results](int replica) { - StatusOr device_ordinal_status = - client->ReplicaNumberToDeviceOrdinal(replica); - if (!device_ordinal_status.ok()) { - results[replica] = device_ordinal_status.status(); - return; - } - const int device_ordinal = device_ordinal_status.ValueOrDie(); + const int device_ordinal = device_assignment(replica, 0); VLOG(3) << "Replica " << replica << " mapped to device ordinal for execution: " << device_ordinal; @@ -408,12 +411,6 @@ StatusOr CompiledLocalComputation::ExecutePerReplica( argument_buffers.push_back(handle->shaped_buffer()); } - DeviceAssignment device_assignment = - client->backend() - .computation_placer() - ->AssignDevices(num_replicas, /*computation_count=*/1) - .ConsumeValueOrDie(); - ExecutableRunOptions options; options.set_device_ordinal(device_ordinal); options.set_allocator(client->backend().memory_allocator()); @@ -426,23 +423,23 @@ StatusOr CompiledLocalComputation::ExecutePerReplica( results[replica] = std::move(result_buffer_status); }; - if (num_replicas == 1) { + if (num_replicas() == 1) { // Fast-path if there is only one replica — run the computation on the // current thread. execute(0); } else { // TODO(phawkins): don't recreate the threadpool for each execution. tensorflow::thread::ThreadPool pool(tensorflow::Env::Default(), "xlarun", - num_replicas - 1); + num_replicas() - 1); - for (int replica = 0; replica < num_replicas - 1; ++replica) { + for (int replica = 0; replica < num_replicas() - 1; ++replica) { pool.Schedule([&execute, replica] { execute(replica); }); } - execute(num_replicas - 1); + execute(num_replicas() - 1); } - std::vector wrapped_results(num_replicas); - for (int replica = 0; replica < num_replicas; ++replica) { + std::vector wrapped_results(num_replicas()); + for (int replica = 0; replica < num_replicas(); ++replica) { auto& statusor = results[replica]; if (!statusor.ok()) { return InternalError( @@ -549,7 +546,7 @@ StatusOr LocalComputation::Compile( argument_shape_pointers.push_back(&argument_shape); } - LocalClient* client = GetOrCreateLocalClient(); + TF_ASSIGN_OR_RETURN(LocalClient * client, GetOrCreateLocalClient()); ExecutableBuildOptions options; if (build_options != nullptr) { options = *build_options; @@ -698,8 +695,20 @@ LocalOp LocalComputationBuilder::Collapse(const LocalOp& operand, return xla::Collapse(operand.op(), dimensions); } -LocalOp LocalComputationBuilder::CrossReplicaSum(const LocalOp& operand) { - return xla::CrossReplicaSum(operand.op()); +LocalOp LocalComputationBuilder::AllToAll( + const LocalOp& operand, int64 split_dimension, int64 concat_dimension, + int64 split_count, absl::Span replica_groups) { + std::vector rg(replica_groups.size()); + for (int i = 0; i < replica_groups.size(); ++i) { + rg.push_back(replica_groups[i]); + } + return xla::AllToAll(operand.op(), split_dimension, concat_dimension, + split_count, rg); +} + +LocalOp LocalComputationBuilder::CrossReplicaSum( + const LocalOp& operand, absl::Span replica_groups) { + return xla::CrossReplicaSum(operand.op(), replica_groups); } LocalOp LocalComputationBuilder::Slice(const LocalOp& operand, @@ -921,10 +930,27 @@ LocalOp LocalComputationBuilder::QR(const LocalOp& a, bool full_matrices) { LocalOp LocalComputationBuilder::TriangularSolve(const LocalOp& a, const LocalOp& b, bool left_side, bool lower, - bool transpose_a, - bool conjugate_a) { - return xla::TriangularSolve(a.op(), b.op(), left_side, lower, transpose_a, - conjugate_a); + bool unit_diagonal, + int transpose_a) { + return xla::TriangularSolve( + a.op(), b.op(), left_side, lower, unit_diagonal, + xla::TriangularSolveOptions::Transpose(transpose_a)); +} + +LocalOp LocalComputationBuilder::Gather( + const LocalOp& input, const LocalOp& start_indices, + const GatherDimensionNumbers& dimension_numbers, + absl::Span slice_sizes) { + return xla::Gather(input.op(), start_indices.op(), dimension_numbers, + slice_sizes); +} + +LocalOp LocalComputationBuilder::Scatter( + const LocalOp& input, const LocalOp& scatter_indices, + const LocalOp& updates, const LocalComputation& update_computation, + const ScatterDimensionNumbers& dimension_numbers) { + return xla::Scatter(input.op(), scatter_indices.op(), updates.op(), + update_computation.computation(), dimension_numbers); } StatusOr LocalComputationBuilder::BuildConstantSubGraph( diff --git a/tensorflow/compiler/xla/python/local_computation_builder.h b/tensorflow/compiler/xla/python/local_computation_builder.h index 5e8341592100bc1eba4d1c17b0c2dd0e0888fdb1..bc8b7e610c02524f5b1c9bb2b0482ec3ab483e8e 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.h +++ b/tensorflow/compiler/xla/python/local_computation_builder.h @@ -180,6 +180,10 @@ class CompiledLocalComputation { public: CompiledLocalComputation(std::unique_ptr executable); + int num_replicas() const { + return executable_->build_options().num_replicas(); + } + StatusOr Execute( absl::Span argument_handles); @@ -312,7 +316,12 @@ class LocalComputationBuilder { LocalOp Collapse(const LocalOp& operand, absl::Span dimensions); - LocalOp CrossReplicaSum(const LocalOp& operand); + LocalOp AllToAll(const LocalOp& operand, int64 split_dimension, + int64 concat_dimension, int64 split_count, + absl::Span replica_groups); + + LocalOp CrossReplicaSum(const LocalOp& operand, + absl::Span replica_groups); LocalOp Slice(const LocalOp& operand, absl::Span start_indices, absl::Span limit_indices, @@ -415,8 +424,20 @@ class LocalComputationBuilder { LocalOp Cholesky(const LocalOp& a); + // `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 + // enum. LocalOp TriangularSolve(const LocalOp& a, const LocalOp& b, bool left_side, - bool lower, bool transpose_a, bool conjugate_a); + bool lower, bool unit_diagonal, int transpose_a); + + LocalOp Gather(const LocalOp& input, const LocalOp& start_indices, + const GatherDimensionNumbers& dimension_numbers, + absl::Span slice_sizes); + + LocalOp Scatter(const LocalOp& input, const LocalOp& scatter_indices, + const LocalOp& updates, + const LocalComputation& update_computation, + const ScatterDimensionNumbers& dimension_numbers); StatusOr BuildConstantSubGraph(const LocalOp& operand); diff --git a/tensorflow/compiler/xla/python/local_computation_builder.i b/tensorflow/compiler/xla/python/local_computation_builder.i index bf5d667c6a12972845735983a74264ea05675971..df2ab0b539be76ad9b974d506b06f1f24a381f1d 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.i +++ b/tensorflow/compiler/xla/python/local_computation_builder.i @@ -34,6 +34,9 @@ limitations under the License. // 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. @@ -167,8 +170,41 @@ bool HandleStringAttribute(PyObject* o, return true; // Handled string attribute, ok! } +bool HandleRepeatedInt64Attribute( + PyObject* o, const char* attr_name, + tensorflow::protobuf::RepeatedField* field) { + PyObject* seq = PyObject_GetAttrString(o, attr_name); + if (!seq) { + return false; + } + + int length = PySequence_Size(seq); + if (length == -1) { + Py_DECREF(seq); + return false; + } + + for (int i = 0; i < length; ++i) { + PyObject* item = PySequence_GetItem(seq, i); + if (!item) { + Py_DECREF(seq); + return false; + } + const int64 dimension = numpy::PyIntOrPyLongToLong(item); + if (dimension == -1 && PyErr_Occurred()) { + Py_DECREF(item); + Py_DECREF(seq); + return false; + } + *field->Add() = dimension; + Py_DECREF(item); + } + Py_DECREF(seq); + return true; } -} + +} // namespace swig +} // namespace xla %} // Required to use PyArray_* functions. @@ -417,16 +453,6 @@ tensorflow::ImportNumpy(); // Literal -%typemap(out) StatusOr { - if ($1.ok()) { - Literal value = $1.ConsumeValueOrDie(); - $result = numpy::PyObjectFromXlaLiteral(*value); - } else { - PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); - SWIG_fail; - } -} - %typemap(in) const Literal& (StatusOr literal_status) { literal_status = numpy::XlaLiteralFromPyObject($input); if (!literal_status.ok()) { @@ -436,16 +462,26 @@ tensorflow::ImportNumpy(); $1 = &literal_status.ValueOrDie(); } -%typemap(out) Literal { - $result = numpy::PyObjectFromXlaLiteral(*$1); +%typemap(out) Literal (StatusOr obj_status) { + obj_status = numpy::PyObjectFromXlaLiteral(*$1); + if (!obj_status.ok()) { + PyErr_SetString(PyExc_RuntimeError, obj_status.status().ToString().c_str()); + SWIG_fail; + } + $result = obj_status.ValueOrDie().release(); } -%typemap(out) StatusOr { +%typemap(out) StatusOr (StatusOr obj_status) { if (!$1.ok()) { PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); SWIG_fail; } - $result = numpy::PyObjectFromXlaLiteral($1.ValueOrDie()); + obj_status = numpy::PyObjectFromXlaLiteral($1.ValueOrDie()); + if (!obj_status.ok()) { + PyErr_SetString(PyExc_RuntimeError, obj_status.status().ToString().c_str()); + SWIG_fail; + } + $result = obj_status.ValueOrDie().release(); } %typemap(in) const std::vector& (std::vector temps) { @@ -657,128 +693,27 @@ tensorflow::ImportNumpy(); %typemap(in) const DotDimensionNumbers& (DotDimensionNumbers dimension_numbers) { - int length; - - /* lhs_contracting_dimensions */ - PyObject* lhs_contracting_dimensions = PyObject_GetAttrString( - $input, "lhs_contracting_dimensions"); - if (!lhs_contracting_dimensions) { - SWIG_fail; - } - - length = PySequence_Size(lhs_contracting_dimensions); - if (length == -1) { - Py_DECREF(lhs_contracting_dimensions); - SWIG_fail; - } - - for (int i = 0; i < length; ++i) { - PyObject* item = PySequence_GetItem(lhs_contracting_dimensions, i); - if (!item) { - Py_DECREF(lhs_contracting_dimensions); - SWIG_fail; - } - const int64 dimension = numpy::PyIntOrPyLongToLong(item); - if (dimension == -1 && PyErr_Occurred()) { - Py_DECREF(item); - Py_DECREF(lhs_contracting_dimensions); - SWIG_fail; - } - dimension_numbers.add_lhs_contracting_dimensions(dimension); - Py_DECREF(item); - } - Py_DECREF(lhs_contracting_dimensions); - - /* rhs_contracting_dimensions */ - PyObject* rhs_contracting_dimensions = PyObject_GetAttrString( - $input, "rhs_contracting_dimensions"); - if (!lhs_contracting_dimensions) { + if (!HandleRepeatedInt64Attribute( + $input, "lhs_contracting_dimensions", + dimension_numbers.mutable_lhs_contracting_dimensions())) { SWIG_fail; } - - length = PySequence_Size(rhs_contracting_dimensions); - if (length == -1) { - Py_DECREF(rhs_contracting_dimensions); + if (!HandleRepeatedInt64Attribute( + $input, "rhs_contracting_dimensions", + dimension_numbers.mutable_rhs_contracting_dimensions())) { SWIG_fail; } - - for (int i = 0; i < length; ++i) { - PyObject* item = PySequence_GetItem(rhs_contracting_dimensions, i); - if (!item) { - Py_DECREF(rhs_contracting_dimensions); - SWIG_fail; - } - const int64 dimension = numpy::PyIntOrPyLongToLong(item); - if (dimension == -1 && PyErr_Occurred()) { - Py_DECREF(item); - Py_DECREF(rhs_contracting_dimensions); - SWIG_fail; - } - dimension_numbers.add_rhs_contracting_dimensions(dimension); - Py_DECREF(item); - } - Py_DECREF(rhs_contracting_dimensions); - - /* lhs_batch_dimensions */ - PyObject* lhs_batch_dimensions = PyObject_GetAttrString( - $input, "lhs_batch_dimensions"); - if (!lhs_batch_dimensions) { + if (!HandleRepeatedInt64Attribute( + $input, "lhs_batch_dimensions", + dimension_numbers.mutable_lhs_batch_dimensions())) { SWIG_fail; } - - length = PySequence_Size(lhs_batch_dimensions); - if (length == -1) { - Py_DECREF(lhs_batch_dimensions); + if (!HandleRepeatedInt64Attribute( + $input, "rhs_batch_dimensions", + dimension_numbers.mutable_rhs_batch_dimensions())) { SWIG_fail; } - for (int i = 0; i < length; ++i) { - PyObject* item = PySequence_GetItem(lhs_batch_dimensions, i); - if (!item) { - Py_DECREF(lhs_batch_dimensions); - SWIG_fail; - } - const int64 dimension = numpy::PyIntOrPyLongToLong(item); - if (dimension == -1 && PyErr_Occurred()) { - Py_DECREF(item); - Py_DECREF(lhs_batch_dimensions); - SWIG_fail; - } - dimension_numbers.add_lhs_batch_dimensions(dimension); - Py_DECREF(item); - } - Py_DECREF(lhs_batch_dimensions); - - /* rhs_batch_dimensions */ - PyObject* rhs_batch_dimensions = PyObject_GetAttrString( - $input, "rhs_batch_dimensions"); - if (!rhs_batch_dimensions) { - SWIG_fail; - } - - length = PySequence_Size(rhs_batch_dimensions); - if (length == -1) { - Py_DECREF(rhs_batch_dimensions); - SWIG_fail; - } - - for (int i = 0; i < length; ++i) { - PyObject* item = PySequence_GetItem(rhs_batch_dimensions, i); - if (!item) { - Py_DECREF(rhs_batch_dimensions); - SWIG_fail; - } - const int64 dimension = numpy::PyIntOrPyLongToLong(item); - if (dimension == -1 && PyErr_Occurred()) { - Py_DECREF(item); - Py_DECREF(rhs_batch_dimensions); - SWIG_fail; - } - dimension_numbers.add_rhs_batch_dimensions(dimension); - Py_DECREF(item); - } - Py_DECREF(rhs_batch_dimensions); - $1 = &dimension_numbers; } @@ -860,90 +795,108 @@ tensorflow::ImportNumpy(); } dimension_numbers.set_kernel_input_feature_dimension(value); - PyObject* o; - int length; - - o = PyObject_GetAttrString($input, "input_spatial_dimensions"); - if (!o) { + if (!HandleRepeatedInt64Attribute( + $input, "input_spatial_dimensions", + dimension_numbers.mutable_input_spatial_dimensions())) { SWIG_fail; } - length = PySequence_Size(o); - if (length == -1) { - Py_DECREF(o); + if (!HandleRepeatedInt64Attribute( + $input, "kernel_spatial_dimensions", + dimension_numbers.mutable_kernel_spatial_dimensions())) { SWIG_fail; } - for (int i = 0; i < length; ++i) { - PyObject* item = PySequence_GetItem(o, i); - if (!item) { - Py_DECREF(o); - SWIG_fail; - } - const int64 dimension = numpy::PyIntOrPyLongToLong(item); - if (dimension == -1 && PyErr_Occurred()) { - Py_DECREF(item); - Py_DECREF(o); - SWIG_fail; - } - dimension_numbers.add_input_spatial_dimensions(dimension); - Py_DECREF(item); + if (!HandleRepeatedInt64Attribute( + $input, "output_spatial_dimensions", + dimension_numbers.mutable_output_spatial_dimensions())) { + SWIG_fail; } - Py_DECREF(o); - o = PyObject_GetAttrString($input, "kernel_spatial_dimensions"); - if (!o) { + $1 = &dimension_numbers; +} + +// GatherDimensionNumbers + +%typemap(in) const GatherDimensionNumbers& + (GatherDimensionNumbers dimension_numbers) { + if (!HandleRepeatedInt64Attribute( + $input, "offset_dims", + dimension_numbers.mutable_offset_dims())) { SWIG_fail; } - length = PySequence_Size(o); - if (length == -1) { - Py_DECREF(o); + if (!HandleRepeatedInt64Attribute( + $input, "collapsed_slice_dims", + dimension_numbers.mutable_collapsed_slice_dims())) { SWIG_fail; } - for (int i = 0; i < length; ++i) { - PyObject* item = PySequence_GetItem(o, i); - if (!item) { - Py_DECREF(o); - SWIG_fail; - } - const int64 dimension = numpy::PyIntOrPyLongToLong(item); - if (dimension == -1 && PyErr_Occurred()) { - Py_DECREF(item); - Py_DECREF(o); - SWIG_fail; - } - dimension_numbers.add_kernel_spatial_dimensions(dimension); - Py_DECREF(item); + if (!HandleRepeatedInt64Attribute( + $input, "start_index_map", + dimension_numbers.mutable_start_index_map())) { + SWIG_fail; } - Py_DECREF(o); - o = PyObject_GetAttrString($input, "output_spatial_dimensions"); - if (!o) { + int64 value; + if (!GetIntAttr($input, "index_vector_dim", &value)) { SWIG_fail; } - length = PySequence_Size(o); - if (length == -1) { - Py_DECREF(o); + dimension_numbers.set_index_vector_dim(value); + + $1 = &dimension_numbers; +} + +// ScatterDimensionNumbers + +%typemap(in) const ScatterDimensionNumbers& + (ScatterDimensionNumbers dimension_numbers) { + if (!HandleRepeatedInt64Attribute( + $input, "update_window_dims", + dimension_numbers.mutable_update_window_dims())) { SWIG_fail; } - for (int i = 0; i < length; ++i) { - PyObject* item = PySequence_GetItem(o, i); - if (!item) { - Py_DECREF(o); - SWIG_fail; - } - const int64 dimension = numpy::PyIntOrPyLongToLong(item); - if (dimension == -1 && PyErr_Occurred()) { - Py_DECREF(item); - Py_DECREF(o); - SWIG_fail; - } - dimension_numbers.add_output_spatial_dimensions(dimension); - Py_DECREF(item); + if (!HandleRepeatedInt64Attribute( + $input, "inserted_window_dims", + dimension_numbers.mutable_inserted_window_dims())) { + SWIG_fail; + } + if (!HandleRepeatedInt64Attribute( + $input, "scatter_dims_to_operand_dims", + dimension_numbers.mutable_scatter_dims_to_operand_dims())) { + SWIG_fail; } - Py_DECREF(o); + + int64 value; + if (!GetIntAttr($input, "index_vector_dim", &value)) { + SWIG_fail; + } + dimension_numbers.set_index_vector_dim(value); $1 = &dimension_numbers; } +// Span + +%typemap(in) absl::Span + (std::vector temps) { + if (!PySequence_Check($input)) { + PyErr_SetString(PyExc_TypeError, "Argument is not a sequence"); + SWIG_fail; + } + const int size = PySequence_Size($input); + temps.reserve(size); + for (int i = 0; i < size; ++i) { + PyObject* o = PySequence_GetItem($input, i); + ReplicaGroup rgrp; + if (!HandleRepeatedInt64Attribute( + o, "replica_ids", + rgrp.mutable_replica_ids())) { + SWIG_fail; + } + temps.push_back(rgrp); + Py_DECREF(o); + } + $1 = temps; +} + + // ExecutableBuildOptions %typemap(in) const ExecutableBuildOptions* @@ -1000,6 +953,12 @@ tensorflow::ImportNumpy(); } Py_DECREF(o); + int64 num_replicas; + if (!GetIntAttr($input, "num_replicas", &num_replicas)) { + SWIG_fail; + } + build_options.set_num_replicas(num_replicas); + $1 = &build_options; } } @@ -1059,6 +1018,7 @@ tensorflow::ImportNumpy(); %unignore xla::swig::LocalComputationBuilder::Pad; %unignore xla::swig::LocalComputationBuilder::Reshape; %unignore xla::swig::LocalComputationBuilder::Collapse; +%unignore xla::swig::LocalComputationBuilder::AllToAll; %unignore xla::swig::LocalComputationBuilder::CrossReplicaSum; %unignore xla::swig::LocalComputationBuilder::Slice; %unignore xla::swig::LocalComputationBuilder::SliceInDim; @@ -1151,6 +1111,8 @@ tensorflow::ImportNumpy(); %unignore xla::swig::LocalComputationBuilder::QR; %unignore xla::swig::LocalComputationBuilder::TriangularSolve; %unignore xla::swig::LocalComputationBuilder::CustomCall; +%unignore xla::swig::LocalComputationBuilder::Gather; +%unignore xla::swig::LocalComputationBuilder::Scatter; %unignore xla::swig::DeleteLocalComputation; %unignore xla::swig::DestructureLocalShapedBufferTuple; %unignore xla::swig::DestructureXrtAllocationTuple; diff --git a/tensorflow/compiler/xla/python/numpy_bridge.cc b/tensorflow/compiler/xla/python/numpy_bridge.cc index 9c9055082559fbf21df931e94866068e1e27f25e..8e056f97255174b63dbb3c7cd4a107978358feb6 100644 --- a/tensorflow/compiler/xla/python/numpy_bridge.cc +++ b/tensorflow/compiler/xla/python/numpy_bridge.cc @@ -26,6 +26,10 @@ namespace swig { namespace numpy { +Safe_PyObjectPtr make_safe(PyObject* object) { + return Safe_PyObjectPtr(object); +} + int PrimitiveTypeToNumpyType(PrimitiveType primitive_type) { switch (primitive_type) { case PRED: @@ -54,6 +58,8 @@ int PrimitiveTypeToNumpyType(PrimitiveType primitive_type) { return NPY_FLOAT64; case C64: return NPY_COMPLEX64; + case C128: + return NPY_COMPLEX128; case TUPLE: return NPY_OBJECT; default: @@ -89,6 +95,8 @@ PrimitiveType NumpyTypeToPrimitiveType(int np_type) { return F64; case NPY_COMPLEX64: return C64; + case NPY_COMPLEX128: + return C128; case NPY_OBJECT: return TUPLE; default: @@ -111,6 +119,7 @@ bool NumpyTypeIsValid(int np_type) { case NPY_FLOAT32: case NPY_FLOAT64: case NPY_COMPLEX64: + case NPY_COMPLEX128: case NPY_OBJECT: return true; default: @@ -344,13 +353,17 @@ StatusOr OpMetadataFromPyObject(PyObject* o) { return result; } -PyObject* PyObjectFromXlaLiteral(const LiteralSlice& literal) { +StatusOr PyObjectFromXlaLiteral(const LiteralSlice& literal) { if (literal.shape().IsTuple()) { int num_elements = ShapeUtil::TupleElementCount(literal.shape()); - PyObject* tuple = PyTuple_New(num_elements); + std::vector elems(num_elements); + for (int i = 0; i < num_elements; i++) { + TF_ASSIGN_OR_RETURN(elems[i], + PyObjectFromXlaLiteral(LiteralSlice(literal, {i}))); + } + Safe_PyObjectPtr tuple = make_safe(PyTuple_New(num_elements)); for (int i = 0; i < num_elements; i++) { - PyTuple_SET_ITEM(tuple, i, - PyObjectFromXlaLiteral(LiteralSlice(literal, {i}))); + PyTuple_SET_ITEM(tuple.get(), i, elems[i].release()); } return tuple; } else { @@ -360,10 +373,10 @@ PyObject* PyObjectFromXlaLiteral(const LiteralSlice& literal) { dimensions[i] = ShapeUtil::GetDimension(literal.shape(), i); } int np_type = PrimitiveTypeToNumpyType(literal.shape().element_type()); - PyObject* array = - PyArray_EMPTY(rank, dimensions.data(), np_type, /*fortran=*/0); - CopyLiteralToNumpyArray(np_type, literal, - reinterpret_cast(array)); + Safe_PyObjectPtr array = make_safe( + PyArray_EMPTY(rank, dimensions.data(), np_type, /*fortran=*/0)); + TF_RETURN_IF_ERROR(CopyLiteralToNumpyArray( + np_type, literal, reinterpret_cast(array.get()))); return array; } } @@ -403,6 +416,12 @@ Status CopyNumpyArrayToLiteral(int np_type, PyArrayObject* py_array, case NPY_BOOL: CopyNumpyArrayToLiteral(py_array, literal); break; + case NPY_INT8: + CopyNumpyArrayToLiteral(py_array, literal); + break; + case NPY_INT16: + CopyNumpyArrayToLiteral(py_array, literal); + break; case NPY_INT32: CopyNumpyArrayToLiteral(py_array, literal); break; @@ -412,6 +431,9 @@ Status CopyNumpyArrayToLiteral(int np_type, PyArrayObject* py_array, case NPY_UINT8: CopyNumpyArrayToLiteral(py_array, literal); break; + case NPY_UINT16: + CopyNumpyArrayToLiteral(py_array, literal); + break; case NPY_UINT32: CopyNumpyArrayToLiteral(py_array, literal); break; @@ -430,6 +452,9 @@ Status CopyNumpyArrayToLiteral(int np_type, PyArrayObject* py_array, case NPY_COMPLEX64: CopyNumpyArrayToLiteral(py_array, literal); break; + case NPY_COMPLEX128: + CopyNumpyArrayToLiteral(py_array, literal); + break; default: return InvalidArgument( "No XLA literal container for Numpy type number: %d", np_type); @@ -437,12 +462,18 @@ Status CopyNumpyArrayToLiteral(int np_type, PyArrayObject* py_array, return Status::OK(); } -void CopyLiteralToNumpyArray(int np_type, const LiteralSlice& literal, - PyArrayObject* py_array) { +Status CopyLiteralToNumpyArray(int np_type, const LiteralSlice& literal, + PyArrayObject* py_array) { switch (np_type) { case NPY_BOOL: CopyLiteralToNumpyArray(literal, py_array); break; + case NPY_INT8: + CopyLiteralToNumpyArray(literal, py_array); + break; + case NPY_INT16: + CopyLiteralToNumpyArray(literal, py_array); + break; case NPY_INT32: CopyLiteralToNumpyArray(literal, py_array); break; @@ -452,6 +483,9 @@ void CopyLiteralToNumpyArray(int np_type, const LiteralSlice& literal, case NPY_UINT8: CopyLiteralToNumpyArray(literal, py_array); break; + case NPY_UINT16: + CopyLiteralToNumpyArray(literal, py_array); + break; case NPY_UINT32: CopyLiteralToNumpyArray(literal, py_array); break; @@ -470,9 +504,14 @@ void CopyLiteralToNumpyArray(int np_type, const LiteralSlice& literal, case NPY_COMPLEX64: CopyLiteralToNumpyArray(literal, py_array); break; + case NPY_COMPLEX128: + CopyLiteralToNumpyArray(literal, py_array); + break; default: - LOG(FATAL) << "No XLA literal container for Numpy type" << np_type; + return InvalidArgument( + "No XLA literal container for Numpy type number: %d", np_type); } + return Status::OK(); } PyObject* LongToPyIntOrPyLong(long x) { // NOLINT diff --git a/tensorflow/compiler/xla/python/numpy_bridge.h b/tensorflow/compiler/xla/python/numpy_bridge.h index 40ff2d9ad214cc4dcad42234fa296834cbc92882..737fc4b29c162b77d5b204f503cb835d7f6c84e3 100644 --- a/tensorflow/compiler/xla/python/numpy_bridge.h +++ b/tensorflow/compiler/xla/python/numpy_bridge.h @@ -36,6 +36,16 @@ namespace swig { namespace numpy { +struct PyDecrefDeleter { + void operator()(PyObject* p) const { Py_DECREF(p); } +}; + +// Safe container for an owned PyObject. On destruction, the reference count of +// the contained object will be decremented. +using Safe_PyObjectPtr = std::unique_ptr; + +Safe_PyObjectPtr make_safe(PyObject* object); + // Maps XLA primitive types (PRED, S8, F32, ..., and TUPLE) to numpy // dtypes (NPY_BOOL, NPY_INT8, NPY_FLOAT32, ..., and NPY_OBJECT), and // vice versa. @@ -74,7 +84,7 @@ StatusOr OpMetadataFromPyObject(PyObject* o); // array data. // // The return value is a new reference. -PyObject* PyObjectFromXlaLiteral(const LiteralSlice& literal); +StatusOr PyObjectFromXlaLiteral(const LiteralSlice& literal); // Converts a Numpy ndarray or a nested Python tuple thereof to a // corresponding XLA literal. @@ -90,8 +100,8 @@ StatusOr XlaLiteralFromPyObject(PyObject* o); Status CopyNumpyArrayToLiteral(int np_type, PyArrayObject* py_array, Literal* literal); -void CopyLiteralToNumpyArray(int np_type, const LiteralSlice& literal, - PyArrayObject* py_array); +Status CopyLiteralToNumpyArray(int np_type, const LiteralSlice& literal, + PyArrayObject* py_array); template void CopyNumpyArrayToLiteral(PyArrayObject* py_array, Literal* literal) { diff --git a/tensorflow/compiler/xla/python/pywrap_xla_exported_symbols.lds b/tensorflow/compiler/xla/python/pywrap_xla_exported_symbols.lds new file mode 100644 index 0000000000000000000000000000000000000000..ef77ed3d95850fdfc7145e6fe1df4833d20bb7df --- /dev/null +++ b/tensorflow/compiler/xla/python/pywrap_xla_exported_symbols.lds @@ -0,0 +1,2 @@ +_PyInit__pywrap_xla +_init_pywrap_xla diff --git a/tensorflow/compiler/xla/python/pywrap_xla_version_script.lds b/tensorflow/compiler/xla/python/pywrap_xla_version_script.lds new file mode 100644 index 0000000000000000000000000000000000000000..d31cfce7be7b6accf05ef77f3485904099965afc --- /dev/null +++ b/tensorflow/compiler/xla/python/pywrap_xla_version_script.lds @@ -0,0 +1,6 @@ +xla { + global: + PyInit_*; + local: + *; +}; diff --git a/tensorflow/compiler/xla/python/xla_client.py b/tensorflow/compiler/xla/python/xla_client.py index 378bbdcb175f10d73da87f5286cf5129477a124c..c7afed300b678c16754541b1308a36c500282292 100644 --- a/tensorflow/compiler/xla/python/xla_client.py +++ b/tensorflow/compiler/xla/python/xla_client.py @@ -12,12 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""An in-process, local XLA client in Python, supporting AOT compilation.""" +"""An XLA client in Python, supporting AOT compilation.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function +import abc import collections import enum # pylint: disable=g-bad-import-order import inspect @@ -40,6 +41,18 @@ from tensorflow.compiler.xla.service import hlo_pb2 # pylint: disable=invalid-name +# Version of the XLA Python client. +# +# JAX packages the XLA python plugin as a binary pip module (jaxlib) that is +# packaged separately from the Python code that consumes it (jax). +# +# We occasionally need to make backwards-incompatible changes to jaxlib, in +# which case we need to be able to detect when incompatible versions are +# installed. +def version(): + return (0, 1, 7) + + _OP_METADATA_FIELDS = [ 'op_type', 'op_name', @@ -49,13 +62,123 @@ _OP_METADATA_FIELDS = [ OpMetadata = collections.namedtuple('OpMetadata', _OP_METADATA_FIELDS) +@six.add_metaclass(abc.ABCMeta) +class Backend(object): + """Abstract base class for XLA backends.""" + + @abc.abstractmethod + def buffer_from_pyval(self, pyval, device=0): + """Allocates a fresh buffer and populates it with `pyval`.""" + + @abc.abstractmethod + def delete_buffer(self, c_buffer): + """Deletes buffer `c_buffer`.""" + + @abc.abstractmethod + def destructure_tuple(self, c_buffer): + """Destructures a tuple buffer into a sequence of buffers.""" + + @abc.abstractmethod + def compile(self, computation, argument_shapes, compile_options): + """Compiles a computation. Returns an executable.""" + + @abc.abstractmethod + def delete_executable(self, executable): + """Deletes an executable.""" + + @abc.abstractmethod + def execute(self, executable, args): + """Runs an executable without replication.""" + + @abc.abstractmethod + def execute_replicated(self, executable, per_replica_args): + """Runs an executable in a replicated manner.""" + + +class XlaLocalBackend(Backend): + """XLA backend implemented using the in-process xla::LocalClient API.""" + + def buffer_from_pyval(self, pyval, device=0): + return c_api.LocalShapedBuffer.FromLiteral(pyval, None, device) + + def delete_buffer(self, c_buffer): + c_api.DeleteLocalShapedBuffer(c_buffer) + + def destructure_tuple(self, c_buffer): + result = c_api.DestructureLocalShapedBufferTuple(c_buffer) + return [result.Release(i) for i in xrange(result.size())] + + def compile(self, c_computation, argument_shapes, compile_options): + return c_computation.Compile(argument_shapes, compile_options) + + def delete_executable(self, executable): + assert isinstance(executable, c_api.CompiledLocalComputation) + c_api.DeleteCompiledLocalComputation(executable) + + def execute(self, executable, args): + return executable.Execute(args) + + def execute_replicated(self, executable, per_replica_args): + output_buffer_tup = executable.ExecutePerReplica(per_replica_args) + size = output_buffer_tup.size() + return [output_buffer_tup.Release(i) for i in xrange(size)] + + +class XrtBackend(Backend): + """XLA backend implemented using XRT.""" + + def __init__(self, target): + self.target = target + + def buffer_from_pyval(self, pyval, device=0): + if device != 0: + raise NotImplementedError( + 'Multi-replica execution is not yet supported via the XRT backend.') + return c_api.XrtAllocation.FromLiteral(pyval, + _maybe_encode_string(self.target)) + + def delete_buffer(self, c_buffer): + c_api.DeleteXrtAllocation(c_buffer) + + def destructure_tuple(self, c_buffer): + result = c_api.DestructureXrtAllocationTuple( + c_buffer, _maybe_encode_string(self.target)) + return [result.Release(i) for i in xrange(result.size())] + + def compile(self, c_computation, argument_shapes, compile_options): + return c_computation.CompileForXrt(argument_shapes, + _maybe_encode_string(self.target)) + + def delete_executable(self, executable): + assert isinstance(executable, c_api.CompiledXrtComputation) + c_api.DeleteCompiledXrtComputation(executable) + + def execute(self, executable, args): + return executable.Execute(args) + + def execute_replicated(self, executable, per_replica_args): + if len(per_replica_args) != 1: + raise NotImplementedError( + 'Multi-replica execution is not yet supported via the XRT backend.') + return [executable.Execute(per_replica_args[0])] + + +XLA_LOCAL_BACKEND = XlaLocalBackend() + + class BackendType(enum.Enum): XLA_LOCAL = 1 XRT = 2 -BackendSpec = collections.namedtuple('Backend', ('backend_type', 'target')) -XLA_LOCAL_BACKEND = BackendSpec(BackendType.XLA_LOCAL, 'local') +def BackendSpec(backend, target): + """Compatibility wrapper to support older clients. Do not use in new code.""" + if backend == BackendType.XLA_LOCAL: + return XLA_LOCAL_BACKEND + elif backend == BackendType.XRT: + return XrtBackend(target) + else: + raise ValueError('Unknown backend {}'.format(backend)) def OpMetadataToProto(pyobj): @@ -199,6 +322,7 @@ XLA_ELEMENT_TYPE_TO_DTYPE = { 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), } @@ -226,10 +350,6 @@ class LocalBuffer(object): self.c_buffer = c_buffer self._backend = backend self._replica = replica - if backend.backend_type == BackendType.XRT: - self._delete = c_api.DeleteXrtAllocation - else: - self._delete = c_api.DeleteLocalShapedBuffer @staticmethod def from_pyval(pyval, replica=0, backend=XLA_LOCAL_BACKEND): @@ -240,14 +360,7 @@ class LocalBuffer(object): raise ValueError( 'Attempt to place buffer on replica {} when the replica count is {}' .format(replica, num_replicas)) - if backend.backend_type == BackendType.XRT: - if replica != 0: - raise NotImplementedError( - 'Multi-replica execution is not yet supported via the XRT backend.') - cbuf = c_api.XrtAllocation.FromLiteral( - pyval, _maybe_encode_string(backend.target)) - else: - cbuf = c_api.LocalShapedBuffer.FromLiteral(pyval, None, replica) + cbuf = backend.buffer_from_pyval(pyval, replica) return LocalBuffer(cbuf, backend, replica) def to_py(self): @@ -261,24 +374,17 @@ class LocalBuffer(object): def delete(self): if self.c_buffer is not None: - self._delete(self.c_buffer) + self._backend.delete_buffer(self.c_buffer) self.c_buffer = None def destructure(self): """Assuming a tuple buffer, unpack it into constituent tuple elements.""" assert self.c_buffer is not None - if self._backend.backend_type == BackendType.XRT: - result = c_api.DestructureXrtAllocationTuple( - self.c_buffer, _maybe_encode_string(self._backend.target)) - else: - result = c_api.DestructureLocalShapedBufferTuple(self.c_buffer) + result = self._backend.destructure_tuple(self.c_buffer) self.delete() - size = result.size() - destructured = tuple( - LocalBuffer( - result.Release(i), replica=self._replica, backend=self._backend) - for i in xrange(size)) - return destructured + return tuple( + LocalBuffer(sub_buffer, replica=self._replica, backend=self._backend) + for sub_buffer in result) def is_deleted(self): return self.c_buffer is None @@ -427,6 +533,20 @@ class Shape(object): updated._check_minor_to_major() # pylint: disable=protected-access return updated + def serialize(self, proto): + """Serializes 'shape' into proto.""" + if self.is_tuple(): + proto.element_type = xla_data_pb2.TUPLE + for shape in self.tuple_shapes(): + shape.serialize(proto.tuple_shapes.add()) + else: + proto.element_type = dtype_to_etype(self.element_type()) + 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.minor_to_major.extend(self.minor_to_major()) + def _wrap_shape(shape_info): dtype, dims = shape_info @@ -458,6 +578,7 @@ class CompileOptions(object): self.dump_unoptimized_hlo_proto_to = None self.dump_per_pass_hlo_proto_to = None self.hlo_profile = False + self.num_replicas = get_replica_count() def transfer_to_infeed(value, replica_number=None): @@ -507,18 +628,6 @@ class LocalComputation(object): self._backend = backend self._is_compiled = is_compiled - # Ensure a reference to C-based destructor for use in __del__. - if is_compiled: - if backend.backend_type == BackendType.XRT: - assert isinstance(c_computation, c_api.CompiledXrtComputation) - self._delete = c_api.DeleteCompiledXrtComputation - else: - assert isinstance(c_computation, c_api.CompiledLocalComputation) - self._delete = c_api.DeleteCompiledLocalComputation - else: - assert isinstance(c_computation, c_api.LocalComputation) - self._delete = c_api.DeleteLocalComputation - @property def computation(self): if self._is_compiled: @@ -572,11 +681,8 @@ class LocalComputation(object): compile_options = compile_options or CompileOptions() compile_options.result_shape = result_shape - if self._backend.backend_type == BackendType.XRT: - c = self.computation.CompileForXrt( - argument_shapes, _maybe_encode_string(self._backend.target)) - else: - c = self.computation.Compile(argument_shapes, compile_options) + c = self._backend.compile(self.computation, argument_shapes, + compile_options) return LocalComputation(c, is_compiled=True, backend=self._backend) def CompileWithExampleArguments(self, @@ -596,7 +702,7 @@ class LocalComputation(object): if check_for_deleted_args and any(arg.is_deleted() for arg in arguments): raise ValueError('Executing with deleted local buffer argument') raw_args = [arg.c_buffer for arg in arguments] - output_buffer = self._c_computation.Execute(raw_args) + output_buffer = self._backend.execute(self._c_computation, raw_args) return LocalBuffer(output_buffer, backend=self._backend, replica=0) def ExecutePerReplica(self, arguments=None): @@ -634,15 +740,8 @@ class LocalComputation(object): ] # Execute - if self._backend.backend_type == BackendType.XRT: - if len(stripped_args) > 1: - raise NotImplementedError( - 'Multi-replica execution is not yet supported via the XRT backend.') - output_buffers = [self._c_computation.Execute(stripped_args[0])] - else: - output_buffer_tup = self._c_computation.ExecutePerReplica(stripped_args) - size = output_buffer_tup.size() - output_buffers = [output_buffer_tup.Release(i) for i in xrange(size)] + output_buffers = self._backend.execute_replicated( + self._c_computation, stripped_args) # Wrap output handles in LocalBuffer instances return tuple( @@ -670,7 +769,19 @@ class LocalComputation(object): return [out.to_py() for out in self.ExecutePerReplica(arguments)] def __del__(self): - self._delete(self._c_computation) + # Python may have freed c_api first. + if c_api and self._c_computation: + if self._is_compiled: + self._backend.delete_executable(self._c_computation) + else: + assert isinstance(self._c_computation, c_api.LocalComputation) + c_api.DeleteLocalComputation(self._c_computation) + + +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): @@ -963,16 +1074,60 @@ class ComputationBuilder(object): dimensions = tuple(range(ndim)) return self._client.Reshape(operand, dimensions, new_sizes) - def CrossReplicaSum(self, operand): + def AllToAll(self, + operand, + split_dimension, + concat_dimension, + replica_groups=None): + """AllToAll op. + + Args: + operand: LocalOp representing the input array + split_dimension: the dimension along which the operand is split + concat_dimension: the dimension along which the split blocks are + concatenated + replica_groups: optional, list of lists of ints encoding a partition of + the set {0, 1, ..., num_replicas} into equally-sized replica groups + within which the all-to-all is performed. If not supplied or None (the + default), all replicas belong to the same group. + + Returns: + A LocalOp that represents the all-to-all concatenation. + """ + if replica_groups is None: + replica_groups_protos = [] # special value for XLA API + else: + replica_groups = list(replica_groups) + replica_groups_protos = [ + _make_replica_group_proto(group) for group in replica_groups] + if not replica_groups: + split_count = get_replica_count() + else: + split_count = len(replica_groups[0]) + if not all(split_count == len(g) for g in replica_groups): + raise ValueError('Replica groups must be equally sized') + return self._client.AllToAll(operand, split_dimension, concat_dimension, + split_count, replica_groups_protos) + + def CrossReplicaSum(self, operand, replica_groups=None): """CrossReplicaSum op. Args: operand: the operand to sum across replica instances. + replica_groups: optional, list of lists of ints encoding a partition of + the set {0, 1, ..., num_replicas} into equally-sized replica groups + within which the cross-replica sum is performed. If not supplied or None + (the default), all replicas belong to the same group. Returns: - A LocalOp that has the sum of the value among all replicas. + A LocalOp that represents on each replica the sum of its group's values. """ - return self._client.CrossReplicaSum(operand) + if replica_groups is None: + replica_groups = [] # special value for XLA API + else: + replica_groups = [ + _make_replica_group_proto(group) for group in replica_groups] + return self._client.CrossReplicaSum(operand, replica_groups) def Collapse(self, operand, dimensions): """Collapse op.""" @@ -1471,11 +1626,35 @@ class ComputationBuilder(object): """Enqueues a QR decomposition onto the computation.""" return self._client.QR(a, full_matrices) - def TriangularSolve(self, a, b, left_side=False, lower=False, - transpose_a=False, conjugate_a=False): + def TriangularSolve(self, + a, + b, + left_side=False, + lower=False, + transpose_a=False, + conjugate_a=False, + unit_diagonal=False): """Enqueues a triangular-solve operation onto the computation.""" - return self._client.TriangularSolve( - a, b, left_side, lower, transpose_a, conjugate_a) + if not transpose_a: + transpose = 1 + if conjugate_a: + a = self.Conj(a) + else: + transpose = 3 if conjugate_a else 2 + return self._client.TriangularSolve(a, b, left_side, lower, unit_diagonal, + transpose) + + 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) + + 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,) def _forward_methods_to_local_builder(): @@ -1538,6 +1717,8 @@ def initialize_platform_name(platform_name): Raises: A runtime exception if the XLA service has already been initialized. + A runtime exception if the platform does not exist, or there are no devices + with that platform. """ platform_name = _maybe_encode_string(platform_name) c_api.InitializePlatformName(platform_name) diff --git a/tensorflow/compiler/xla/python/xla_client_test.py b/tensorflow/compiler/xla/python/xla_client_test.py index 002a20e60a9fbe117af991731a555e60eef9397a..aa38c06cf908079e627156f51264965892de7ff0 100644 --- a/tensorflow/compiler/xla/python/xla_client_test.py +++ b/tensorflow/compiler/xla/python/xla_client_test.py @@ -88,6 +88,12 @@ def NumpyArrayBool(*args, **kwargs): class ComputationsWithConstantsTest(LocalComputationTest): """Tests focusing on Constant ops.""" + def testConstantScalarSumS8(self): + c = self._NewComputation() + root = c.Add(c.Constant(np.int8(1)), c.Constant(np.int8(2))) + self.assertEqual(c.GetShape(root), c.GetReturnValueShape()) + self._ExecuteAndCompareExact(c, expected=np.int8(3)) + def testConstantScalarSumF32(self): c = self._NewComputation() root = c.Add(c.ConstantF32Scalar(1.11), c.ConstantF32Scalar(3.14)) @@ -553,6 +559,18 @@ class SingleOpTest(LocalComputationTest): 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 + def DISABLED_testAllToAllOneReplica(self): + samples = [ + NumpyArrayF32([97.0]), + NumpyArrayF32([64.0, 117.0]), + NumpyArrayF32([[2.0, 3.0], [4.0, 5.0]]), + ] + for lhs in samples[:1]: + c = self._NewComputation() + c.AllToAll(c.Constant(lhs), 0, 0) + self._ExecuteAndCompareExact(c, expected=lhs) + def testCrossReplicaSumOneReplica(self): samples = [ NumpyArrayF32(42.0), @@ -565,6 +583,18 @@ class SingleOpTest(LocalComputationTest): c.CrossReplicaSum(c.Constant(lhs)) self._ExecuteAndCompareExact(c, expected=lhs) + def testCrossReplicaSumOneReplicaWithSingletonGroup(self): + samples = [ + NumpyArrayF32(42.0), + NumpyArrayF32([97.0]), + NumpyArrayF32([64.0, 117.0]), + NumpyArrayF32([[2.0, 3.0], [4.0, 5.0]]), + ] + for lhs in samples: + c = self._NewComputation() + c.CrossReplicaSum(c.Constant(lhs), [[0]]) + self._ExecuteAndCompareExact(c, expected=lhs) + def testDotMatrixVectorF32(self): c = self._NewComputation() lhs = NumpyArrayF32([[2.0, 3.0], [4.0, 5.0]]) @@ -1129,6 +1159,21 @@ class SingleOpTest(LocalComputationTest): self.assertFalse(c.IsConstant(non_const_expr)) # self.assertTrue(c.IsConstant(c.Sub(c.Add(x, a), x))) # TODO(b/77245564) + 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.offset_dims.append(1) + dnums.offset_dims.append(2) + dnums.start_index_map.append(0) + dnums.start_index_map.append(1) + dnums.index_vector_dim = 2 + c = self._NewComputation() + c.Gather(c.Constant(a), c.Constant(indices), dnums, slice_sizes=[1, 1]) + g = self._Execute(c, ()) + expected = np.array([[[[2, 7]]], [[[5, 6]]]], dtype=np.int32) + np.testing.assert_allclose(g, expected, rtol=1e-4) + class EmbeddedComputationsTest(LocalComputationTest): """Tests for XLA graphs with embedded computations (such as maps).""" @@ -1186,6 +1231,14 @@ class EmbeddedComputationsTest(LocalComputationTest): c.Mul(c.ParameterFromNumpy(NumpyArrayF64(0)), c.ConstantF64Scalar(2.0)) return c.Build() + def _CreateBinaryAddS32Computation(self): + """Computation (s32, s32) -> s32 that adds its two parameters.""" + c = self._NewComputation("add_param0_by_param1") + c.Add( + c.ParameterFromNumpy(NumpyArrayS32(0)), + c.ParameterFromNumpy(NumpyArrayS32(0))) + return c.Build() + def _CreateBinaryAddF32Computation(self): """Computation (f32, f32) -> f32 that adds its two parameters.""" c = self._NewComputation("add_param0_by_param1") @@ -1568,6 +1621,23 @@ class EmbeddedComputationsTest(LocalComputationTest): execution.join() self.assertEqual(want, got) + def testScatter(self): + a = np.arange(9).astype(np.int32).reshape((3, 3)) + 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.update_window_dims.append(1) + dnums.inserted_window_dims.append(0) + dnums.scatter_dims_to_operand_dims.append(0) + dnums.index_vector_dim = 1 + + c = self._NewComputation() + c.Scatter(c.Constant(a), c.Constant(scatter_indices), c.Constant(updates), + self._CreateBinaryAddS32Computation(), dnums) + expected = np.array([[10, 21, 32], [3, 4, 5], [76, 87, 98]], dtype=np.int32) + self._ExecuteAndCompareClose(c, expected=expected) + class ErrorTest(LocalComputationTest): diff --git a/tensorflow/compiler/xla/python_api/xla_literal.py b/tensorflow/compiler/xla/python_api/xla_literal.py index 757e41a78ad2b57d2ef6e1f3055160be22c7b3ed..19bd685ab2260485d2a86f0a682d0cdd36712fdb 100644 --- a/tensorflow/compiler/xla/python_api/xla_literal.py +++ b/tensorflow/compiler/xla/python_api/xla_literal.py @@ -69,7 +69,7 @@ def _ConvertNumpyArrayToLiteral(ndarray): if ndarray.ndim == 0: getattr(literal, type_record.literal_field_name).append( - _np.asscalar(ndarray.astype(type_record.literal_field_type))) + ndarray.astype(type_record.literal_field_type).item()) else: # Ndarrays with boolean dtypes need special type conversion with protobufs if ndarray.dtype in {_np.bool_, _np.dtype('bool')}: diff --git a/tensorflow/compiler/xla/reference_util.cc b/tensorflow/compiler/xla/reference_util.cc index 033f6c9b0f2dadbdeeb49cae27e11fc592e183ee..08b78ee244844f41d551d7e249cec0cbf157d639 100644 --- a/tensorflow/compiler/xla/reference_util.cc +++ b/tensorflow/compiler/xla/reference_util.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_set.h" #include "absl/memory/memory.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal_util.h" @@ -605,24 +606,26 @@ ReferenceUtil::ReduceToRowArray2D( const std::function& reduce_function) { std::vector result; CHECK_EQ(dims.size(), 3); - const std::set dim_set(dims.begin(), dims.end()); + const absl::flat_hash_set dim_set(dims.begin(), dims.end()); CHECK_EQ(dim_set.size(), 3); - for (int64 a0 = 0; a0 == 0 || (!dim_set.count(0) && a0 < array.n1()); ++a0) { - for (int64 a1 = 0; a1 == 0 || (!dim_set.count(1) && a1 < array.n2()); + for (int64 a0 = 0; a0 == 0 || (!dim_set.contains(0) && a0 < array.n1()); + ++a0) { + for (int64 a1 = 0; a1 == 0 || (!dim_set.contains(1) && a1 < array.n2()); ++a1) { - for (int64 a2 = 0; a2 == 0 || (!dim_set.count(2) && a2 < array.n3()); + for (int64 a2 = 0; a2 == 0 || (!dim_set.contains(2) && a2 < array.n3()); ++a2) { - for (int64 a3 = 0; a3 == 0 || (!dim_set.count(3) && a3 < array.n4()); + for (int64 a3 = 0; a3 == 0 || (!dim_set.contains(3) && a3 < array.n4()); ++a3) { float accumulator = init; - for (int64 i0 = 0; i0 == 0 || (dim_set.count(0) && i0 < array.n1()); - ++i0) { - for (int64 i1 = 0; i1 == 0 || (dim_set.count(1) && i1 < array.n2()); - ++i1) { + for (int64 i0 = 0; + i0 == 0 || (dim_set.contains(0) && i0 < array.n1()); ++i0) { + for (int64 i1 = 0; + i1 == 0 || (dim_set.contains(1) && i1 < array.n2()); ++i1) { for (int64 i2 = 0; - i2 == 0 || (dim_set.count(2) && i2 < array.n3()); ++i2) { + i2 == 0 || (dim_set.contains(2) && i2 < array.n3()); ++i2) { for (int64 i3 = 0; - i3 == 0 || (dim_set.count(3) && i3 < array.n4()); ++i3) { + i3 == 0 || (dim_set.contains(3) && i3 < array.n4()); + ++i3) { // Handle zero-sized arrays. if (array.n1() > 0 && array.n2() > 0 && array.n3() > 0 && array.n4() > 0) { diff --git a/tensorflow/compiler/xla/rpc/grpc_service.cc b/tensorflow/compiler/xla/rpc/grpc_service.cc index d8123a6de28ca532819ece4a75cd0b725f8c1bbd..22b4218fbd5e9bc59a0de22735eb51db46670f09 100644 --- a/tensorflow/compiler/xla/rpc/grpc_service.cc +++ b/tensorflow/compiler/xla/rpc/grpc_service.cc @@ -47,6 +47,14 @@ namespace xla { }); } +::grpc::Status GRPCService::GetDeviceHandles(::grpc::ServerContext* context, + const GetDeviceHandlesRequest* arg, + GetDeviceHandlesResponse* result) { + return DelegateRPC([this, arg, result]() { + return service_->GetDeviceHandles(arg, result); + }); +} + ::grpc::Status GRPCService::Compile(::grpc::ServerContext* /*context*/, const CompileRequest* arg, CompileResponse* result) { @@ -61,6 +69,14 @@ namespace xla { [this, arg, result]() { return service_->Execute(arg, result); }); } +::grpc::Status GRPCService::ExecuteGraphParallel( + ::grpc::ServerContext* /*context*/, const ExecuteGraphParallelRequest* arg, + ExecuteParallelResponse* result) { + return DelegateRPC([this, arg, result]() { + return service_->ExecuteGraphParallel(arg, result); + }); +} + ::grpc::Status GRPCService::WaitForExecution(::grpc::ServerContext* context, const WaitForExecutionRequest* arg, WaitForExecutionResponse* result) { diff --git a/tensorflow/compiler/xla/rpc/grpc_service.h b/tensorflow/compiler/xla/rpc/grpc_service.h index 3e586b288a56a22573d0c3b9ae7b2f25fdbf851a..b546704f73e34941cbf7bc2fe08062aa438039f7 100644 --- a/tensorflow/compiler/xla/rpc/grpc_service.h +++ b/tensorflow/compiler/xla/rpc/grpc_service.h @@ -39,6 +39,10 @@ class GRPCService : public grpc::XlaService::Service { const DeconstructTupleRequest* arg, DeconstructTupleResponse* result) override; + ::grpc::Status GetDeviceHandles(::grpc::ServerContext* context, + const GetDeviceHandlesRequest* arg, + GetDeviceHandlesResponse* result) override; + ::grpc::Status Compile(::grpc::ServerContext* context, const CompileRequest* arg, CompileResponse* result) override; @@ -46,6 +50,9 @@ class GRPCService : public grpc::XlaService::Service { ::grpc::Status Execute(::grpc::ServerContext* context, const ExecuteRequest* arg, ExecuteResponse* result) override; + ::grpc::Status ExecuteGraphParallel(::grpc::ServerContext* context, + const ExecuteGraphParallelRequest* arg, + ExecuteParallelResponse* result) override; ::grpc::Status WaitForExecution(::grpc::ServerContext* context, const WaitForExecutionRequest* arg, diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD index eed958dc6b52882e13de1f09663bd7f52ce38c0d..a5eae6d3962255d25a72362df45f0d8af52b1011 100644 --- a/tensorflow/compiler/xla/service/BUILD +++ b/tensorflow/compiler/xla/service/BUILD @@ -114,6 +114,7 @@ tf_cc_test( ":bfloat16_normalization", ":bfloat16_support", ":hlo", + ":hlo_creation_utils", ":hlo_verifier", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -223,19 +224,23 @@ cc_library( "hlo_evaluator_typed_visitor.h", "hlo_evaluator_typed_visitor_bfloat16.cc", "hlo_evaluator_typed_visitor_bool.cc", + "hlo_evaluator_typed_visitor_complex128.cc", "hlo_evaluator_typed_visitor_complex64.cc", "hlo_evaluator_typed_visitor_double.cc", "hlo_evaluator_typed_visitor_float.cc", "hlo_evaluator_typed_visitor_half.cc", + "hlo_evaluator_typed_visitor_int16.cc", "hlo_evaluator_typed_visitor_int32.cc", "hlo_evaluator_typed_visitor_int64.cc", "hlo_evaluator_typed_visitor_int8.cc", + "hlo_evaluator_typed_visitor_uint16.cc", "hlo_evaluator_typed_visitor_uint32.cc", "hlo_evaluator_typed_visitor_uint64.cc", "hlo_evaluator_typed_visitor_uint8.cc", ], hdrs = ["hlo_evaluator.h"], deps = [ + ":dynamic_dimension_inference", ":hlo", ":hlo_casting_utils", ":hlo_query", @@ -256,6 +261,7 @@ cc_library( "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/container:node_hash_map", "@com_google_absl//absl/memory", + "@com_google_absl//absl/meta:type_traits", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", @@ -267,6 +273,7 @@ tf_cc_test( srcs = ["hlo_evaluator_test.cc"], deps = [ ":hlo", + ":hlo_element_type_converter", ":hlo_evaluator", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:reference_util", @@ -279,7 +286,6 @@ tf_cc_test( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:xla_builder", - "//tensorflow/compiler/xla/service:hlo_element_type_converter", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:test_utils", @@ -515,6 +521,7 @@ tf_cc_test( "//tensorflow/compiler/xla:window_util", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "@com_google_absl//absl/container:flat_hash_map", ], ) @@ -677,6 +684,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", "//third_party/eigen3", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", @@ -695,6 +703,7 @@ cc_library( ":compiler", ":computation_layout", ":device_memory_allocator", + ":dynamic_dimension_inference", ":executable", ":execution_tracker", ":hlo", @@ -1002,6 +1011,7 @@ cc_library( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", @@ -1053,7 +1063,6 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", - "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", @@ -1091,7 +1100,6 @@ cc_library( ":buffer_value_containers", ":heap_simulator", ":hlo", - ":hlo_memory_scheduler", ":hlo_proto", ":logical_buffer", ":tuple_points_to_analysis", @@ -1136,6 +1144,7 @@ tf_cc_test( "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/memory", ], ) @@ -1195,7 +1204,6 @@ cc_library( ":tuple_points_to_analysis", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:util", - "//tensorflow/core:lib", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/memory", @@ -1230,7 +1238,6 @@ cc_library( deps = [ ":hlo", ":hlo_proto", - "//tensorflow/compiler/xla:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", ], @@ -1454,11 +1461,15 @@ cc_library( hdrs = ["hlo_creation_utils.h"], deps = [ ":hlo", + ":hlo_module_config", ":shape_inference", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:literal_util", "//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:comparators", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", @@ -1498,12 +1509,25 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", - "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", ], ) +cc_library( + name = "op_expander_pass", + srcs = ["op_expander_pass.cc"], + hdrs = ["op_expander_pass.h"], + deps = [ + ":hlo", + ":hlo_creation_utils", + ":hlo_pass", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:util", + "@com_google_absl//absl/algorithm:container", + ], +) + cc_library( name = "gather_expander", srcs = ["gather_expander.cc"], @@ -1512,6 +1536,7 @@ cc_library( ":hlo", ":hlo_creation_utils", ":hlo_pass", + ":op_expander_pass", ":while_util", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:statusor", @@ -1535,6 +1560,28 @@ cc_library( ], ) +cc_library( + name = "triangular_solve_expander", + srcs = ["triangular_solve_expander.cc"], + hdrs = ["triangular_solve_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: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", @@ -1580,6 +1627,7 @@ cc_library( "//tensorflow/core:lib", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", @@ -1595,7 +1643,7 @@ tf_cc_test( ":algebraic_simplifier", ":hlo", ":hlo_casting_utils", - ":hlo_matchers", + ":hlo_creation_utils", ":hlo_parser", ":hlo_pass", ":pattern_matcher", @@ -1720,9 +1768,9 @@ cc_library( ) tf_cc_test( - name = "convolution_feature_group_converter_test", + name = "convolution_group_converter_test", size = "small", - srcs = ["convolution_feature_group_converter_test.cc"], + srcs = ["convolution_group_converter_test.cc"], deps = [ ":convolution_group_converter", ":hlo", @@ -1866,8 +1914,9 @@ cc_library( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/strings", ], ) @@ -1922,6 +1971,7 @@ cc_library( hdrs = ["dynamic_dimension_inference.h"], deps = [ ":hlo", + ":while_util", "//tensorflow/compiler/xla:status", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", @@ -1931,6 +1981,46 @@ cc_library( ], ) +cc_library( + name = "dynamic_padder", + srcs = ["dynamic_padder.cc"], + hdrs = ["dynamic_padder.h"], + deps = [ + ":dynamic_dimension_inference", + ":hlo_dce", + ":hlo_pass", + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", + ], +) + +tf_cc_test( + name = "dynamic_padder_test", + srcs = ["dynamic_padder_test.cc"], + deps = [ + ":dynamic_padder", + "//tensorflow/compiler/xla:debug_options_flags", + "//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:xla_builder", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_matchers", + "//tensorflow/compiler/xla/service:hlo_runner", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/core:test", + ], +) + tf_cc_test( name = "dynamic_dimension_inference_test", srcs = ["dynamic_dimension_inference_test.cc"], @@ -2017,7 +2107,6 @@ cc_library( "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", - "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", @@ -2058,6 +2147,7 @@ tf_cc_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client", "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:local_client", @@ -2116,6 +2206,8 @@ tf_cc_test( "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", ], ) @@ -2255,6 +2347,7 @@ tf_cc_test( srcs = ["hlo_dataflow_analysis_test.cc"], deps = [ ":hlo", + ":hlo_creation_utils", ":hlo_dataflow_analysis", ":hlo_graph_dumper", ":hlo_matchers", @@ -2288,6 +2381,7 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", ], @@ -2424,6 +2518,7 @@ tf_cc_test( srcs = ["tuple_points_to_analysis_test.cc"], deps = [ ":hlo", + ":hlo_creation_utils", ":hlo_matchers", ":instruction_fusion", ":tuple_points_to_analysis", @@ -2548,6 +2643,7 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_set", ], ) @@ -2798,7 +2894,6 @@ cc_library( "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/container:inlined_vector", @@ -2971,15 +3066,11 @@ cc_library( srcs = ["hlo_get_dimension_size_rewriter.cc"], hdrs = ["hlo_get_dimension_size_rewriter.h"], deps = [ + ":dynamic_dimension_inference", ":hlo", ":hlo_pass", ":shape_inference", - "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:literal_util", - "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/core:lib", "@com_google_absl//absl/algorithm:container", ], ) @@ -3164,6 +3255,7 @@ tf_cc_test( "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", ], ) @@ -3188,6 +3280,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:regexp_internal", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/types:optional", @@ -3221,7 +3314,6 @@ cc_library( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", ], ) @@ -3348,7 +3440,6 @@ cc_library( ":hlo_pass_pipeline", "//tensorflow/compiler/xla:shape_util", "//tensorflow/core:lib", - "@com_google_absl//absl/container:flat_hash_map", ], ) @@ -3405,10 +3496,39 @@ cc_library( ":hlo_profile_printer_data", ":human_readable_profile_builder", "//tensorflow/compiler/xla:types", + "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/strings", ], ) +cc_library( + name = "sort_simplifier", + srcs = ["sort_simplifier.cc"], + hdrs = ["sort_simplifier.h"], + deps = [ + ":hlo", + ":hlo_pass", + "//tensorflow/compiler/xla:statusor", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", + ], +) + +tf_cc_test( + name = "sort_simplifier_test", + srcs = ["sort_simplifier_test.cc"], + deps = [ + ":hlo_matchers", + ":hlo_parser", + ":pattern_matcher", + ":pattern_matcher_gmock", + ":sort_simplifier", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/core:test", + ], +) + cc_library( name = "tuple_util", srcs = ["tuple_util.cc"], @@ -3505,9 +3625,7 @@ cc_library( ":while_util", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:util", - "//tensorflow/core:lib", "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:inlined_vector", ], ) @@ -3562,7 +3680,6 @@ cc_library( ":hlo_evaluator", ":hlo_pass", "//tensorflow/compiler/xla:util", - "//tensorflow/core:lib", "//tensorflow/core:ptr_util", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", @@ -3576,14 +3693,16 @@ cc_library( tf_cc_test( name = "indexed_array_analysis_test", srcs = ["indexed_array_analysis_test.cc"], + extra_copts = ["-Wno-string-plus-int"], deps = [ ":hlo_matchers", + ":hlo_parser", ":indexed_array_analysis", "//tensorflow/compiler/xla:test", - "//tensorflow/compiler/xla/service:hlo_parser", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/core:test", + "@com_google_absl//absl/strings", ], ) @@ -3605,6 +3724,7 @@ cc_library( "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/types:variant", ], ) @@ -3668,6 +3788,47 @@ cc_library( ], ) +cc_library( + name = "optimize_input_output_buffer_alias", + srcs = ["optimize_input_output_buffer_alias.cc"], + hdrs = ["optimize_input_output_buffer_alias.h"], + deps = [ + ":hlo", + ":hlo_pass", + "//tensorflow/compiler/xla:shape_tree", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/memory", + "@com_google_absl//absl/strings:str_format", + ], +) + +tf_cc_test( + name = "optimize_input_output_buffer_alias_test", + srcs = ["optimize_input_output_buffer_alias_test.cc"], + deps = [ + ":optimize_input_output_buffer_alias", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:test_utils", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:lib", + "//tensorflow/core:test", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/memory", + "@com_google_absl//absl/strings:str_format", + ], +) + cc_library( name = "ar_crs_combiner", srcs = ["ar_crs_combiner.cc"], @@ -3677,10 +3838,10 @@ cc_library( ":pattern_matcher", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_pass", "@com_google_absl//absl/container:flat_hash_map", diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/tensorflow/compiler/xla/service/algebraic_simplifier.cc index c573b6016141f6b8e6bd0413d637dee68eee104b..c5deb74e96ad5d9ff62d1407f84064dad63e61fb 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier.cc @@ -27,6 +27,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" #include "absl/container/inlined_vector.h" #include "absl/memory/memory.h" #include "absl/strings/str_cat.h" @@ -35,6 +36,7 @@ limitations under the License. #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" #include "tensorflow/compiler/xla/service/hlo_casting_utils.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" @@ -51,6 +53,7 @@ limitations under the License. #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/window_util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/bits.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/logging.h" @@ -121,23 +124,37 @@ bool TransposeIsBitcast(const HloInstruction* transpose) { transpose->dimensions()); } -// Returns true if the given reshape/copy produces a result which is bit-wise -// identical to its operand and thus may be replaced with a bitcast. -// -// This function is conservative -- even if this function returns false, the -// reshape may still be a bitcast. For example, a reshape from [28x28] to [784]. -bool ReshapeOrCopyIsBitcast( - const HloInstruction* instr, - const AlgebraicSimplifierOptions::ValidBitcastCallback& - valid_bitcast_callback) { +// Recursive helper for method below. +HloInstruction* BitcastingOperandOfReshapeOrCopyChainHelper( + HloInstruction* instr, HloInstruction* operand, + const AlgebraicSimplifierOptions& options) { + // Can't replace chain of copies and reshapes with bitcasts if the compiler + // used a memory layout which isn't compatible. + if (options.ReshapeIsBitcast(operand->shape(), instr->shape())) { + return operand; + } + + // If the operand is a copy or reshape try to see if the operand's operand + // would produce a bitcast with initial instruction. + if (HloOpcode::kReshape == operand->opcode() || + HloOpcode::kCopy == operand->opcode()) { + return BitcastingOperandOfReshapeOrCopyChainHelper( + instr, operand->mutable_operand(0), options); + } + return nullptr; +} + +// Returns an operand of a chain of reshapes and copies that is bit-wise +// identical to first reshape or copy in the chain. +HloInstruction* BitcastingOperandOfReshapeOrCopyChain( + HloInstruction* instr, const AlgebraicSimplifierOptions& options) { + if (!options.is_layout_sensitive()) { + return nullptr; + } CHECK(HloOpcode::kReshape == instr->opcode() || HloOpcode::kCopy == instr->opcode()); - - const HloInstruction* operand = instr->operand(0); - // Can't insert bitcasts if the compiler used a memory layout which isn't - // compatible. - return ShapeUtil::ReshapeIsBitcast(operand->shape(), instr->shape()) && - valid_bitcast_callback(operand->shape(), instr->shape()); + return BitcastingOperandOfReshapeOrCopyChainHelper( + instr, instr->mutable_operand(0), options); } bool IsUnstridedSlice(const HloInstruction* hlo) { @@ -204,6 +221,8 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault { Status HandlePower(HloInstruction* power) override; + Status HandleRemainder(HloInstruction* remainder) override; + Status HandleReshape(HloInstruction* reshape) override; Status HandleReduce(HloInstruction* reduce) override; @@ -272,8 +291,11 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault { shape, hlo, zero, {dim}, AddReduce_computation)); } - // Convenience method for replacing an instruction with a bitcast. - void ReplaceWithBitcast(HloInstruction* instruction); + // Convenience method for replacing an instruction with a bitcast. If operand + // is not null, then the bitcast will use the specified operand instead of the + // operand of the instruction. + void ReplaceWithBitcast(HloInstruction* instruction, + HloInstruction* operand = nullptr); // Replace old instruction with new instruction if old and new instructions // have the same shape. Updates uses and root instruction. Returns whether a @@ -370,11 +392,6 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault { // Tries to convert slice(reshape(X)) into reshape(slice(X)) StatusOr TryToReorderSliceAndReshape(HloInstruction* slice); - // If the sort instruction has a tuple shape then looks for unused output - // values and removes them from the sort instruction. Returns true if the - // graph have been modified. - StatusOr RemoveUnusedOperandFromSort(HloInstruction* sort); - // Current HloComputation instance the AlgebraicSimplifierVisitor is // traversing. HloComputation* computation_; @@ -407,17 +424,19 @@ bool AlgebraicSimplifierVisitor::SameShape(const HloInstruction* lhs, } } -void AlgebraicSimplifierVisitor::ReplaceWithBitcast( - HloInstruction* instruction) { +void AlgebraicSimplifierVisitor::ReplaceWithBitcast(HloInstruction* instruction, + HloInstruction* operand) { CHECK_EQ(1, instruction->operand_count()); + if (operand == nullptr) { + operand = instruction->mutable_operand(0); + } CHECK_EQ(ShapeUtil::ElementsIn(instruction->shape()), - ShapeUtil::ElementsIn(instruction->operand(0)->shape())); + ShapeUtil::ElementsIn(operand->shape())); CHECK_EQ(ShapeUtil::ByteSizeOf(instruction->shape()), - ShapeUtil::ByteSizeOf(instruction->operand(0)->shape())); + ShapeUtil::ByteSizeOf(operand->shape())); - auto bitcast = computation_->AddInstruction( - HloInstruction::CreateUnary(instruction->shape(), HloOpcode::kBitcast, - instruction->mutable_operand(0))); + auto bitcast = computation_->AddInstruction(HloInstruction::CreateUnary( + instruction->shape(), HloOpcode::kBitcast, operand)); TF_CHECK_OK(ReplaceInstruction(instruction, bitcast)); } @@ -578,9 +597,9 @@ Status AlgebraicSimplifierVisitor::HandleCopy(HloInstruction* copy) { return Status::OK(); } - if (options_.is_layout_sensitive() && - ReshapeOrCopyIsBitcast(copy, options_.valid_bitcast_callback())) { - ReplaceWithBitcast(copy); + if (HloInstruction* bitcast_operand = + BitcastingOperandOfReshapeOrCopyChain(copy, options_)) { + ReplaceWithBitcast(copy, bitcast_operand); } return Status::OK(); @@ -797,10 +816,82 @@ Status InvertConstant(const HloInstruction& constant, Literal* result) { return T{1.0} / constant.literal().Get(indices); }); } + +template +std::unique_ptr TryDivideToShift(HloInstruction* divide, + HloComputation* computation) { + HloInstruction *a, *b, *c; + CHECK(Match(divide, m::Divide(m::Op(&a), m::Op(&b)))); + + if (ShapeUtil::ElementIsIntegral(divide->shape()) && + !Match(b, m::ConstantEffectiveScalar(&c)) && + !Match(b, m::Broadcast(m::ConstantEffectiveScalar(&c)))) { + return nullptr; + } + + if (ShapeUtil::ElementIsSigned(divide->shape())) { + int64 b_value = c->literal().GetFirstElement(); + if (b_value > 0 && IsPowerOfTwo(static_cast(b_value))) { + // Handle negative dividends by negating the result of the division. + HloInstruction* zero_like_a = BroadcastZeros( + computation, a->shape().element_type(), a->shape().dimensions()); + + auto* dividend_is_negative = + computation->AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::ChangeElementType(a->shape(), PRED), HloOpcode::kLt, a, + zero_like_a)); + + auto* negated_dividend = computation->AddInstruction( + HloInstruction::CreateUnary(a->shape(), HloOpcode::kNegate, a)); + + auto* abs_dividend = + computation->AddInstruction(HloInstruction::CreateTernary( + a->shape(), HloOpcode::kSelect, dividend_is_negative, + negated_dividend, a)); + + int log2_abs_b_value = tensorflow::Log2Floor64(b_value); + + auto* shift_amount = + computation->AddInstruction(HloInstruction::CreateConstant( + LiteralUtil::CreateR0(log2_abs_b_value))); + if (!ShapeUtil::IsScalar(b->shape())) { + shift_amount = computation->AddInstruction( + HloInstruction::CreateBroadcast(b->shape(), shift_amount, {})); + } + + auto* quotient = computation->AddInstruction(HloInstruction::CreateBinary( + divide->shape(), HloOpcode::kShiftRightLogical, abs_dividend, + shift_amount)); + + auto* neqated_quotient = + computation->AddInstruction(HloInstruction::CreateUnary( + quotient->shape(), HloOpcode::kNegate, quotient)); + + return HloInstruction::CreateTernary(divide->shape(), HloOpcode::kSelect, + dividend_is_negative, + neqated_quotient, quotient); + } + } else { + uint64 b_value = c->literal().GetFirstElement(); + if (IsPowerOfTwo(b_value)) { + int log2_abs_b_value = tensorflow::Log2Floor64(b_value); + HloInstruction* shift_amount = + computation->AddInstruction(HloInstruction::CreateConstant( + LiteralUtil::CreateR0(log2_abs_b_value))); + if (!ShapeUtil::IsScalar(b->shape())) { + shift_amount = computation->AddInstruction( + HloInstruction::CreateBroadcast(b->shape(), shift_amount, {})); + } + return HloInstruction::CreateBinary( + divide->shape(), HloOpcode::kShiftRightLogical, a, shift_amount); + } + } + + return nullptr; +} } // namespace Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { - Shape* shape; HloInstruction *a, *b, *c, *d; CHECK(Match(divide, m::Divide(m::Op(&a), m::Op(&b)))); // A/1 => A @@ -809,6 +900,61 @@ Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { return Status::OK(); } + // A / B => A >> log2(B) if B is a power of 2. + switch (divide->shape().element_type()) { + case S8: + if (std::unique_ptr shift = + TryDivideToShift(divide, computation_)) { + return ReplaceWithNewInstruction(divide, std::move(shift)); + } + break; + case S16: + if (std::unique_ptr shift = + TryDivideToShift(divide, computation_)) { + return ReplaceWithNewInstruction(divide, std::move(shift)); + } + break; + case S32: + if (std::unique_ptr shift = + TryDivideToShift(divide, computation_)) { + return ReplaceWithNewInstruction(divide, std::move(shift)); + } + break; + case S64: + if (std::unique_ptr shift = + TryDivideToShift(divide, computation_)) { + return ReplaceWithNewInstruction(divide, std::move(shift)); + } + break; + case U8: + if (std::unique_ptr shift = + TryDivideToShift(divide, computation_)) { + return ReplaceWithNewInstruction(divide, std::move(shift)); + } + break; + case U16: + if (std::unique_ptr shift = + TryDivideToShift(divide, computation_)) { + return ReplaceWithNewInstruction(divide, std::move(shift)); + } + break; + case U32: + if (std::unique_ptr shift = + TryDivideToShift(divide, computation_)) { + return ReplaceWithNewInstruction(divide, std::move(shift)); + } + break; + case U64: + if (std::unique_ptr shift = + TryDivideToShift(divide, computation_)) { + return ReplaceWithNewInstruction(divide, std::move(shift)); + } + break; + default: + break; + } + + Shape* shape; // exp(A)/exp(B) => exp(A-B) if (Match(divide, m::Divide(m::Exp(m::Op(&a)), m::Exp(m::Op(&b))) .WithShape(m::Shape(&shape)))) { @@ -859,8 +1005,9 @@ Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { // (Backends can do this transformation, but generally only if the constant is // a scalar.) if (Match(divide, m::Divide(m::NonConstant(&a), m::Constant(&b)))) { - Literal new_literal(b->shape()); - switch (b->shape().element_type()) { + Shape result_shape = b->literal().shape(); + Literal new_literal(result_shape); + switch (result_shape.element_type()) { case F16: TF_RETURN_IF_ERROR(InvertConstant(*b, &new_literal)); break; @@ -876,6 +1023,9 @@ Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { case C64: TF_RETURN_IF_ERROR(InvertConstant(*b, &new_literal)); break; + case C128: + TF_RETURN_IF_ERROR(InvertConstant(*b, &new_literal)); + break; default: return Status::OK(); } @@ -940,7 +1090,7 @@ StatusOr AlgebraicSimplifierVisitor::HandleDotStrengthReduction( const int64 rhs_rank = rhs->shape().rank(); const int64 lhs_rank = lhs->shape().rank(); const auto& dnums = dot->dot_dimension_numbers(); - if (dnums.rhs_contracting_dimensions_size() > 1) { + if (dnums.rhs_contracting_dimensions_size() != 1) { return false; } if (dot_rank > 2 && (lhs_rank != rhs_rank || lhs_rank != dot_rank)) { @@ -1379,6 +1529,9 @@ StatusOr AlgebraicSimplifierVisitor::OptimizeDotOfGather( // => output dimensions: DS ({M x N}, {0, start}, {M, 1}) => {M x 1}. bool lhs_is_dynamic_slice = lhs->opcode() == HloOpcode::kDynamicSlice; + HloDynamicSliceInstruction* dynamic_slice = + lhs_is_dynamic_slice ? Cast(lhs) + : Cast(rhs); // ctA: HloInstruction* left_operand = @@ -1396,8 +1549,6 @@ StatusOr AlgebraicSimplifierVisitor::OptimizeDotOfGather( HloInstruction::CreateDot(memoized_shape, left_operand, right_operand, dnums, dot->precision_config())); // Get pair {start, 0} or {0, start}. - HloInstruction* original_start_indices = - lhs_is_dynamic_slice ? lhs->mutable_operand(1) : rhs->mutable_operand(1); // Position of start: int index_of_non_zero_start = lhs_is_dynamic_slice ? 1 - lhs_contracting_dimension @@ -1406,23 +1557,19 @@ StatusOr AlgebraicSimplifierVisitor::OptimizeDotOfGather( int index_of_zero_start = 1 - index_of_non_zero_start; // Slice out start and 0 components and reorder if necessary. - auto indices_type = original_start_indices->shape().element_type(); + auto indices_type = dynamic_slice->operand(1)->shape().element_type(); Shape s_shape = ShapeUtil::MakeShape(indices_type, {1}); Shape d_shape = ShapeUtil::MakeShape(indices_type, {2}); HloInstruction* non_zero_start = - computation_->AddInstruction(HloInstruction::CreateSlice( - s_shape, original_start_indices, {index_of_non_zero_start}, - {index_of_non_zero_start + 1}, {1})); + dynamic_slice->mutable_operand(1 + index_of_non_zero_start); HloInstruction* zero_start = - computation_->AddInstruction(HloInstruction::CreateSlice( - s_shape, original_start_indices, {index_of_zero_start}, - {index_of_zero_start + 1}, {1})); - HloInstruction* new_start_indices = - lhs_is_dynamic_slice - ? computation_->AddInstruction(HloInstruction::CreateConcatenate( - d_shape, {non_zero_start, zero_start}, 0)) - : computation_->AddInstruction(HloInstruction::CreateConcatenate( - d_shape, {zero_start, non_zero_start}, 0)); + dynamic_slice->mutable_operand(1 + index_of_zero_start); + std::vector new_start_indices; + if (lhs_is_dynamic_slice) { + new_start_indices = {non_zero_start, zero_start}; + } else { + new_start_indices = {zero_start, non_zero_start}; + } // Build DynamicSlice(ctA x ctB). const int new_slice_m = lhs_is_dynamic_slice ? 1 : m; @@ -1438,7 +1585,9 @@ StatusOr AlgebraicSimplifierVisitor::OptimizeDotOfGather( Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot) { HloInstruction *lhs, *rhs; CHECK(Match(dot, m::Dot(m::Op(&lhs), m::Op(&rhs)))); - + if (options_.is_layout_sensitive()) { + return Status::OK(); + } // Replace a zero element dot with a broadcast of the constant 0. if (ShapeUtil::IsZeroElementArray(dot->shape()) || ShapeUtil::IsZeroElementArray(lhs->shape()) || @@ -1455,6 +1604,53 @@ Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot) { dot->shape().element_type() != BF16) { return Status::OK(); } + + // If there are no contracting dimensions, a dot can be rewritten as + // mul(broadcast(transpose(x)),broadcast(transpose(y))) + if (dot->dot_dimension_numbers().lhs_contracting_dimensions_size() == 0) { + std::vector lhs_transpose( + dot->dot_dimension_numbers().lhs_batch_dimensions().begin(), + dot->dot_dimension_numbers().lhs_batch_dimensions().end()); + for (int64 i = 0; i < lhs->shape().rank(); ++i) { + if (!absl::c_linear_search( + dot->dot_dimension_numbers().lhs_batch_dimensions(), i)) { + lhs_transpose.push_back(i); + } + } + TF_ASSIGN_OR_RETURN(HloInstruction * new_lhs, + MakeTransposeHlo(lhs, lhs_transpose)); + if (dot->shape().rank() != lhs->shape().rank()) { + std::vector lhs_broadcast_dims(lhs->shape().rank()); + absl::c_iota(lhs_broadcast_dims, 0); + new_lhs = computation_->AddInstruction(HloInstruction::CreateBroadcast( + dot->shape(), new_lhs, lhs_broadcast_dims)); + } + std::vector rhs_transpose( + dot->dot_dimension_numbers().rhs_batch_dimensions().begin(), + dot->dot_dimension_numbers().rhs_batch_dimensions().end()); + for (int64 i = 0; i < rhs->shape().rank(); ++i) { + if (!absl::c_linear_search( + dot->dot_dimension_numbers().rhs_batch_dimensions(), i)) { + rhs_transpose.push_back(i); + } + } + TF_ASSIGN_OR_RETURN(HloInstruction * new_rhs, + MakeTransposeHlo(rhs, rhs_transpose)); + if (dot->shape().rank() != rhs->shape().rank()) { + std::vector rhs_broadcast_dims( + dot->dot_dimension_numbers().lhs_batch_dimensions_size()); + absl::c_iota(rhs_broadcast_dims, 0); + for (int64 i = lhs->shape().rank(); i < dot->shape().rank(); ++i) { + rhs_broadcast_dims.push_back(i); + } + new_rhs = computation_->AddInstruction(HloInstruction::CreateBroadcast( + dot->shape(), new_rhs, rhs_broadcast_dims)); + } + return ReplaceWithNewInstruction( + dot, HloInstruction::CreateBinary(dot->shape(), HloOpcode::kMultiply, + new_lhs, new_rhs)); + } + if (lhs->shape().rank() > 2 || rhs->shape().rank() > 2 || dot->shape().rank() > 2) { if (options_.enable_dot_strength_reduction() && @@ -1493,7 +1689,11 @@ Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot) { } // Simplify dot(transpose(a), transpose(b)) to transpose(dot(b,a)). - if (lhs->IsRank2Transpose() && rhs->IsRank2Transpose()) { + if (dot->dot_dimension_numbers().lhs_batch_dimensions_size() == 0 && + dot->dot_dimension_numbers().lhs_contracting_dimensions_size() == 1 && + dot->dot_dimension_numbers().lhs_contracting_dimensions(0) == 1 && + dot->dot_dimension_numbers().rhs_contracting_dimensions(0) == 0 && + lhs->IsRank2Transpose() && rhs->IsRank2Transpose()) { DotDimensionNumbers dot_dimension_numbers; dot_dimension_numbers.add_lhs_contracting_dimensions(1); dot_dimension_numbers.add_rhs_contracting_dimensions(0); @@ -2144,14 +2344,151 @@ AlgebraicSimplifierVisitor::TryToSinkBroadcastAfterOpWithUniqueNonScalarOperand( return changed; } +namespace { +template +std::unique_ptr TryRemainderToAnd(HloInstruction* remainder, + HloComputation* computation) { + HloInstruction *a, *b, *c; + CHECK(Match(remainder, m::Remainder(m::Op(&a), m::Op(&b)))); + + if (ShapeUtil::ElementIsIntegral(remainder->shape()) && + !Match(b, m::ConstantEffectiveScalar(&c)) && + !Match(b, m::Broadcast(m::ConstantEffectiveScalar(&c)))) { + return nullptr; + } + + if (ShapeUtil::ElementIsSigned(remainder->shape())) { + int64 b_value = c->literal().GetFirstElement(); + if (b_value > 0 && IsPowerOfTwo(static_cast(b_value))) { + // Handle negative dividends by negating the result of the division. + HloInstruction* zero_like_a = BroadcastZeros( + computation, a->shape().element_type(), a->shape().dimensions()); + + auto* dividend_is_negative = + computation->AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::ChangeElementType(a->shape(), PRED), HloOpcode::kLt, a, + zero_like_a)); + + auto* negated_dividend = computation->AddInstruction( + HloInstruction::CreateUnary(a->shape(), HloOpcode::kNegate, a)); + + auto* abs_dividend = + computation->AddInstruction(HloInstruction::CreateTernary( + a->shape(), HloOpcode::kSelect, dividend_is_negative, + negated_dividend, a)); + + auto* mask_amount = + computation->AddInstruction(HloInstruction::CreateConstant( + LiteralUtil::CreateR0(b_value - 1))); + if (!ShapeUtil::IsScalar(b->shape())) { + mask_amount = computation->AddInstruction( + HloInstruction::CreateBroadcast(b->shape(), mask_amount, {})); + } + + auto* quotient = computation->AddInstruction(HloInstruction::CreateBinary( + remainder->shape(), HloOpcode::kAnd, abs_dividend, mask_amount)); + + auto* neqated_quotient = + computation->AddInstruction(HloInstruction::CreateUnary( + quotient->shape(), HloOpcode::kNegate, quotient)); + + return HloInstruction::CreateTernary( + remainder->shape(), HloOpcode::kSelect, dividend_is_negative, + neqated_quotient, quotient); + } + } else { + uint64 b_value = c->literal().GetFirstElement(); + if (IsPowerOfTwo(b_value)) { + HloInstruction* mask_amount = + computation->AddInstruction(HloInstruction::CreateConstant( + LiteralUtil::CreateR0(b_value - 1))); + if (!ShapeUtil::IsScalar(b->shape())) { + mask_amount = computation->AddInstruction( + HloInstruction::CreateBroadcast(b->shape(), mask_amount, {})); + } + return HloInstruction::CreateBinary(remainder->shape(), HloOpcode::kAnd, + a, mask_amount); + } + } + return nullptr; +} +} // namespace + +Status AlgebraicSimplifierVisitor::HandleRemainder(HloInstruction* remainder) { + HloInstruction *a, *b; + CHECK(Match(remainder, m::Remainder(m::Op(&a), m::Op(&b)))); + + // A % B => A & (B - 1) if B is a power of 2. + switch (remainder->shape().element_type()) { + case S8: + if (std::unique_ptr shift = + TryRemainderToAnd(remainder, computation_)) { + return ReplaceWithNewInstruction(remainder, std::move(shift)); + } + break; + case S16: + if (std::unique_ptr shift = + TryRemainderToAnd(remainder, computation_)) { + return ReplaceWithNewInstruction(remainder, std::move(shift)); + } + break; + case S32: + if (std::unique_ptr shift = + TryRemainderToAnd(remainder, computation_)) { + return ReplaceWithNewInstruction(remainder, std::move(shift)); + } + break; + case S64: + if (std::unique_ptr shift = + TryRemainderToAnd(remainder, computation_)) { + return ReplaceWithNewInstruction(remainder, std::move(shift)); + } + break; + case U8: + if (std::unique_ptr shift = + TryRemainderToAnd(remainder, computation_)) { + return ReplaceWithNewInstruction(remainder, std::move(shift)); + } + break; + case U16: + if (std::unique_ptr shift = + TryRemainderToAnd(remainder, computation_)) { + return ReplaceWithNewInstruction(remainder, std::move(shift)); + } + break; + case U32: + if (std::unique_ptr shift = + TryRemainderToAnd(remainder, computation_)) { + return ReplaceWithNewInstruction(remainder, std::move(shift)); + } + break; + case U64: + if (std::unique_ptr shift = + TryRemainderToAnd(remainder, computation_)) { + return ReplaceWithNewInstruction(remainder, std::move(shift)); + } + break; + default: + break; + } + + return Status::OK(); +} + Status AlgebraicSimplifierVisitor::HandleReshape(HloInstruction* reshape) { auto operand = reshape->mutable_operand(0); // Reshape directly to empty constant if the shape contains zero-element // dimension. if (ShapeUtil::IsZeroElementArray(reshape->shape())) { + // If the instruction doesn't have a layout, use a default layout for + // the literal result. + Shape reshaped_shape = reshape->shape(); + if (!LayoutUtil::HasLayout(reshaped_shape)) { + LayoutUtil::SetToDefaultLayout(&reshaped_shape); + } auto empty_constant = HloInstruction::CreateConstant( - Literal::CreateFromShape(reshape->shape())); + Literal::CreateFromShape(reshaped_shape)); return ReplaceWithNewInstruction(reshape, std::move(empty_constant)); } @@ -2200,12 +2537,10 @@ Status AlgebraicSimplifierVisitor::HandleReshape(HloInstruction* reshape) { } // Make this a bitcast if possible. - if (options_.is_layout_sensitive() && - ReshapeOrCopyIsBitcast(reshape, options_.valid_bitcast_callback())) { - ReplaceWithBitcast(reshape); - return Status::OK(); + if (HloInstruction* bitcast_operand = + BitcastingOperandOfReshapeOrCopyChain(reshape, options_)) { + ReplaceWithBitcast(reshape, bitcast_operand); } - return Status::OK(); } @@ -2215,8 +2550,7 @@ Status AlgebraicSimplifierVisitor::HandleReverse(HloInstruction* reverse) { auto dim_is_one = [&](int64 i) -> bool { return reverse->shape().dimensions(i) == 1; }; - if (std::all_of(reverse->dimensions().begin(), reverse->dimensions().end(), - dim_is_one)) { + if (absl::c_all_of(reverse->dimensions(), dim_is_one)) { return ReplaceInstruction(reverse, reverse->mutable_operand(0)); } return Status::OK(); @@ -2440,28 +2774,72 @@ Status AlgebraicSimplifierVisitor::HandleDynamicUpdateSlice( return Status::OK(); } -Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* reduce) { - // TODO(b/112040122): Most of those optimizations can be done for multi-output - // reduces. - if (reduce->shape().IsTuple()) { - return Status::OK(); - } +Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* hlo) { + HloReduceInstruction* reduce = Cast(hlo); + bool multi_output_reduce = reduce->shape().IsTuple(); + + // For tuple reduce, we require all reduce shapes to be the same, up to the + // element types, so we can just the first operand and the first result as a + // representative. + auto arg = reduce->inputs()[0]; + auto init_value = reduce->init_values()[0]; + const Shape& reduce_result_shape = + multi_output_reduce ? reduce->shape().tuple_shapes(0) : reduce->shape(); - auto arg = reduce->mutable_operand(0); - auto init_value = reduce->mutable_operand(1); absl::Span dimensions(reduce->dimensions()); HloComputation* function = reduce->to_apply(); if (ShapeUtil::IsZeroElementArray(arg->shape()) || - ShapeUtil::IsZeroElementArray(reduce->shape())) { - return ReplaceWithNewInstruction( - reduce, - HloInstruction::CreateBroadcast(reduce->shape(), init_value, {})); + ShapeUtil::IsZeroElementArray(reduce_result_shape)) { + if (multi_output_reduce) { + std::vector broadcast_inits; + int64 inputs = reduce->input_count(); + for (int64 i = 0; i < inputs; ++i) { + broadcast_inits.push_back(computation_->AddInstruction( + HloInstruction::CreateBroadcast(reduce->shape().tuple_shapes(i), + reduce->init_values()[i], {}))); + } + return ReplaceWithNewInstruction( + reduce, HloInstruction::CreateTuple(broadcast_inits)); + } else { + return ReplaceWithNewInstruction( + reduce, + HloInstruction::CreateBroadcast(reduce_result_shape, init_value, {})); + } + } + + // If the reduction results in the same number of elements, then the only + // possible side effect would be a reshape. Since the init_value is an + // identity of the reduction function, we can therefore replace the reduce + // with a simple reshape, ignoring the reduction function completely. + if (ShapeUtil::ElementsIn(reduce_result_shape) == + ShapeUtil::ElementsIn(arg->shape())) { + if (multi_output_reduce) { + std::vector reshaped_args; + int64 inputs = reduce->input_count(); + for (int64 i = 0; i < inputs; ++i) { + reshaped_args.push_back( + computation_->AddInstruction(HloInstruction::CreateReshape( + reduce->shape().tuple_shapes(i), reduce->inputs()[i]))); + } + return ReplaceWithNewInstruction( + reduce, HloInstruction::CreateTuple(reshaped_args)); + } else { + return ReplaceWithNewInstruction( + reduce, HloInstruction::CreateReshape(reduce_result_shape, arg)); + } + } + + // TODO(b/112040122): Most of those optimizations below can be done for + // multi-output reduces. + if (multi_output_reduce) { + return Status::OK(); } // A Transpose feeding a reduce can simply permute the reduction dimensions // field if the output of the reduce is a vector or scalar. Higher ranked // result may require a transpose of the output. - if (reduce->shape().rank() <= 1 && arg->opcode() == HloOpcode::kTranspose) { + if (reduce_result_shape.rank() <= 1 && + arg->opcode() == HloOpcode::kTranspose) { auto transpose_dimensions = arg->dimensions(); std::vector new_reduce_dimensions; for (auto dim : dimensions) { @@ -2469,20 +2847,10 @@ Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* reduce) { } return ReplaceWithNewInstruction( reduce, HloInstruction::CreateReduce( - reduce->shape(), arg->mutable_operand(0), init_value, + reduce_result_shape, arg->mutable_operand(0), init_value, new_reduce_dimensions, function)); } - // If the reduction results in the same number of elements, then the only - // possible side effect would be a reshape. Since the init_value is an - // identity of the reduction function, we can therefore replace the reduce - // with a simple reshape, ignoring the reduction function completely. - if (ShapeUtil::ElementsIn(reduce->shape()) == - ShapeUtil::ElementsIn(arg->shape())) { - return ReplaceWithNewInstruction( - reduce, HloInstruction::CreateReshape(reduce->shape(), arg)); - } - // If a reduce feeds a reduce with the same computation and initial value, // they can be combined into a single reduce. if (arg->opcode() == HloOpcode::kReduce && @@ -2491,9 +2859,9 @@ Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* reduce) { // Create a new reduce with the combined reduction dimensions of both // reduces. std::vector arg_dims = arg->dimensions(); - std::sort(arg_dims.begin(), arg_dims.end()); + absl::c_sort(arg_dims); std::vector reduce_dims = reduce->dimensions(); - std::sort(reduce_dims.begin(), reduce_dims.end()); + absl::c_sort(reduce_dims); // Transform reduce_dims to the same rank as the operand of the operand. for (int64 arg_dim : arg_dims) { for (int64& dim : reduce_dims) { @@ -2508,9 +2876,9 @@ Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* reduce) { std::merge(arg_dims.begin(), arg_dims.end(), reduce_dims.begin(), reduce_dims.end(), std::back_inserter(new_dimensions)); return ReplaceWithNewInstruction( - reduce, - HloInstruction::CreateReduce(reduce->shape(), arg->mutable_operand(0), - init_value, new_dimensions, function)); + reduce, HloInstruction::CreateReduce( + reduce_result_shape, arg->mutable_operand(0), init_value, + new_dimensions, function)); } // A reshape that collapses multiple dimensions into a dimension being @@ -2539,7 +2907,7 @@ Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* reduce) { } if (can_move_reshape_into_reduce) { changed_ = true; - std::unordered_set dimensions_not_to_reduce; + absl::flat_hash_set dimensions_not_to_reduce; for (auto dim_pair : unmodified_dims) { if (arg_dim_in_output[dim_pair.second]) { dimensions_not_to_reduce.insert(dim_pair.first); @@ -2547,13 +2915,13 @@ Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* reduce) { } std::vector new_reduce_dimensions; for (int64 i = 0; i < arg->operand(0)->shape().rank(); ++i) { - if (dimensions_not_to_reduce.count(i) == 0) { + if (!dimensions_not_to_reduce.contains(i)) { new_reduce_dimensions.push_back(i); } } return ReplaceWithNewInstruction( reduce, HloInstruction::CreateReduce( - reduce->shape(), arg->mutable_operand(0), init_value, + reduce_result_shape, arg->mutable_operand(0), init_value, new_reduce_dimensions, function)); } } @@ -2568,11 +2936,11 @@ Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* reduce) { HloInstruction* old_reduce = nullptr; for (HloInstruction* operand : arg->operands()) { HloInstruction* new_reduce = computation_->AddInstruction( - HloInstruction::CreateReduce(reduce->shape(), operand, init_value, + HloInstruction::CreateReduce(reduce_result_shape, operand, init_value, reduce->dimensions(), function)); if (old_reduce != nullptr) { new_reduce = computation_->AddInstruction(HloInstruction::CreateMap( - reduce->shape(), {old_reduce, new_reduce}, function)); + reduce_result_shape, {old_reduce, new_reduce}, function)); } old_reduce = new_reduce; } @@ -2822,69 +3190,6 @@ Status AlgebraicSimplifierVisitor::HandleSelect(HloInstruction* select) { return Status::OK(); } -StatusOr AlgebraicSimplifierVisitor::RemoveUnusedOperandFromSort( - HloInstruction* sort) { - if (!sort->shape().IsTuple()) { - return false; - } - - if (sort->parent()->root_instruction() == sort) { - // Can't analyse users of the root instruction. - return false; - } - - // Index 0 is the sorting key used by the sort HLO itself. - absl::flat_hash_set used_indices{0}; - for (const HloInstruction* user : sort->users()) { - if (user->opcode() != HloOpcode::kGetTupleElement) { - // Can't analyse users other then get-tuple-element. - return false; - } - used_indices.insert(user->tuple_index()); - } - - if (used_indices.size() == sort->operand_count()) { - // All operands are used. - return false; - } - - std::vector operands{sort->mutable_operand(0)}; - std::vector new_shapes{sort->operand(0)->shape()}; - for (int64 i = 1; i < sort->operand_count(); ++i) { - if (used_indices.count(i)) { - operands.push_back(sort->mutable_operand(i)); - new_shapes.push_back(sort->operand(i)->shape()); - } - } - Shape new_sort_shape = new_shapes.size() == 1 - ? new_shapes[0] - : ShapeUtil::MakeTupleShape(new_shapes); - HloInstruction* new_sort = computation_->AddInstruction( - sort->CloneWithNewOperands(new_sort_shape, operands)); - - // Map from original get-tuple-element tuple index to new HLO instruction - absl::flat_hash_map result_map; - if (new_sort->shape().IsTuple()) { - // Old sort key maps to new sort key. - int64 new_index = 0; - for (int64 i = 0; i < sort->operand_count(); ++i) { - if (used_indices.count(i)) { - result_map[i] = - computation_->AddInstruction(HloInstruction::CreateGetTupleElement( - new_shapes[new_index], new_sort, new_index)); - ++new_index; - } - } - } else { - result_map[0] = new_sort; - } - for (HloInstruction* user : sort->users()) { - TF_RETURN_IF_ERROR( - user->ReplaceAllUsesWith(result_map.at(user->tuple_index()))); - } - return true; -} - Status AlgebraicSimplifierVisitor::HandleSort(HloInstruction* sort) { auto operand = sort->mutable_operand(0); int64 dimension_to_sort = sort->dimensions(0); @@ -2897,116 +3202,6 @@ Status AlgebraicSimplifierVisitor::HandleSort(HloInstruction* sort) { return ReplaceWithNewInstruction( sort, HloInstruction::CreateTuple(sort->operands())); } - - // Remove the unused values from a key-value sort. - TF_ASSIGN_OR_RETURN(bool removed_operand, RemoveUnusedOperandFromSort(sort)); - if (removed_operand) { - changed_ = true; - return Status::OK(); - } - - if (!options_.enable_permutation_sort_replacement()) { - return Status::OK(); - } - // Check if we are sorting a permutation. In that case, we know that the keys - // will be sorted to the identity permutation, and we can represent the - // changes to the 'values' parameter as a scatter. - if (sort->operand_count() == 2 && - operand->opcode() == HloOpcode::kGetTupleElement) { - const HloInstruction* other_sort = operand->operand(0); - // Check whether the 'values' parameter is the result of another sort with - // the same sort dimension. - if (other_sort->opcode() == HloOpcode::kSort && - other_sort->operand_count() >= 2 && - other_sort->dimensions(0) == dimension_to_sort && - other_sort->operand(operand->tuple_index())->opcode() == - HloOpcode::kIota) { - auto* iota = - Cast(other_sort->operand(operand->tuple_index())); - // The sort operand needs to be an integral iota, and the iota dimension - // needs to be the dimension that was sorted. - if (iota->iota_dimension() == dimension_to_sort && - ShapeUtil::ElementIsIntegral(iota->shape())) { - // We use the following construction method for a Scatter that applies - // the permutation from 'keys' to the 'values' parameter. - // - Take the "keys" parameter of the second sort and reshape it to have - // another "1" dimension at the end. - // - Concatenate it with iotas of the same extended shape with all - // different iota_dimensions except the dimension_to_sort in the order - // of iota_dimensions/dimension_to_sort, so e.g. with rank 3 and - // dimension_to_sort = 1, we would have concatenate of (iota with - // iota_dimension=0, keys, iota with iota_dimension = 2) - // - Use this as the indices parameter of scatter, and set updates - // of the scatter to be a reshaped 'values' parameter of sort (adding - // 'rank' many 1 dimensions at the end). - int64 rank = operand->shape().rank(); - Shape extended_shape = operand->shape(); - extended_shape.add_dimensions(1); - extended_shape.mutable_layout()->add_minor_to_major(rank); - auto reshaped_permutation = computation_->AddInstruction( - HloInstruction::CreateReshape(extended_shape, operand)); - std::vector concat_operands; - for (int64 i = 0; i < rank; ++i) { - if (i == dimension_to_sort) { - concat_operands.push_back(reshaped_permutation); - } else { - concat_operands.push_back(computation_->AddInstruction( - HloInstruction::CreateIota(extended_shape, i))); - } - } - Shape concat_shape = operand->shape(); - concat_shape.add_dimensions(rank); - concat_shape.mutable_layout()->add_minor_to_major(rank); - auto scatter_indices = - rank > 1 ? computation_->AddInstruction( - HloInstruction::CreateConcatenate( - concat_shape, concat_operands, rank)) - : reshaped_permutation; - - // We don't care about the operand, it will be completely overridden by - // the updates. - auto scatter_operand = computation_->AddInstruction( - HloInstruction::CreateIota(sort->operand(1)->shape(), 0)); - - // Construct the updates operand of scatter. - Shape update_shape = sort->operand(1)->shape(); - for (int64 i = 0; i < rank; ++i) { - update_shape.add_dimensions(1); - update_shape.mutable_layout()->add_minor_to_major(rank + i); - } - auto scatter_updates = - computation_->AddInstruction(HloInstruction::CreateReshape( - update_shape, sort->mutable_operand(1))); - - // Construct the updates computation, which simply replaces the operand - // values with the update values. - HloComputation::Builder b("update_replace_computation"); - Shape scalar_shape = ShapeUtil::MakeShape(S32, {}); - b.AddInstruction( - HloInstruction::CreateParameter(0, scalar_shape, "scalar_lhs")); - auto scalar_rhs = b.AddInstruction( - HloInstruction::CreateParameter(1, scalar_shape, "scalar_rhs")); - auto update_replace_computation = - computation_->parent()->AddEmbeddedComputation(b.Build(scalar_rhs)); - - ScatterDimensionNumbers dim_numbers; - dim_numbers.set_index_vector_dim(rank); - for (int64 i = 0; i < rank; ++i) { - dim_numbers.add_update_window_dims(rank + i); - dim_numbers.add_scatter_dims_to_operand_dims(i); - } - auto scatter = - computation_->AddInstruction(HloInstruction::CreateScatter( - sort->operand(1)->shape(), scatter_operand, scatter_indices, - scatter_updates, update_replace_computation, dim_numbers)); - return ReplaceWithNewInstruction( - sort, HloInstruction::CreateTuple( - {computation_->AddInstruction(HloInstruction::CreateIota( - operand->shape(), dimension_to_sort)), - scatter})); - } - } - } return Status::OK(); } @@ -3298,15 +3493,6 @@ StatusOr AlgebraicSimplifierVisitor::SimplifyConvToDot( const Shape dot_output_shape = ShapeUtil::MakeShapeWithDescendingLayout( convolution_shape.element_type(), {conv_width, output_channels}); - // We cannot insert bitcasts if the layouts will not be compatible. - // TODO(b/33178038): Consider inserting a transpose if a bitcast would be - // invalid. - if (!options_.valid_bitcast_callback()(input_shape, new_input_shape) || - !options_.valid_bitcast_callback()(filter_shape, new_filter_shape) || - !options_.valid_bitcast_callback()(dot_output_shape, convolution_shape)) { - return false; - } - auto new_lhs = add_bitcast(new_input_shape, lhs); auto new_rhs = add_bitcast(new_filter_shape, rhs); DotDimensionNumbers dot_dimension_numbers; diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.h b/tensorflow/compiler/xla/service/algebraic_simplifier.h index 1acaadaaa03803e467fd9bf6b22f5793fd6a1ec9..df5a8c2ec141458a95fafb76b1e99e4b04a61b28 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier.h +++ b/tensorflow/compiler/xla/service/algebraic_simplifier.h @@ -25,21 +25,25 @@ namespace xla { class AlgebraicSimplifierOptions { public: - // Given shapes 'from_shape' and 'to_shape', determines if it is valid to - // bitcast from 'from_shape' to 'to_shape' after considering platform - // dependent effects on layout like alignment restrictions. Precondition: the - // two shapes have layouts, the same number of elements and - // ShapeUtil::ReshapeIsBitcast returns true. - using ValidBitcastCallback = + AlgebraicSimplifierOptions() {} + // Platform dependent callback to determine if a reshape `from_shape` to + // `to_shape` is a bitcast. + using ReshapeIsBitcastCallback = std::function; - explicit AlgebraicSimplifierOptions( - ValidBitcastCallback valid_bitcast_callback) - : valid_bitcast_callback_(std::move(valid_bitcast_callback)) {} - // If valid_bitcast_callback returns true, then the pass will replace reshapes - // and transposes with bitcasts. - const ValidBitcastCallback& valid_bitcast_callback() const { - return valid_bitcast_callback_; + ReshapeIsBitcastCallback reshape_is_bitcast_callback) + : reshape_is_bitcast_callback_(std::move(reshape_is_bitcast_callback)) {} + + // Use the platform specific callback if set. It is not sensible to return + // true here if the options are not layout sensitive. + bool ReshapeIsBitcast(const Shape& from_shape, const Shape& to_shape) const { + if (!is_layout_sensitive_) { + return false; + } + if (!reshape_is_bitcast_callback_) { + return ShapeUtil::ReshapeIsBitcast(from_shape, to_shape); + } + return reshape_is_bitcast_callback_(from_shape, to_shape); } // If is_layout_sensitive is true, then the simplifier preserves layout during @@ -47,12 +51,14 @@ class AlgebraicSimplifierOptions { void set_is_layout_sensitive(bool is_layout_sensitive) { is_layout_sensitive_ = is_layout_sensitive; } + bool is_layout_sensitive() const { return is_layout_sensitive_; } // Enable dot simplification on platforms where it is profitable. void set_enable_dot_strength_reduction(bool enable_dot_strength_reduction) { enable_dot_strength_reduction_ = enable_dot_strength_reduction; } + bool enable_dot_strength_reduction() const { return enable_dot_strength_reduction_; } @@ -65,16 +71,6 @@ class AlgebraicSimplifierOptions { return enable_conv_simplification_; } - // If enable_permutation_sort_replacement is true, a sort op that is known to - // sort a permutation will be replaced with a scatter op. - void set_enable_permutation_sort_replacement( - bool enable_permutation_sort_replacement) { - enable_permutation_sort_replacement_ = enable_permutation_sort_replacement; - } - bool enable_permutation_sort_replacement() const { - return enable_permutation_sort_replacement_; - } - // If enable_window_reduce_replacement is true, the kReduceWindow instruction // can be optimized by replacement with simpler operations. void set_enable_window_reduce_to_reduce_replacement( @@ -82,16 +78,16 @@ class AlgebraicSimplifierOptions { enable_window_reduce_to_reduce_replacement_ = enable_window_reduce_to_reduce_replacement; } + bool enable_window_reduce_to_reduce_replacement() const { return enable_window_reduce_to_reduce_replacement_; } private: - ValidBitcastCallback valid_bitcast_callback_; + ReshapeIsBitcastCallback reshape_is_bitcast_callback_; bool is_layout_sensitive_{false}; bool enable_dot_strength_reduction_{true}; bool enable_conv_simplification_{true}; - bool enable_permutation_sort_replacement_{false}; bool enable_window_reduce_to_reduce_replacement_{true}; }; diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc index 211c5bf05a0d7ac533ce483f194db309991076be..feb6a0fb79538b38afa1110296b52061ec7f2259 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc @@ -25,6 +25,7 @@ limitations under the License. #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/service/hlo_casting_utils.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_creation_utils.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_instructions.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" @@ -46,17 +47,9 @@ namespace { using ::testing::ElementsAre; namespace m = match; -AlgebraicSimplifierOptions::ValidBitcastCallback bitcasting_callback() { - return [](const Shape&, const Shape&) { return true; }; -} - -AlgebraicSimplifierOptions::ValidBitcastCallback non_bitcasting_callback() { - return [](const Shape&, const Shape&) { return false; }; -} - class AlgebraicSimplifierTest : public HloTestBase { protected: - AlgebraicSimplifierOptions default_options_{non_bitcasting_callback()}; + AlgebraicSimplifierOptions default_options_; }; // Test that A + 0 is simplified to A @@ -202,6 +195,86 @@ TEST_F(AlgebraicSimplifierTest, FactorFpAdditionBfloat16) { m::Broadcast(m::ConstantScalar(0.125))))); } +TEST_F(AlgebraicSimplifierTest, UnsignedDivideByPowerOf2) { + const char* kModuleStr = R"( + HloModule m + test { + p = u32[4] parameter(0) + c = u32[] constant(8) + b = u32[4] broadcast(c), dimensions={} + ROOT d = u32[4] divide(p, b) + } + )"; + 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::ShiftRightLogical( + m::Parameter(0), m::Broadcast(m::ConstantScalar(3))))); +} + +TEST_F(AlgebraicSimplifierTest, SignedDivideByPowerOf2) { + const char* kModuleStr = R"( + HloModule m + test { + p = s32[4] parameter(0) + c = s32[] constant(8) + b = s32[4] broadcast(c), dimensions={} + ROOT d = s32[4] divide(p, b) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).ValueOrDie()); + auto match_dividend_is_negative = + m::Lt(m::Parameter(0), m::Broadcast(m::ConstantScalar(0))); + auto match_abs = m::Select(match_dividend_is_negative, + m::Negate(m::Parameter(0)), m::Parameter(0)); + auto match_shift = + m::ShiftRightLogical(match_abs, m::Broadcast(m::ConstantScalar(3))); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::Select(match_dividend_is_negative, + m::Negate(match_shift), match_shift))); +} + +TEST_F(AlgebraicSimplifierTest, UnsignedRemainderByPowerOf2) { + const char* kModuleStr = R"( + HloModule m + test { + p = u32[4] parameter(0) + c = u32[] constant(8) + b = u32[4] broadcast(c), dimensions={} + ROOT r = u32[4] remainder(p, b) + } + )"; + 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::AndAnyOrder(m::Parameter(0), + m::Broadcast(m::ConstantScalar(7))))); +} + +TEST_F(AlgebraicSimplifierTest, SignedRemainderByPowerOf2) { + const char* kModuleStr = R"( + HloModule m + test { + p = s32[4] parameter(0) + c = s32[] constant(8) + b = s32[4] broadcast(c), dimensions={} + ROOT r = s32[4] remainder(p, b) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).ValueOrDie()); + auto match_dividend_is_negative = + m::Lt(m::Parameter(0), m::Broadcast(m::ConstantScalar(0))); + auto match_abs = m::Select(match_dividend_is_negative, + m::Negate(m::Parameter(0)), m::Parameter(0)); + auto match_and = + m::AndAnyOrder(match_abs, m::Broadcast(m::ConstantScalar(7))); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::Select(match_dividend_is_negative, + m::Negate(match_and), match_and))); +} + // Test that A * 0 is simplified to 0 TEST_F(AlgebraicSimplifierTest, MulZero) { auto m = CreateNewVerifiedModule(); @@ -1464,23 +1537,77 @@ TEST_F(AlgebraicSimplifierTest, RemoveCopy) { EXPECT_THAT(computation->root_instruction(), param0); } -TEST_F(AlgebraicSimplifierTest, CopyEqualsBitcast) { +TEST_F(AlgebraicSimplifierTest, CopyOfReshapeOfCopyEqualsBitcast) { auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloInstruction* param = builder.AddInstruction(HloInstruction::CreateParameter( - 0, ShapeUtil::MakeShape(F32, {1, 14, 14, 64}), "param")); - *param->mutable_shape()->mutable_layout() = - LayoutUtil::MakeLayout({0, 1, 2, 3}); + 0, ShapeUtil::MakeShapeWithLayout(F32, {1, 14, 14, 64}, {3, 2, 1, 0}), + "param")); HloInstruction* copy = builder.AddInstruction(HloInstruction::CreateUnary( - ShapeUtil::MakeShape(F32, {1, 14, 14, 64}), HloOpcode::kCopy, param)); - *copy->mutable_shape()->mutable_layout() = - LayoutUtil::MakeLayout({1, 2, 0, 3}); + ShapeUtil::MakeShapeWithLayout(F32, {1, 14, 14, 64}, {0, 1, 2, 3}), + HloOpcode::kCopy, param)); + HloInstruction* reshape = + builder.AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShapeWithLayout(F32, {14 * 14, 64}, {0, 1}), copy)); + builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShapeWithLayout(F32, {14 * 14, 64}, {1, 0}), + HloOpcode::kCopy, reshape)); + auto computation = m->AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Copy(m::Reshape(m::Copy(m::Parameter(0)))))); + + AlgebraicSimplifierOptions options; + options.set_is_layout_sensitive(true); + AlgebraicSimplifier simplifier(options); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + // Verify that the copy of reshape of copy is replaced. + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Bitcast(m::Parameter(0)))); +} + +TEST_F(AlgebraicSimplifierTest, ReshapeOfCopyEqualsBitcast) { + auto m = CreateNewVerifiedModule(); + HloComputation::Builder builder(TestName()); + HloInstruction* param = + builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShapeWithLayout(F32, {1, 14, 14, 64}, {3, 2, 1, 0}), + "param")); + HloInstruction* copy = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShapeWithLayout(F32, {1, 14, 14, 64}, {0, 1, 2, 3}), + HloOpcode::kCopy, param)); + builder.AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShapeWithLayout(F32, {14 * 14, 64}, {1, 0}), copy)); + + auto computation = m->AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Reshape(m::Copy(m::Parameter(0))))); + + AlgebraicSimplifierOptions options; + options.set_is_layout_sensitive(true); + AlgebraicSimplifier simplifier(options); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + // Verify that the copy of reshape of copy is replaced. + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Bitcast(m::Parameter(0)))); +} + +TEST_F(AlgebraicSimplifierTest, CopyEqualsBitcast) { + auto m = CreateNewVerifiedModule(); + HloComputation::Builder builder(TestName()); + HloInstruction* param = + builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShapeWithLayout(F32, {1, 14, 14, 64}, {0, 1, 2, 3}), + "param")); + builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShapeWithLayout(F32, {1, 14, 14, 64}, {1, 2, 0, 3}), + HloOpcode::kCopy, param)); auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Copy(m::Parameter(0)))); - AlgebraicSimplifierOptions options(non_bitcasting_callback()); + AlgebraicSimplifierOptions options( + [](const Shape&, const Shape&) { return false; }); options.set_is_layout_sensitive(true); AlgebraicSimplifier simplifier1(options); ASSERT_FALSE(simplifier1.Run(m.get()).ValueOrDie()); @@ -1488,10 +1615,10 @@ TEST_F(AlgebraicSimplifierTest, CopyEqualsBitcast) { EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Copy(m::Parameter(0)))); - AlgebraicSimplifierOptions options2(bitcasting_callback()); + AlgebraicSimplifierOptions options2; options2.set_is_layout_sensitive(true); AlgebraicSimplifier simplifier2(options2); - ASSERT_TRUE(simplifier2.Run(m.get()).ValueOrDie()); + EXPECT_TRUE(simplifier2.Run(m.get()).ValueOrDie()); // Verify that the copy is replaced. EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Bitcast(m::Parameter(0)))); @@ -1744,7 +1871,7 @@ TEST_F(AlgebraicSimplifierTest, CopyWithDifferentLayout) { EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Copy(m::Parameter(0)))); - AlgebraicSimplifierOptions options(non_bitcasting_callback()); + AlgebraicSimplifierOptions options; options.set_is_layout_sensitive(true); AlgebraicSimplifier simplifier(options); EXPECT_FALSE(simplifier.Run(m.get()).ValueOrDie()); @@ -1774,7 +1901,7 @@ TEST_F(AlgebraicSimplifierTest, CopyWithSameLayout) { EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Copy(m::Parameter(0)))); - AlgebraicSimplifierOptions options(non_bitcasting_callback()); + AlgebraicSimplifierOptions options; options.set_is_layout_sensitive(true); AlgebraicSimplifier simplifier(options); ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); @@ -1804,7 +1931,8 @@ TEST_F(AlgebraicSimplifierTest, NoBitcastAdded) { EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Reshape(m::Parameter(0)))); - AlgebraicSimplifierOptions options(non_bitcasting_callback()); + AlgebraicSimplifierOptions options( + [](const Shape&, const Shape&) { return false; }); options.set_is_layout_sensitive(true); AlgebraicSimplifier simplifier(options); EXPECT_FALSE(simplifier.Run(m.get()).ValueOrDie()); @@ -1835,8 +1963,7 @@ TEST_F(AlgebraicSimplifierTest, ReshapeOfTransposeOfRngToRng) { auto computation = m->AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier( - (AlgebraicSimplifierOptions(bitcasting_callback()))); + AlgebraicSimplifier simplifier(AlgebraicSimplifierOptions{}); EXPECT_TRUE(simplifier.Run(m.get()).ValueOrDie()); // Verify that reshape(transpose(rng)) is replace by a single rng of the @@ -1887,7 +2014,7 @@ TEST_F(AlgebraicSimplifierTest, ReshapeReplacedWithBitcast) { m::Op().Is(dimensions_wrong_reshape), m::Op().Is(layout_wrong_reshape)))); - AlgebraicSimplifierOptions options(bitcasting_callback()); + AlgebraicSimplifierOptions options; options.set_is_layout_sensitive(true); AlgebraicSimplifier simplifier(options); simplifier.Run(m.get()).ValueOrDie(); @@ -1917,8 +2044,7 @@ TEST_F(AlgebraicSimplifierTest, FailureToSinkReshapeDoesntAffectChangedBit) { builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(F32, {4}), add)); - AlgebraicSimplifier simplifier( - (AlgebraicSimplifierOptions(bitcasting_callback()))); + AlgebraicSimplifier simplifier(AlgebraicSimplifierOptions{}); m->AddEntryComputation(builder.Build()); EXPECT_TRUE(simplifier.Run(m.get()).ValueOrDie()); } @@ -1942,8 +2068,7 @@ TEST_F(AlgebraicSimplifierTest, FailureToSinkBroadcastDoesntAffectChangedBit) { HloInstruction::CreateBroadcast(ShapeUtil::MakeShape(F32, {2, 2, 2}), add, /*broadcast_dimensions=*/{0, 1})); - AlgebraicSimplifier simplifier( - (AlgebraicSimplifierOptions(bitcasting_callback()))); + AlgebraicSimplifier simplifier(AlgebraicSimplifierOptions{}); m->AddEntryComputation(builder.Build()); EXPECT_TRUE(simplifier.Run(m.get()).ValueOrDie()); } @@ -1968,7 +2093,7 @@ TEST_F(AlgebraicSimplifierTest, TransposeEqualsBitcast1) { EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Transpose(m::Parameter(0)))); - AlgebraicSimplifierOptions options(bitcasting_callback()); + AlgebraicSimplifierOptions options; options.set_is_layout_sensitive(true); AlgebraicSimplifier simplifier(options); ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); @@ -1998,7 +2123,7 @@ TEST_F(AlgebraicSimplifierTest, TransposeEqualsBitcast2) { EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Transpose(m::Parameter(0)))); - AlgebraicSimplifierOptions options(bitcasting_callback()); + AlgebraicSimplifierOptions options; options.set_is_layout_sensitive(true); AlgebraicSimplifier simplifier(options); ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); @@ -2055,7 +2180,7 @@ TEST_F(AlgebraicSimplifierTest, CopiesMerged) { EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Copy(m::Copy(m::Parameter(0))))); - AlgebraicSimplifierOptions options(non_bitcasting_callback()); + AlgebraicSimplifierOptions options; options.set_is_layout_sensitive(true); AlgebraicSimplifier simplifier(options); ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); @@ -2103,9 +2228,8 @@ TEST_F(AlgebraicSimplifierTest, TransposeIsReshape) { ROOT reshaped_again = f32[10] reshape(f32[10,1,1] transposed) } )"; - TF_ASSERT_OK_AND_ASSIGN( - auto module, - HloRunner::CreateModuleFromString(hlo_string, GetDebugOptionsForTest())); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); HloPassFix simplifier(default_options_); EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); @@ -2624,161 +2748,22 @@ TEST_F(AlgebraicSimplifierTest, SliceOfReshapeUnchanged) { TEST_F(AlgebraicSimplifierTest, RemoveNoopSort) { auto builder = HloComputation::Builder(TestName()); + auto module = CreateNewVerifiedModule(); Shape keys_shape = ShapeUtil::MakeShape(F32, {1}); auto keys = builder.AddInstruction( HloInstruction::CreateParameter(0, keys_shape, "keys")); - builder.AddInstruction(HloInstruction::CreateSort(keys_shape, 0, keys)); - auto module = CreateNewVerifiedModule(); + TF_ASSERT_OK( + MakeSortHlo(keys_shape, {keys}, 0, &builder, module.get()).status()); HloComputation* computation = module->AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(default_options_); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), keys); } -TEST_F(AlgebraicSimplifierTest, ReplacePermutationSortWithScatter) { - const char* hlo_string = R"( - HloModule permutation_sort - - ENTRY sort_computation { - keys = f32[64,8732]{1,0} parameter(0) - values = s32[64,8732]{1,0} iota(), iota_dimension=1 - sort = (f32[64,8732]{1,0}, s32[64,8732]{1,0}) sort(keys, values), dimensions={1} - gte = s32[64,8732]{1,0} get-tuple-element(sort), index=1 - ROOT sort2 = (s32[64,8732]{1,0}, s32[64,8732]{1,0}) sort(gte, values), dimensions={1} - } - )"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); - - AlgebraicSimplifierOptions options(non_bitcasting_callback()); - options.set_enable_permutation_sort_replacement(true); - AlgebraicSimplifier simplifier(options); - EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - auto root = module->entry_computation()->root_instruction(); - EXPECT_THAT(root, - GmockMatch(m::Tuple( - m::Iota(), - m::Scatter(m::Iota(), m::Concatenate(m::Iota(), m::Reshape()), - m::Reshape())))); -} - -TEST_F(AlgebraicSimplifierTest, DontReplacePermutationSortIfNonIntegral) { - // Same as ReplacePermutationSortWithScatter except that the iota has F32 - // type. - const char* hlo_string = R"( - HloModule permutation_sort - - ENTRY sort_computation { - keys = f32[64,8732]{1,0} parameter(0) - values = f32[64,8732]{1,0} iota(), iota_dimension=1 - sort = (f32[64,8732]{1,0}, f32[64,8732]{1,0}) sort(keys, values), dimensions={1} - gte = f32[64,8732]{1,0} get-tuple-element(sort), index=1 - ROOT sort2 = (f32[64,8732]{1,0}, f32[64,8732]{1,0}) sort(gte, values), dimensions={1} - } - )"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); - - AlgebraicSimplifierOptions options(non_bitcasting_callback()); - options.set_enable_permutation_sort_replacement(true); - AlgebraicSimplifier simplifier(options); - EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); -} - -TEST_F(AlgebraicSimplifierTest, DontReplacePermutationSortWrongDimensions) { - // Same as ReplacePermutationSortWithScatter except that the sort dimensions - // don't match. - const char* hlo_string = R"( - HloModule permutation_sort - - ENTRY sort_computation { - keys = f32[64,8732]{1,0} parameter(0) - values = s32[64,8732]{1,0} iota(), iota_dimension=1 - sort = (f32[64,8732]{1,0}, s32[64,8732]{1,0}) sort(keys, values), dimensions={1} - gte = s32[64,8732]{1,0} get-tuple-element(sort), index=1 - ROOT sort2 = (s32[64,8732]{1,0}, s32[64,8732]{1,0}) sort(gte, values), dimensions={0} - } - )"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); - - AlgebraicSimplifierOptions options(non_bitcasting_callback()); - options.set_enable_permutation_sort_replacement(true); - AlgebraicSimplifier simplifier(options); - EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); -} - -TEST_F(AlgebraicSimplifierTest, RemoveUnusedSortOperandArrayResult) { - const char* hlo_string = R"( - HloModule permutation_sort - - ENTRY sort_computation { - keys = f32[64,8732]{1,0} parameter(0) - values = s32[64,8732]{1,0} parameter(1) - sort = (f32[64,8732]{1,0}, s32[64,8732]{1,0}) sort(keys, values), - dimensions={1} - ROOT gte = f32[64,8732]{1,0} get-tuple-element(sort), index=0 - })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); - - AlgebraicSimplifierOptions options(bitcasting_callback()); - AlgebraicSimplifier simplifier(options); - EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - auto root = module->entry_computation()->root_instruction(); - EXPECT_THAT(root, GmockMatch(m::Sort(m::Parameter(0)))); -} - -TEST_F(AlgebraicSimplifierTest, RemoveUnusedSortOperandTuple) { - const char* hlo_string = R"( - HloModule permutation_sort - - ENTRY sort_computation { - keys = f32[64,87] parameter(0) - values.0 = s32[64,87] parameter(1) - values.1 = u32[64,87] parameter(2) - sort = (f32[64,87], s32[64,87], u32[64,87]) sort( - keys, values.0, values.1), - dimensions={1} - gte.0 = f32[64,87] get-tuple-element(sort), index=0 - gte.1 = u32[64,87] get-tuple-element(sort), index=2 - ROOT tuple = (f32[64,87], u32[64,87]) tuple(gte.0, gte.1) - })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); - - AlgebraicSimplifierOptions options(bitcasting_callback()); - AlgebraicSimplifier simplifier(options); - EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - auto root = module->entry_computation()->root_instruction(); - EXPECT_THAT( - root, - GmockMatch(m::Tuple( - m::GetTupleElement(m::Sort(m::Parameter(0), m::Parameter(2)), 0), - m::GetTupleElement(m::Sort(m::Parameter(0), m::Parameter(2)), 1)))); -} - -TEST_F(AlgebraicSimplifierTest, DontRemoveUnusedSortKey) { - const char* hlo_string = R"( - HloModule permutation_sort - - ENTRY sort_computation { - keys = f32[64,8732]{1,0} parameter(0) - values = s32[64,8732]{1,0} parameter(1) - sort = (f32[64,8732]{1,0}, s32[64,8732]{1,0}) sort(keys, values), dimensions={1} - ROOT gte = s32[64,8732]{1,0} get-tuple-element(sort), index=1 - })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); - - AlgebraicSimplifierOptions options(bitcasting_callback()); - AlgebraicSimplifier simplifier(options); - EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); -} - TEST_F(AlgebraicSimplifierTest, ReplaceEffectiveScalarKeyValueSortWithTuple) { auto builder = HloComputation::Builder(TestName()); + auto module = CreateNewVerifiedModule(); Shape keys_shape = ShapeUtil::MakeShape(F32, {5, 0}); Shape values_shape = ShapeUtil::MakeShape(S32, {5, 0}); @@ -2788,10 +2773,10 @@ TEST_F(AlgebraicSimplifierTest, ReplaceEffectiveScalarKeyValueSortWithTuple) { HloInstruction::CreateParameter(1, values_shape, "values0")); auto values1 = builder.AddInstruction( HloInstruction::CreateParameter(2, values_shape, "values1")); - builder.AddInstruction(HloInstruction::CreateSort( - ShapeUtil::MakeTupleShape({keys_shape, values_shape, values_shape}), 0, - keys, {values0, values1})); - auto module = CreateNewVerifiedModule(); + TF_ASSERT_OK(MakeSortHlo(ShapeUtil::MakeTupleShape( + {keys_shape, values_shape, values_shape}), + {keys, values0, values1}, 0, &builder, module.get()) + .status()); HloComputation* computation = module->AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(default_options_); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); @@ -3013,7 +2998,7 @@ class ConvInputPaddingTest : public AlgebraicSimplifierTest, public ::testing::WithParamInterface {}; -INSTANTIATE_TEST_CASE_P( +INSTANTIATE_TEST_SUITE_P( ConvInputPaddingTestCases, ConvInputPaddingTest, ::testing::ValuesIn(std::vector{ // Merge this edge padding into the conv. @@ -3121,7 +3106,7 @@ class ConvFilterPaddingTest : public AlgebraicSimplifierTest, public ::testing::WithParamInterface {}; -INSTANTIATE_TEST_CASE_P( +INSTANTIATE_TEST_SUITE_P( ConvFilterPaddingTestCases, ConvFilterPaddingTest, ::testing::ValuesIn(std::vector{ // Can only merge interior padding on the filter's spatial dimensions; @@ -3360,7 +3345,7 @@ TEST_F(AlgebraicSimplifierTest, ConvertConvToMatmul) { auto module = CreateNewUnverifiedModule(); auto* computation = module->AddEntryComputation(b.Build()); - AlgebraicSimplifierOptions simplifier_options(bitcasting_callback()); + AlgebraicSimplifierOptions simplifier_options; simplifier_options.set_is_layout_sensitive(true); AlgebraicSimplifier simplifier(simplifier_options); if (!simplifier.Run(module.get()).ValueOrDie()) { @@ -3774,12 +3759,16 @@ TEST_F(AlgebraicSimplifierTest, TrivialDynamicSlice) { HloComputation::Builder builder(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {10, 100, 1000}); + std::vector params; + for (int i = 0; i < 3; ++i) { + params.push_back(builder.AddInstruction(HloInstruction::CreateParameter( + i + 1, ShapeUtil::MakeShape(U32, {}), "slice_indices"))); + } builder.AddInstruction(HloInstruction::CreateDynamicSlice( shape, builder.AddInstruction( HloInstruction::CreateParameter(0, shape, "slice_from")), - builder.AddInstruction(HloInstruction::CreateParameter( - 1, ShapeUtil::MakeShape(U32, {3}), "slice_indices")), + params, /*slice_sizes=*/{10, 100, 1000})); auto computation = m->AddEntryComputation(builder.Build()); @@ -3798,28 +3787,35 @@ TEST_F(AlgebraicSimplifierTest, TrivialDynamicUpdateSlice) { Shape full_shape = ShapeUtil::MakeShape(F32, {10, 100, 1000}); Shape slice_shape = ShapeUtil::MakeShape(F32, {10, 1, 1000}); + std::vector slice_indices, update_indices; + for (int i = 0; i < 3; ++i) { + slice_indices.push_back( + builder.AddInstruction(HloInstruction::CreateParameter( + i + 1, ShapeUtil::MakeShape(U32, {}), "slice_indices"))); + update_indices.push_back( + builder.AddInstruction(HloInstruction::CreateParameter( + i + 5, ShapeUtil::MakeShape(U32, {}), "update_indices"))); + } HloInstruction* slice = builder.AddInstruction(HloInstruction::CreateDynamicSlice( slice_shape, builder.AddInstruction( HloInstruction::CreateParameter(0, full_shape, "slice_from")), - builder.AddInstruction(HloInstruction::CreateParameter( - 1, ShapeUtil::MakeShape(U32, {3}), "slice_indices")), + slice_indices, /*slice_sizes=*/{10, 1, 1000})); builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( slice_shape, builder.AddInstruction( - HloInstruction::CreateParameter(2, slice_shape, "to_update")), - slice, - builder.AddInstruction(HloInstruction::CreateParameter( - 3, ShapeUtil::MakeShape(U32, {3}), "update_indices")))); + HloInstruction::CreateParameter(4, slice_shape, "to_update")), + slice, update_indices)); auto computation = m->AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(default_options_); ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), - GmockMatch(m::DynamicSlice(m::Parameter(), m::Parameter()))); + GmockMatch(m::DynamicSlice(m::Parameter(), m::Parameter(), + m::Parameter(), m::Parameter()))); } // Test that two consecutive broadcasts can be merged to one. @@ -3926,7 +3922,7 @@ TEST_F(AlgebraicSimplifierTest, SliceOfPadLow) { TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AlgebraicSimplifierOptions options(bitcasting_callback()); + AlgebraicSimplifierOptions options; AlgebraicSimplifier simplifier(options); EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); auto root = module->entry_computation()->root_instruction(); @@ -3947,7 +3943,7 @@ TEST_F(AlgebraicSimplifierTest, SliceOfPadHigh) { TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AlgebraicSimplifierOptions options(bitcasting_callback()); + AlgebraicSimplifierOptions options; AlgebraicSimplifier simplifier(options); EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); auto root = module->entry_computation()->root_instruction(); @@ -3968,7 +3964,7 @@ TEST_F(AlgebraicSimplifierTest, SliceOfPadMidNonScalar) { TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AlgebraicSimplifierOptions options(bitcasting_callback()); + AlgebraicSimplifierOptions options; AlgebraicSimplifier simplifier(options); EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); } @@ -3987,7 +3983,7 @@ TEST_F(AlgebraicSimplifierTest, SliceOfPadMidScalar) { TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AlgebraicSimplifierOptions options(bitcasting_callback()); + AlgebraicSimplifierOptions options; AlgebraicSimplifier simplifier(options); EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); auto root = module->entry_computation()->root_instruction(); @@ -4009,7 +4005,7 @@ TEST_F(AlgebraicSimplifierTest, SliceOfConcatScalarInput) { TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AlgebraicSimplifierOptions options(bitcasting_callback()); + AlgebraicSimplifierOptions options; AlgebraicSimplifier simplifier(options); EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); auto root = module->entry_computation()->root_instruction(); @@ -4031,7 +4027,7 @@ TEST_F(AlgebraicSimplifierTest, SliceOfConcatNonScalarInput) { TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AlgebraicSimplifierOptions options(bitcasting_callback()); + AlgebraicSimplifierOptions options; AlgebraicSimplifier simplifier(options); EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); auto root = module->entry_computation()->root_instruction(); @@ -4053,7 +4049,7 @@ TEST_F(AlgebraicSimplifierTest, NegateNegate) { TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AlgebraicSimplifierOptions options(bitcasting_callback()); + AlgebraicSimplifierOptions options; AlgebraicSimplifier simplifier(options); EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); auto root = module->entry_computation()->root_instruction(); @@ -4073,7 +4069,7 @@ TEST_F(AlgebraicSimplifierTest, NotNot) { TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AlgebraicSimplifierOptions options(bitcasting_callback()); + AlgebraicSimplifierOptions options; AlgebraicSimplifier simplifier(options); EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); auto root = module->entry_computation()->root_instruction(); @@ -4210,7 +4206,7 @@ PadReduceWindowEffectiveBroadcastCases() { return *cases; } -INSTANTIATE_TEST_CASE_P( +INSTANTIATE_TEST_SUITE_P( PadReduceWindowEffectiveBroadcastInstantiation, PadReduceWindowEffectiveBroadcastTest, ::testing::ValuesIn(PadReduceWindowEffectiveBroadcastCases())); @@ -4226,8 +4222,10 @@ TEST_P(BatchDotStrengthReductionTest, BatchDotStrengthReduction) { std::tie(m, k, n, element_type) = GetParam(); Shape dot_shape = ShapeUtil::MakeShape(element_type, {1, 3, 5, m, n}); - Shape lhs_shape = ShapeUtil::MakeShape(element_type, {1, 3, 5, m, k}); - Shape rhs_shape = ShapeUtil::MakeShape(element_type, {1, 3, 5, k, n}); + Shape lhs_shape = k > 0 ? ShapeUtil::MakeShape(element_type, {1, 3, 5, m, k}) + : ShapeUtil::MakeShape(element_type, {1, 3, 5, m}); + Shape rhs_shape = k > 0 ? ShapeUtil::MakeShape(element_type, {1, 3, 5, k, n}) + : ShapeUtil::MakeShape(element_type, {1, 3, 5, n}); HloComputation::Builder builder(TestName()); auto lhs = builder.AddInstruction( @@ -4241,14 +4239,16 @@ TEST_P(BatchDotStrengthReductionTest, BatchDotStrengthReduction) { dot_dnums.add_rhs_batch_dimensions(0); dot_dnums.add_rhs_batch_dimensions(1); dot_dnums.add_rhs_batch_dimensions(2); - dot_dnums.add_lhs_contracting_dimensions(4); - dot_dnums.add_rhs_contracting_dimensions(3); + if (k > 0) { + dot_dnums.add_lhs_contracting_dimensions(4); + dot_dnums.add_rhs_contracting_dimensions(3); + } builder.AddInstruction(HloInstruction::CreateDot( dot_shape, lhs, rhs, dot_dnums, DefaultPrecisionConfig(2))); auto computation = module->AddEntryComputation(builder.Build()); AlgebraicSimplifier simplifier(default_options_); TF_ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(module.get())); - const bool dot_should_be_transformed = m == 1 || k == 1 || n == 1; + const bool dot_should_be_transformed = m == 1 || k == 1 || n == 1 || k == -1; const bool computation_should_be_modified = dot_should_be_transformed; EXPECT_EQ(changed, computation_should_be_modified); bool has_no_dot = true; @@ -4261,9 +4261,9 @@ TEST_P(BatchDotStrengthReductionTest, BatchDotStrengthReduction) { EXPECT_EQ(has_no_dot, dot_should_be_transformed); } -INSTANTIATE_TEST_CASE_P( +INSTANTIATE_TEST_SUITE_P( BatchDotStrengthReductionTestInstantiation, BatchDotStrengthReductionTest, - ::testing::Combine(::testing::Values(1, 2), ::testing::Values(1, 2), + ::testing::Combine(::testing::Values(1, 2), ::testing::Values(-1, 1, 2), ::testing::Values(1, 2), ::testing::Values(F32, BF16))); class DotStrengthReductionTest @@ -4318,7 +4318,7 @@ TEST_P(DotStrengthReductionTest, DotStrengthReduction) { EXPECT_EQ(has_no_dot, dot_should_be_transformed); } -INSTANTIATE_TEST_CASE_P( +INSTANTIATE_TEST_SUITE_P( DotStrengthReductionTestInstantiation, DotStrengthReductionTest, ::testing::Combine(::testing::Values(1, 2), ::testing::Values(1, 2), ::testing::Values(1, 2), ::testing::Bool(), @@ -4480,9 +4480,10 @@ TEST_F(AlgebraicSimplifierTest, DynamicUpdateSliceZeroUpdate) { HloInstruction* const update = builder.AddInstruction( HloInstruction::CreateParameter(1, update_shape, "update")); HloInstruction* const start_indices = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({0}))); + HloInstruction::CreateConstant(LiteralUtil::CreateR0({}))); builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - dslice_shape, operand, update, start_indices)); + dslice_shape, operand, update, + std::initializer_list({start_indices}))); const HloComputation* const computation = m->AddEntryComputation(builder.Build()); @@ -4491,9 +4492,9 @@ TEST_F(AlgebraicSimplifierTest, DynamicUpdateSliceZeroUpdate) { EXPECT_THAT(computation->root_instruction(), operand); } -INSTANTIATE_TEST_CASE_P(DotOfConcatSimplificationTestInstantiation, - DotOfConcatSimplificationTest, - ::testing::ValuesIn(kDotOfConcatTestSpecs)); +INSTANTIATE_TEST_SUITE_P(DotOfConcatSimplificationTestInstantiation, + DotOfConcatSimplificationTest, + ::testing::ValuesIn(kDotOfConcatTestSpecs)); struct DotOfGatherTestSpec { int64 m; @@ -4535,14 +4536,17 @@ TEST_P(DotOfGatherSimplificationTest, ConstantRHS) { int32 start_row = (spec.lcd == 0) ? 0 : spec.s; int32 start_col = (spec.lcd == 0) ? spec.s : 0; - const auto start_indices = + std::vector start_indices = { builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({start_row, start_col}))); + LiteralUtil::CreateR0(start_row))), + builder.AddInstruction(HloInstruction::CreateConstant( + LiteralUtil::CreateR0(start_col)))}; int64 slice_row_size = (spec.lcd == 0) ? spec.k : 1; int64 slice_col_size = (spec.lcd == 0) ? 1 : spec.k; - Shape ds_shape = ShapeUtil::MakeShape(F32, {slice_row_size, slice_col_size}); + std::vector slice_sizes = {slice_row_size, slice_col_size}; + Shape ds_shape = ShapeUtil::MakeShape(F32, slice_sizes); auto* ds = builder.AddInstruction(HloInstruction::CreateDynamicSlice( - ds_shape, lhs, start_indices, {slice_row_size, slice_col_size})); + ds_shape, lhs, start_indices, slice_sizes)); int64 rhs_rows = (spec.rcd == 0) ? spec.k : spec.n; int64 rhs_cols = (spec.rcd == 0) ? spec.n : spec.k; @@ -4575,7 +4579,7 @@ TEST_P(DotOfGatherSimplificationTest, ConstantRHS) { } else { EXPECT_THAT(computation->root_instruction(), GmockMatch(m::DynamicSlice(m::Dot(m::Constant(), m::Constant()), - m::Concatenate()))); + m::Constant(), m::Constant()))); } } @@ -4613,14 +4617,17 @@ TEST_P(DotOfGatherSimplificationTest, ConstantLHS) { int32 start_row = (spec.rcd == 0) ? 0 : spec.s; int32 start_col = (spec.rcd == 0) ? spec.s : 0; - const auto start_indices = + std::vector start_indices = { + builder.AddInstruction(HloInstruction::CreateConstant( + LiteralUtil::CreateR0(start_row))), builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({start_row, start_col}))); + LiteralUtil::CreateR0(start_col)))}; int64 slice_row_size = (spec.rcd == 0) ? spec.k : 1; int64 slice_col_size = (spec.rcd == 0) ? 1 : spec.k; - Shape ds_shape = ShapeUtil::MakeShape(F32, {slice_row_size, slice_col_size}); + std::vector slice_sizes = {slice_row_size, slice_col_size}; + Shape ds_shape = ShapeUtil::MakeShape(F32, slice_sizes); auto* ds = builder.AddInstruction(HloInstruction::CreateDynamicSlice( - ds_shape, rhs, start_indices, {slice_row_size, slice_col_size})); + ds_shape, rhs, start_indices, slice_sizes)); DotDimensionNumbers dot_dnums; dot_dnums.add_lhs_contracting_dimensions(spec.lcd); @@ -4645,7 +4652,7 @@ TEST_P(DotOfGatherSimplificationTest, ConstantLHS) { } else { EXPECT_THAT(computation->root_instruction(), GmockMatch(m::DynamicSlice(m::Dot(m::Constant(), m::Constant()), - m::Concatenate()))); + m::Constant(), m::Constant()))); } } @@ -4693,9 +4700,124 @@ std::vector DotOfGatherPositiveNegativeTests() { return all; } -INSTANTIATE_TEST_CASE_P( +INSTANTIATE_TEST_SUITE_P( DotOfGatherSimplificationTestInstantiation, DotOfGatherSimplificationTest, ::testing::ValuesIn(DotOfGatherPositiveNegativeTests())); +TEST_F(AlgebraicSimplifierTest, TupleReduceReshape) { + const char* hlo_string = R"( +HloModule module + +reducer { + parameter.1 = f32[] parameter(0) + parameter.3 = f32[] parameter(2) + add.2 = f32[] add(parameter.1, parameter.3) + parameter.0 = f32[] parameter(1) + parameter.2 = f32[] parameter(3) + add.3 = f32[] add(parameter.0, parameter.2) + ROOT tuple.4 = (f32[], f32[]) tuple(add.2, add.3) +} + +ENTRY entry { + parameter.6 = (f32[], f32[]) parameter(0) + get-tuple-element.10 = f32[] get-tuple-element(parameter.6), index=0 + get-tuple-element.11 = f32[] get-tuple-element(parameter.6), index=1 + constant = f32[] constant(0) + ROOT reduce = (f32[], f32[]) reduce(get-tuple-element.10, get-tuple-element.11, constant, constant), dimensions={}, to_apply=reducer +} +)"; + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); + + AlgebraicSimplifierOptions options; + AlgebraicSimplifier simplifier(options); + EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + auto root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, GmockMatch(m::Tuple( + m::Reshape(m::GetTupleElement(m::Parameter(), 0)), + m::Reshape(m::GetTupleElement(m::Parameter(), 1))))); +} + +TEST_F(AlgebraicSimplifierTest, TupleReduceBroadcast) { + const char* hlo_string = R"( +HloModule module + +reducer { + parameter.1 = f32[] parameter(0) + parameter.3 = f32[] parameter(2) + mul.2 = f32[] add(parameter.1, parameter.3) + parameter.0 = f32[] parameter(1) + parameter.2 = f32[] parameter(3) + add.3 = f32[] add(parameter.0, parameter.2) + ROOT tuple.4 = (f32[], f32[]) tuple(mul.2, add.3) +} + +ENTRY entry { + parameter.6 = (f32[0, 10, 10], f32[0, 10, 10]) parameter(0) + get-tuple-element.10 = f32[0, 10, 10] get-tuple-element(parameter.6), index=0 + get-tuple-element.11 = f32[0, 10, 10] get-tuple-element(parameter.6), index=1 + constant.0 = f32[] constant(0) + constant.1 = f32[] constant(1) + ROOT reduce = (f32[10, 10], f32[10, 10]) reduce(get-tuple-element.10, get-tuple-element.11, constant.0, constant.1), dimensions={0}, to_apply=reducer +} +)"; + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); + + AlgebraicSimplifierOptions options; + AlgebraicSimplifier simplifier(options); + EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + auto root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, GmockMatch(m::Tuple(m::Broadcast(m::ConstantScalar(0)), + m::Broadcast(m::ConstantScalar(1))))); +} + +TEST_F(AlgebraicSimplifierTest, ZeroSizedReshapeWithoutLayout) { + auto builder = HloComputation::Builder(TestName()); + HloInstruction* param = + builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {1}), "param")); + HloInstruction* broadcast = + builder.AddInstruction(HloInstruction::CreateBroadcast( + ShapeUtil::MakeShape(F32, {0, 1}), param, {1})); + + // Create a reshape with zero sized result and without layout. + Shape reshaped_shape = ShapeUtil::MakeShape(F32, {0}); + reshaped_shape.clear_layout(); + builder.AddInstruction( + HloInstruction::CreateReshape(reshaped_shape, broadcast)); + + std::unique_ptr module = CreateNewVerifiedModule(); + module->AddEntryComputation(builder.Build()); + + AlgebraicSimplifierOptions options; + AlgebraicSimplifier simplifier(options); + EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, GmockMatch(m::Constant())); +} + +TEST_F(AlgebraicSimplifierTest, DividedByConstantInstructionWithoutLayout) { + Shape shape = ShapeUtil::MakeShape(F32, {}); + shape.clear_layout(); + auto builder = HloComputation::Builder(TestName()); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + + HloInstruction* const_value = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(20.0f))); + builder.AddInstruction(HloInstruction::CreateBinary(shape, HloOpcode::kDivide, + param, const_value)); + + std::unique_ptr module = CreateNewVerifiedModule(); + module->AddEntryComputation(builder.Build()); + + AlgebraicSimplifierOptions options; + AlgebraicSimplifier simplifier(options); + EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, GmockMatch(m::Multiply())); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/ar_crs_combiner.cc b/tensorflow/compiler/xla/service/ar_crs_combiner.cc index 47d2c7e35705698d49950c2fa042af1c6327d521..99373dc107ade342da590e9d875a98d789aad485 100644 --- a/tensorflow/compiler/xla/service/ar_crs_combiner.cc +++ b/tensorflow/compiler/xla/service/ar_crs_combiner.cc @@ -26,9 +26,9 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/pattern_matcher.h" +#include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" namespace xla { @@ -36,19 +36,34 @@ namespace { namespace m = match; -// Returns true iff the argument instruction is an AllReduce, followed by a -// certain sequence of instructions and then a CRS. It must be possible to move -// the AR past each instruction in the sequence. -bool MatchesArCrsPattern(HloInstruction* instruction) { +// Checks if the argument instruction is an AllReduce, followed by a certain +// sequence of instructions and then a CRS. It must be possible to move +// the AR past each instruction in the sequence. Returns the CRS, which is the +// last instruction in the sequence. +absl::optional MatchesArCrsPattern( + HloInstruction* instruction) { auto can_ar_move_past_instruction = [](HloInstruction* instruction) -> bool { if (instruction->user_count() != 1) { return false; } - auto opcode = instruction->opcode(); - return opcode == HloOpcode::kBitcast || opcode == HloOpcode::kTranspose || - opcode == HloOpcode::kReshape || opcode == HloOpcode::kConvert || - opcode == HloOpcode::kAdd || opcode == HloOpcode::kSubtract || - opcode == HloOpcode::kMultiply; + switch (instruction->opcode()) { + case HloOpcode::kBitcast: + case HloOpcode::kTranspose: + case HloOpcode::kReshape: + return true; + case HloOpcode::kConvert: + // Can be moved across if both input and output is either float or + // integer (e.g. S32<->U32 or F32<->BF16) + return ShapeUtil::ElementIsFloating(instruction->shape()) == + ShapeUtil::ElementIsFloating(instruction->operand(0)->shape()); + case HloOpcode::kAdd: + case HloOpcode::kSubtract: + case HloOpcode::kMultiply: + // Only supported for floating point operands. + return ShapeUtil::ElementIsFloating(instruction->shape()); + default: + return false; + } }; auto computation_is_addition = [](HloComputation* c) { @@ -59,17 +74,22 @@ bool MatchesArCrsPattern(HloInstruction* instruction) { if (!instruction->IsCrossModuleAllReduce() || !computation_is_addition(instruction->called_computations()[0]) || instruction->user_count() != 1) { - return false; + return absl::nullopt; } auto next = instruction->users()[0]; while (!next->IsCrossReplicaAllReduce()) { if (can_ar_move_past_instruction(next)) { next = next->users()[0]; } else { - return false; + return absl::nullopt; } } - return computation_is_addition(next->called_computations()[0]); + if (!Cast(next)->IsNoop() && + computation_is_addition(next->called_computations()[0])) { + return absl::optional(next); + } else { + return absl::nullopt; + } } } // namespace @@ -85,7 +105,7 @@ absl::optional ArCrsCombiner::WhileFromBodyParameter( return caller_instruction; } } - return absl::optional(); + return absl::nullopt; } std::vector ArCrsCombiner::GetAllTuples( @@ -176,6 +196,15 @@ bool ArCrsCombiner::InstructionsComputeSameValue( if (opcode1 != i2->opcode() || operands1.size() != i2->operands().size()) { return false; } + auto eq_computations = [](const HloComputation* a, const HloComputation* b) { + return *a == *b; + }; + if (i1->IsCrossModuleAllReduce()) { + return i1->Identical(*i2, + /*eq_operands=*/std::equal_to(), + eq_computations, + /*layout_sensitive=*/false); + } visited_pairs->emplace(min_uid, max_uid); for (int i = 0; i < operands1.size(); ++i) { auto operand1 = operands1[i]; @@ -201,9 +230,6 @@ bool ArCrsCombiner::InstructionsComputeSameValue( // InstructionsComputeSameValue earlier. auto eq_instructions = [](const HloInstruction* i1, const HloInstruction* i2) -> bool { return true; }; - auto eq_computations = [](const HloComputation* a, const HloComputation* b) { - return *a == *b; - }; return i1->Identical(*i2, eq_instructions, eq_computations, /*layout_sensitive=*/false); } @@ -211,8 +237,14 @@ bool ArCrsCombiner::InstructionsComputeSameValue( void ArCrsCombiner::GroupAllReducesById(HloModule* module) { for (HloComputation* computation : module->MakeNonfusionComputations()) { for (HloInstruction* instruction : computation->instructions()) { - if (MatchesArCrsPattern(instruction)) { - all_reduce_map_[*(instruction->all_reduce_id())].push_back(instruction); + auto maybe_crs = MatchesArCrsPattern(instruction); + if (maybe_crs) { + auto crs = *maybe_crs; + int64 ar_id = *(instruction->all_reduce_id()); + if (crs_reserved_map_.find(crs) == crs_reserved_map_.end()) { + all_reduce_map_[ar_id].push_back(instruction); + crs_reserved_map_[crs] = ar_id; + } } } } @@ -229,14 +261,17 @@ void ArCrsCombiner::KeepProvablyEqualInstructionGroups() { auto next_0 = instr_0->users()[0]; auto next_i = instr_i->users()[0]; absl::flat_hash_map visited_pairs; - do { + while (true) { if (!InstructionsComputeSameValue(next_0, next_i, &visited_pairs)) { all_reduce_map_.erase(all_reduce_id); break; } + if (next_0->IsCrossReplicaAllReduce()) { + break; + } next_0 = next_0->users()[0]; next_i = next_i->users()[0]; - } while (!next_0->IsCrossReplicaAllReduce()); + } } } } diff --git a/tensorflow/compiler/xla/service/ar_crs_combiner.h b/tensorflow/compiler/xla/service/ar_crs_combiner.h index 6f54b97615b270bc6b180dd47d9aff6473752b47..e61ef5d4f9072979a6c356a9456c91e19405b01e 100644 --- a/tensorflow/compiler/xla/service/ar_crs_combiner.h +++ b/tensorflow/compiler/xla/service/ar_crs_combiner.h @@ -83,6 +83,11 @@ class ArCrsCombiner : public HloModulePass { // Map from all-reduce ids to the all reduce instructions. absl::flat_hash_map> all_reduce_map_; + // Map from a CRS instruction to the all-reduce ID of the AR paired with the + // CRS. Sometimes, several ARs in the code could be paired with the same CRS. + // We use this map to pick a single AR/CRS path to rewrite. + absl::flat_hash_map crs_reserved_map_; + std::unique_ptr call_graph_; }; diff --git a/tensorflow/compiler/xla/service/ar_crs_combiner_test.cc b/tensorflow/compiler/xla/service/ar_crs_combiner_test.cc index caa57296f465698eb70d7cb8327d4678f394b323..5152f0dc884a153f9b0ade06acd479832d87ff25 100644 --- a/tensorflow/compiler/xla/service/ar_crs_combiner_test.cc +++ b/tensorflow/compiler/xla/service/ar_crs_combiner_test.cc @@ -360,6 +360,7 @@ HloModule foobar ENTRY %entrycomp (p: bf16[]) -> (f32[], f32[]) { %p = bf16[] parameter(0) + %constant.bf16 = bf16[] constant(1) %all-reduce.ar.1 = bf16[] all-reduce(%p), @@ -377,7 +378,7 @@ ENTRY %entrycomp (p: bf16[]) -> (f32[], f32[]) { sharding={maximal device=0} %all-reduce.ar.2 = bf16[] - all-reduce(%p), + all-reduce(%constant.bf16), replica_groups={{0},{1}}, all_reduce_id=1, to_apply=%sum.bf16, @@ -407,7 +408,7 @@ ENTRY %entrycomp (p: bf16[]) -> (f32[], f32[]) { EXPECT_TRUE(changed); EXPECT_THAT(module->entry_computation()->root_instruction(), op::Tuple(op::AllReduce(op::Convert(op::Parameter())), - op::AllReduce(op::Convert(op::Parameter())))); + op::AllReduce(op::Convert(op::Constant())))); auto crs_after = module->entry_computation()->root_instruction()->operands()[0]; auto replica_groups_after = crs_after->replica_groups(); @@ -705,5 +706,470 @@ ENTRY %entrycomp (p: f32[]) -> (f32[], f32[]) { EXPECT_FALSE(changed); } +TEST_F(ArCrsCombinerTest, ArThenCrsDontCrash) { + const char* module_str = R"( +HloModule foobar + +%sum.1 (a: f32[], b: f32[]) -> f32[] { + %a = f32[] parameter(0) + %b = f32[] parameter(1) + ROOT %add = f32[] add(%a, %b) +} + +ENTRY %entrycomp (p: f32[]) -> (f32[], f32[]) { + %p = f32[] parameter(0) + %constant.f32 = f32[] constant(123) + + %all-reduce.ar.1 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum.1, + sharding={maximal device=0} + %all-reduce.1 = f32[] + all-reduce(%all-reduce.ar.1), + replica_groups={{0,1}}, + to_apply=%sum.1, + sharding={maximal device=0} + %multiply.1 = f32[] + multiply(%all-reduce.1, %constant.f32), + sharding={maximal device=0} + + %all-reduce.ar.2 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum.1, + sharding={maximal device=1} + %all-reduce.2 = f32[] + all-reduce(%all-reduce.ar.2), + replica_groups={{0,1}}, + to_apply=%sum.1, + sharding={maximal device=1} + %multiply.2 = f32[] + multiply(%all-reduce.2, %constant.f32), + sharding={maximal device=1} + + ROOT %tuple = (f32[], f32[]) + tuple(%all-reduce.1, %all-reduce.2), + sharding={{maximal device=0}, {maximal device=1}} +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto crs_before = + module->entry_computation()->root_instruction()->operands()[0]; + auto replica_groups_before = crs_before->replica_groups(); + ArCrsCombiner combiner(2); + auto changed = combiner.Run(module.get()).ValueOrDie(); + EXPECT_TRUE(changed); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Tuple(op::AllReduce(op::Parameter()), + op::AllReduce(op::Parameter()))); + auto crs_after = + module->entry_computation()->root_instruction()->operands()[0]; + auto replica_groups_after = crs_after->replica_groups(); + CompareReplicaGroups(replica_groups_before, replica_groups_after); +} + +TEST_F(ArCrsCombinerTest, RewriteMultipleAdds) { + const char* module_str = R"( +HloModule foobar + +%sum (x: f32[], y: f32[]) -> f32[] { + %x = f32[] parameter(0) + %y = f32[] parameter(1) + ROOT %add = f32[] add(%x, %y) +} + +ENTRY %entrycomp (p: f32[]) -> (f32[], f32[]) { + %p = f32[] parameter(0) + %constant.1 = f32[] constant(1) + %constant.2 = f32[] constant(2) + + %all-reduce.ar.1 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum, + sharding={maximal device=0} + %add.11 = f32[] + add(%constant.1, %all-reduce.ar.1), + sharding={maximal device=0} + %add.12 = f32[] + add(%constant.2, %add.11), + sharding={maximal device=0} + %all-reduce.1 = f32[] + all-reduce(%add.12), + replica_groups={{0,1}}, + to_apply=%sum, + sharding={maximal device=0} + + %all-reduce.ar.2 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum, + sharding={maximal device=0} + %add.21 = f32[] + add(%constant.1, %all-reduce.ar.2), + sharding={maximal device=0} + %add.22 = f32[] + add(%constant.2, %add.21), + sharding={maximal device=0} + %all-reduce.2 = f32[] + all-reduce(%add.22), + replica_groups={{0,1}}, + to_apply=%sum, + sharding={maximal device=0} + + ROOT %tuple = (f32[], f32[]) + tuple(%all-reduce.1, %all-reduce.2), + sharding={{maximal device=0}, {maximal device=1}} +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto crs_before = + module->entry_computation()->root_instruction()->operands()[0]; + auto replica_groups_before = crs_before->replica_groups(); + ArCrsCombiner combiner(2); + auto changed = combiner.Run(module.get()).ValueOrDie(); + EXPECT_TRUE(changed); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Tuple(op::AllReduce(op::Add( + op::Divide(op::Constant(), op::Constant()), + op::Add(op::Divide(op::Constant(), op::Constant()), + op::Parameter()))), + op::AllReduce(op::Add( + op::Divide(op::Constant(), op::Constant()), + op::Add(op::Divide(op::Constant(), op::Constant()), + op::Parameter()))))); + auto crs_after = + module->entry_computation()->root_instruction()->operands()[0]; + auto replica_groups_after = crs_after->replica_groups(); + CompareReplicaGroups(replica_groups_before, replica_groups_after); +} + +TEST_F(ArCrsCombinerTest, RewriteArSubtractCrs) { + const char* module_str = R"( +HloModule foobar + +%sum.f32 (x: f32[], y: f32[]) -> f32[] { + %x = f32[] parameter(0) + %y = f32[] parameter(1) + ROOT %add = f32[] add(%x, %y) +} + +ENTRY %entrycomp (p: f32[]) -> (f32[], f32[]) { + %p = f32[] parameter(0) + %constant.f32 = f32[] constant(123) + + %all-reduce.ar.1 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum.f32, + sharding={maximal device=0} + %sub.1 = f32[] + subtract(%constant.f32, %all-reduce.ar.1), + sharding={maximal device=0} + %all-reduce.1 = f32[] + all-reduce(%sub.1), + replica_groups={{0,1}}, + to_apply=%sum.f32, + sharding={maximal device=0} + + %all-reduce.ar.2 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum.f32, + sharding={maximal device=1} + %sub.2 = f32[] + subtract(%constant.f32, %all-reduce.ar.2), + sharding={maximal device=1} + %all-reduce.2 = f32[] + all-reduce(%sub.2), + replica_groups={{0,1}}, + to_apply=%sum.f32, + sharding={maximal device=1} + + ROOT %tuple = (f32[], f32[]) + tuple(%all-reduce.1, %all-reduce.2), + sharding={{maximal device=0}, {maximal device=1}} +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto crs_before = + module->entry_computation()->root_instruction()->operands()[0]; + auto replica_groups_before = crs_before->replica_groups(); + ArCrsCombiner combiner(2); + auto changed = combiner.Run(module.get()).ValueOrDie(); + EXPECT_TRUE(changed); + EXPECT_THAT( + module->entry_computation()->root_instruction(), + op::Tuple( + op::AllReduce(op::Subtract(op::Divide(op::Constant(), op::Constant()), + op::Parameter())), + op::AllReduce(op::Subtract(op::Divide(op::Constant(), op::Constant()), + op::Parameter())))); + auto crs_after = + module->entry_computation()->root_instruction()->operands()[0]; + auto replica_groups_after = crs_after->replica_groups(); + CompareReplicaGroups(replica_groups_before, replica_groups_after); +} + +TEST_F(ArCrsCombinerTest, RewriteMultipleARsLeft) { + const char* module_str = R"( +HloModule foobar + +%sum (x: f32[], y: f32[]) -> f32[] { + %x = f32[] parameter(0) + %y = f32[] parameter(1) + ROOT %add = f32[] add(%x, %y) +} + +ENTRY %entrycomp (p: f32[]) -> (f32[], f32[]) { + %p = f32[] parameter(0) + %const1 = f32[] constant(1) + %const2 = f32[] constant(2) + + %ar11 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum, + sharding={maximal device=0} + %add11 = f32[] + add(%ar11, %const1), + sharding={maximal device=0} + %ar12 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=2, + to_apply=%sum, + sharding={maximal device=0} + %add12 = f32[] + add(%add11, %ar12), + sharding={maximal device=0} + %crs1 = f32[] + all-reduce(%add12), + replica_groups={{0,1}}, + to_apply=%sum, + sharding={maximal device=0} + + %ar21 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum, + sharding={maximal device=1} + %add21 = f32[] + add(%ar21, %const1), + sharding={maximal device=1} + %ar22 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=2, + to_apply=%sum, + sharding={maximal device=1} + %add22 = f32[] + add(%add21, %ar22), + sharding={maximal device=1} + %crs2 = f32[] + all-reduce(%add22), + replica_groups={{0,1}}, + to_apply=%sum, + sharding={maximal device=1} + + ROOT %tuple = (f32[], f32[]) + tuple(%crs1, %crs2), + sharding={{maximal device=0}, {maximal device=1}} +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto crs_before = + module->entry_computation()->root_instruction()->operands()[0]; + auto replica_groups_before = crs_before->replica_groups(); + ArCrsCombiner combiner(2); + auto changed = combiner.Run(module.get()).ValueOrDie(); + EXPECT_TRUE(changed); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Tuple(op::AllReduce(op::Add( + op::Add(op::Parameter(), + op::Divide(op::Constant(), op::Constant())), + op::Divide(op::AllReduce(), op::Constant()))), + op::AllReduce(op::Add( + op::Add(op::Parameter(), + op::Divide(op::Constant(), op::Constant())), + op::Divide(op::AllReduce(), op::Constant()))))); + auto crs_after = + module->entry_computation()->root_instruction()->operands()[0]; + auto replica_groups_after = crs_after->replica_groups(); + CompareReplicaGroups(replica_groups_before, replica_groups_after); +} + +TEST_F(ArCrsCombinerTest, RewriteMultipleARsRight) { + const char* module_str = R"( +HloModule foobar + +%sum (x: f32[], y: f32[]) -> f32[] { + %x = f32[] parameter(0) + %y = f32[] parameter(1) + ROOT %add = f32[] add(%x, %y) +} + +ENTRY %entrycomp (p: f32[]) -> (f32[], f32[]) { + %p = f32[] parameter(0) + %const1 = f32[] constant(1) + %const2 = f32[] constant(2) + + %ar11 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum, + sharding={maximal device=0} + %ar12 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=2, + to_apply=%sum, + sharding={maximal device=0} + %add11 = f32[] + add(%ar12, %const1), + sharding={maximal device=0} + %add12 = f32[] + add(%ar11, %add11), + sharding={maximal device=0} + %crs1 = f32[] + all-reduce(%add12), + replica_groups={{0,1}}, + to_apply=%sum, + sharding={maximal device=0} + + %ar21 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum, + sharding={maximal device=1} + %ar22 = f32[] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=2, + to_apply=%sum, + sharding={maximal device=1} + %add21 = f32[] + add(%ar22, %const1), + sharding={maximal device=1} + %add22 = f32[] + add(%ar21, %add21), + sharding={maximal device=1} + %crs2 = f32[] + all-reduce(%add22), + replica_groups={{0,1}}, + to_apply=%sum, + sharding={maximal device=1} + + ROOT %tuple = (f32[], f32[]) + tuple(%crs1, %crs2), + sharding={{maximal device=0}, {maximal device=1}} +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto crs_before = + module->entry_computation()->root_instruction()->operands()[0]; + auto replica_groups_before = crs_before->replica_groups(); + ArCrsCombiner combiner(2); + auto changed = combiner.Run(module.get()).ValueOrDie(); + EXPECT_TRUE(changed); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Tuple(op::AllReduce(op::Add( + op::Parameter(), + op::Divide(op::Add(op::AllReduce(), op::Constant()), + op::Constant()))), + op::AllReduce(op::Add( + op::Parameter(), + op::Divide(op::Add(op::AllReduce(), op::Constant()), + op::Constant()))))); + auto crs_after = + module->entry_computation()->root_instruction()->operands()[0]; + auto replica_groups_after = crs_after->replica_groups(); + CompareReplicaGroups(replica_groups_before, replica_groups_after); +} + +TEST_F(ArCrsCombinerTest, OneReplicaDontRewrite) { + const char* module_str = R"( +HloModule foobar + +%sum.bf16 (a: bf16[], b: bf16[]) -> bf16[] { + %a = bf16[] parameter(0) + %b = bf16[] parameter(1) + ROOT %add = bf16[] add(%a, %b) +} + +%sum.f32 (x: f32[], y: f32[]) -> f32[] { + %x = f32[] parameter(0) + %y = f32[] parameter(1) + ROOT %add = f32[] add(%x, %y) +} + +ENTRY %entrycomp (p: bf16[]) -> (f32[], f32[]) { + %p = bf16[] parameter(0) + %constant.bf16 = bf16[] constant(1) + + %all-reduce.ar.1 = bf16[] + all-reduce(%p), + replica_groups={{0}}, + all_reduce_id=1, + to_apply=%sum.bf16, + sharding={maximal device=0} + %convert.1 = f32[] + convert(%all-reduce.ar.1), + sharding={maximal device=0} + %all-reduce.1 = f32[] + all-reduce(%convert.1), + replica_groups={{0}}, + to_apply=%sum.f32, + sharding={maximal device=0} + + %all-reduce.ar.2 = bf16[] + all-reduce(%constant.bf16), + replica_groups={{0}}, + all_reduce_id=1, + to_apply=%sum.bf16, + sharding={maximal device=1} + %convert.2 = f32[] + convert(%all-reduce.ar.2), + sharding={maximal device=1} + %all-reduce.2 = f32[] + all-reduce(%convert.2), + replica_groups={{0}}, + to_apply=%sum.f32, + sharding={maximal device=1} + + ROOT %tuple = (f32[], f32[]) + tuple(%all-reduce.1, %all-reduce.2), + sharding={{maximal device=0}, {maximal device=1}} +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + ArCrsCombiner combiner(2); + auto changed = combiner.Run(module.get()).ValueOrDie(); + EXPECT_FALSE(changed); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/backend.cc b/tensorflow/compiler/xla/service/backend.cc index 2cf24a9dd5fa18abe9dde4eb49b03c6586bfef03..215e8ced4bb3f98a26ac4eb9912a7fd4d917852f 100644 --- a/tensorflow/compiler/xla/service/backend.cc +++ b/tensorflow/compiler/xla/service/backend.cc @@ -115,12 +115,10 @@ StatusOr Backend::BorrowStream(int device_ordinal) { StatusOr Backend::BorrowStream(se::StreamExecutor* executor) { tensorflow::mutex_lock l(mu_); - if (0 == stream_pools_.count(executor)) { - stream_pools_.emplace(std::piecewise_construct, - std::forward_as_tuple(executor), - std::forward_as_tuple()); + if (!stream_pools_.contains(executor)) { + stream_pools_.emplace(executor, absl::make_unique()); } - return stream_pools_.at(executor).BorrowStream(executor); + return stream_pools_.at(executor)->BorrowStream(executor); } Backend::Backend(se::Platform* platform, Compiler* compiler, diff --git a/tensorflow/compiler/xla/service/backend.h b/tensorflow/compiler/xla/service/backend.h index 7ca993fb2656037951d98d9c4459a3c3e4c64c61..c35f033dc0180409ae3888c2050021da83f5c72a 100644 --- a/tensorflow/compiler/xla/service/backend.h +++ b/tensorflow/compiler/xla/service/backend.h @@ -22,6 +22,7 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_map.h" #include "absl/strings/str_cat.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/service/compiler.h" @@ -175,7 +176,8 @@ class Backend { tensorflow::mutex mu_; // Mapping from stream executor to stream pools, used by `BorrowStream` above. - std::map stream_pools_ GUARDED_BY(mu_); + absl::flat_hash_map> + stream_pools_ GUARDED_BY(mu_); // The default memory allocator to use. std::unique_ptr memory_allocator_; diff --git a/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc b/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc index 551ac4be73a7630d213a53ca3606aa7f890cd794..2591ff602c8467afef9a1cbacd9aff2e63a8457e 100644 --- a/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc +++ b/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/bfloat16_normalization.h" #include "tensorflow/compiler/xla/service/bfloat16_support.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_creation_utils.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" @@ -282,8 +283,10 @@ TEST_F(BFloat16NormalizationTest, ResolveMixedPrecisionTupleSort) { HloInstruction* value = builder.AddInstruction( HloInstruction::CreateParameter(1, s32_shape, "value")); - HloInstruction* sort = builder.AddInstruction(HloInstruction::CreateSort( - ShapeUtil::MakeTupleShape({bf16_shape, s32_shape}), 0, key, {value})); + TF_ASSERT_OK_AND_ASSIGN( + auto* sort, + MakeSortHlo(ShapeUtil::MakeTupleShape({bf16_shape, s32_shape}), + {key, value}, 0, &builder, module.get())); HloInstruction* gte = builder.AddInstruction( HloInstruction::CreateGetTupleElement(bf16_shape, sort, 0)); @@ -308,8 +311,10 @@ TEST_F(BFloat16NormalizationTest, ResolveMixedPrecisionTupleSortRoot) { HloInstruction* value = builder.AddInstruction( HloInstruction::CreateParameter(1, bf16_shape, "value")); - HloInstruction* sort = builder.AddInstruction(HloInstruction::CreateSort( - ShapeUtil::MakeTupleShape({bf16_shape, bf16_shape}), 0, key, {value})); + TF_ASSERT_OK_AND_ASSIGN( + auto* sort, + MakeSortHlo(ShapeUtil::MakeTupleShape({bf16_shape, f32_shape}), + {key, value}, 0, &builder, module.get())); auto computation = module->AddEntryComputation(builder.Build()); diff --git a/tensorflow/compiler/xla/service/buffer_assignment.cc b/tensorflow/compiler/xla/service/buffer_assignment.cc index 6f4c1104f363146cebea83873b936ea3f8292801..cbebbdc8a2d7d0b65f12accbe424bea383ff5355 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment.cc +++ b/tensorflow/compiler/xla/service/buffer_assignment.cc @@ -86,10 +86,9 @@ std::vector ColorInterferenceGraph( // first, but it would be good to investigate other ordering heuristics too. std::vector nodes(node_count); std::iota(nodes.begin(), nodes.end(), 0); - std::sort(nodes.begin(), nodes.end(), - [&interference_map](const int64 i, const int64 j) { - return interference_map[i].size() > interference_map[j].size(); - }); + absl::c_sort(nodes, [&interference_map](const int64 i, const int64 j) { + return interference_map[i].size() > interference_map[j].size(); + }); const int64 kColorUnassigned = -1; std::vector assigned_colors(node_count, kColorUnassigned); @@ -138,8 +137,8 @@ Status GatherComputationsByAllocationType( worklist.pop_front(); const HloComputation* computation = worklist_front.first; bool is_thread_local = worklist_front.second; - bool in_thread_local_set = thread_local_set.count(computation) > 0; - bool in_global_set = global_set.count(computation) > 0; + bool in_thread_local_set = thread_local_set.contains(computation); + bool in_global_set = global_set.contains(computation); // If the computation has already been added to the respective set, then // nothing to do. @@ -192,6 +191,7 @@ Status GatherComputationsByAllocationType( case HloOpcode::kReduceWindow: case HloOpcode::kScatter: case HloOpcode::kSelectAndScatter: + case HloOpcode::kSort: case HloOpcode::kFusion: // Map/reduce etc computations are always thread-local. worklist.push_back(std::make_pair(subcomputation, @@ -207,9 +207,9 @@ Status GatherComputationsByAllocationType( // Add the computations to the vectors in post order. for (auto* computation : module->MakeComputationPostOrder()) { - if (thread_local_set.count(computation) > 0) { + if (thread_local_set.contains(computation)) { thread_local_computations->push_back(computation); - } else if (global_set.count(computation) > 0) { + } else if (global_set.contains(computation)) { global_computations->push_back(computation); } // If the computation is not reachable from the entry computation, then it @@ -219,13 +219,6 @@ Status GatherComputationsByAllocationType( return Status::OK(); } -size_t BufferAllocation::Slice::Hasher::operator()(Slice s) const { - uint64 h = std::hash()(s.index()); - h = tensorflow::Hash64Combine(h, std::hash()(s.offset())); - h = tensorflow::Hash64Combine(h, std::hash()(s.size())); - return h; -} - string BufferAllocation::Slice::ToString() const { return absl::StrCat("{index:", index(), ", offset:", offset_, ", size:", size_, "}"); @@ -240,7 +233,7 @@ BufferAllocation::Slice BufferAllocation::GetSlice( void BufferAllocation::AddAssignment(const LogicalBuffer& buffer, int64 offset, int64 size) { VLOG(4) << "Trying to add " << buffer << " to allocation #" << index(); - CHECK(assigned_buffers_.count(&buffer) == 0) + CHECK(!assigned_buffers_.contains(&buffer)) << "LogicalBuffer " << buffer << " already assigned to allocation " << index_; CHECK_LE(offset, size_) << "LogicalBuffer " << buffer @@ -279,11 +272,12 @@ BufferAllocationProto BufferAllocation::ToProto() const { proto_assigned->set_offset(buffer_offset_size.second.offset); proto_assigned->set_size(buffer_offset_size.second.size); } - std::sort(proto.mutable_assigned()->begin(), proto.mutable_assigned()->end(), - [](const BufferAllocationProto::Assigned& assign1, - const BufferAllocationProto::Assigned& assign2) { - return assign1.logical_buffer_id() < assign2.logical_buffer_id(); - }); + absl::c_sort(*proto.mutable_assigned(), + [](const BufferAllocationProto::Assigned& assign1, + const BufferAllocationProto::Assigned& assign2) { + return assign1.logical_buffer_id() < + assign2.logical_buffer_id(); + }); return proto; } @@ -315,10 +309,10 @@ string BufferAllocation::ToString() const { for (const auto& buffer_offset_size : assigned_buffers_) { sorted_buffers.push_back(buffer_offset_size.first); } - std::sort(sorted_buffers.begin(), sorted_buffers.end(), - [](const LogicalBuffer* a, const LogicalBuffer* b) { - return a->id() < b->id(); - }); + absl::c_sort(sorted_buffers, + [](const LogicalBuffer* a, const LogicalBuffer* b) { + return a->id() < b->id(); + }); for (const LogicalBuffer* buffer : sorted_buffers) { const OffsetSize& offset_size = FindOrDie(assigned_buffers_, buffer); StrAppend(&output, absl::StrFormat( @@ -346,7 +340,7 @@ const PointsToSet& BufferAssignment::GetPointsToSet( bool BufferAssignment::HasAllocation(const LogicalBuffer& buffer) const { TF_CHECK_OK(points_to_analysis().VerifyBuffer(buffer)); - return allocation_index_for_buffer_.count(&buffer) > 0; + return allocation_index_for_buffer_.contains(&buffer); } const BufferAllocation& BufferAssignment::GetAssignedAllocation( @@ -401,7 +395,7 @@ bool BufferAssignment::HasAllocationAt(const HloInstruction* instruction, const ShapeIndex& index) const { for (const LogicalBuffer* buffer : GetPointsToSet(instruction).element(index)) { - if (allocation_index_for_buffer_.count(buffer) > 0) { + if (allocation_index_for_buffer_.contains(buffer)) { return true; } } @@ -459,8 +453,7 @@ bool BufferAssignment::SharesSliceAtIndex( bool BufferAssignment::HaveDisjointSlices(const HloInstruction* hlo_a, const HloInstruction* hlo_b) const { - using SliceSet = - flat_hash_set; + using SliceSet = flat_hash_set; // Gets the slices all of instr's subshapes. If any subshape doesn't have an // assigned slice, returns the empty set. auto collect_slices = [&](const HloInstruction* instr) -> SliceSet { @@ -487,10 +480,9 @@ bool BufferAssignment::HaveDisjointSlices(const HloInstruction* hlo_a, // didn't return the empty set) for both HLOs, and the two resulting sets of // slices are disjoint. return !slices_a.empty() && !slices_b.empty() && - std::none_of(slices_a.begin(), slices_a.end(), - [&](const BufferAllocation::Slice& slice) { - return slices_b.count(slice) > 0; - }); + absl::c_none_of(slices_a, [&](const BufferAllocation::Slice& slice) { + return slices_b.contains(slice); + }); } StatusOr @@ -519,7 +511,7 @@ BufferAllocation* BufferAssignment::NewAllocation(const LogicalBuffer& buffer, void BufferAssignment::AddAssignment(BufferAllocation* allocation, const LogicalBuffer& buffer, int64 offset, int64 size) { - CHECK_EQ(0, allocation_index_for_buffer_.count(&buffer)) + CHECK(!allocation_index_for_buffer_.contains(&buffer)) << "LogicalBuffer " << buffer << " already has an allocation."; CHECK(allocation->is_reusable() || allocation->assigned_buffers().empty()) << "Non-reusable allocation already assigned a buffer: " @@ -761,7 +753,8 @@ namespace { bool MayInterfereAcrossSubcomputations(BufferAssignment* assignment, const LogicalBuffer& a_buffer, const LogicalBuffer& b_buffer) { - auto call_graph = assignment->liveness().hlo_ordering().call_graph(); + const CallGraph& call_graph = + assignment->liveness().hlo_ordering().call_graph(); const HloInstruction* a_ancestor; const HloInstruction* b_ancestor; std::tie(a_ancestor, b_ancestor) = @@ -960,35 +953,35 @@ Status BufferAssigner::AssignBuffersForComputation( // operands (assuming operands are the same/larger size) enabling the // important reuse case where an elementwise instruction reuses one of its // operand's buffer. This improves locality. - std::sort(sorted_buffers.begin(), sorted_buffers.end(), - [has_sequential_order, &liveness, &post_order_position, assignment]( - const LogicalBuffer* a, const LogicalBuffer* b) { - // Primary sort is by decreasing buffer size. - const int64 a_size = assignment->buffer_size_(*a); - const int64 b_size = assignment->buffer_size_(*b); - if (a_size != b_size) { - return a_size > b_size; // use ">" for decreasing size. - } - // Otherwise live out buffers come before others, if the - // instructions are sequentially ordered. - if (has_sequential_order) { - const bool a_live_out = liveness.MaybeLiveOut(*a); - const bool b_live_out = liveness.MaybeLiveOut(*b); - if (a_live_out != b_live_out) { - return a_live_out; - } - } - // Final tiebreaker is in instruction post order. - return post_order_position.at(a->instruction()) < - post_order_position.at(b->instruction()); - }); + absl::c_sort(sorted_buffers, + [has_sequential_order, &liveness, &post_order_position, + assignment](const LogicalBuffer* a, const LogicalBuffer* b) { + // Primary sort is by decreasing buffer size. + const int64 a_size = assignment->buffer_size_(*a); + const int64 b_size = assignment->buffer_size_(*b); + if (a_size != b_size) { + return a_size > b_size; // use ">" for decreasing size. + } + // Otherwise live out buffers come before others, if the + // instructions are sequentially ordered. + if (has_sequential_order) { + const bool a_live_out = liveness.MaybeLiveOut(*a); + const bool b_live_out = liveness.MaybeLiveOut(*b); + if (a_live_out != b_live_out) { + return a_live_out; + } + } + // Final tiebreaker is in instruction post order. + return post_order_position.at(a->instruction()) < + post_order_position.at(b->instruction()); + }); // BufferAllocations are necessarily created in decreasing size order. Keep // indices of previously created BufferAllocations in allocation_indices. std::vector allocation_indices; for (const LogicalBuffer* buffer : sorted_buffers) { VLOG(3) << "Assigning allocation to: " << *buffer; - if (colocated_buffers.count(buffer) > 0) { + if (colocated_buffers.contains(buffer)) { // Colocated buffers are currently assigned in an earlier pass. VLOG(3) << "Skipping colocated buffer: " << *buffer; continue; @@ -1020,10 +1013,14 @@ Status BufferAssigner::AssignBuffersForComputation( // callers. BufferAllocation* allocation = assignment->NewAllocation(*buffer, buffer_size); + bool parameter_has_alias = + assignment->module().input_output_alias_config().ParameterHasAlias( + instruction->parameter_number(), buffer->index()); allocation->set_entry_computation_parameter( - instruction->parameter_number(), buffer->index()); - VLOG(3) << "New allocation #" << allocation->index() - << " for entry computation parameter: " << *buffer; + instruction->parameter_number(), buffer->index(), + parameter_has_alias); + VLOG(3) << "Mark allocation #" << allocation->index() + << " as entry computation parameter: " << *buffer; continue; } @@ -1056,7 +1053,7 @@ Status BufferAssigner::AssignBuffersForComputation( assignment->GetAllSlices(operand, /*index=*/{})) { BufferAllocation* allocation = assignment->GetMutableAllocation(operand_slice.index()); - if (colocated_allocations.count(allocation->index()) == 0) { + if (!colocated_allocations.contains(allocation->index())) { // TODO(b/32491382) Colocated buffers are currently assigned in an // earlier pass, and so can break the "increasing allocation size" // invariant in this function (causing this CHECK to fail). However, @@ -1087,7 +1084,7 @@ Status BufferAssigner::AssignBuffersForComputation( // Instructions are iterated in increasing buffer size, so any // previously create allocation must be large enough to hold this // instruction's output (with the exception of colocated buffers). - if (colocated_allocations.count(allocation->index()) == 0) { + if (!colocated_allocations.contains(allocation->index())) { // TODO(b/32491382) Colocated buffers are currently assigned in an // earlier pass, and so can break the "increasing allocation size" // invariant in this function (causing this CHECK to fail). However, @@ -1313,10 +1310,10 @@ std::vector ComputePeakMemoryLogicalBuffers( live_buffers.end()); // Stabily sort the live buffers. - std::sort(live_buffers_vector.begin(), live_buffers_vector.end(), - [](const LogicalBuffer* a, const LogicalBuffer* b) { - return a->id() < b->id(); - }); + absl::c_sort(live_buffers_vector, + [](const LogicalBuffer* a, const LogicalBuffer* b) { + return a->id() < b->id(); + }); return live_buffers_vector; } @@ -1376,7 +1373,7 @@ void BufferAssigner::AddSetToColocatedBufferSets( std::vector overlap_set_indices; for (size_t index = 0; index < colocated_buffer_sets->size(); ++index) { for (const LogicalBuffer* buffer : colocated_set) { - if ((*colocated_buffer_sets)[index].count(buffer) > 0) { + if ((*colocated_buffer_sets)[index].contains(buffer)) { VLOG(5) << "Found overlap with existing set on buffer " << buffer->ToString() << "\n" << ColocatedBufferSetsToString((*colocated_buffer_sets)[index], @@ -1425,12 +1422,14 @@ BufferAssigner::MergeColocatedBufferSets( << colocated_buffer_sets.size(); // Returns true if the given buffer is for the entry parameter. - auto is_entry_parameter = [](const LogicalBuffer& buffer) { + auto is_readonly_entry_parameter = [](const LogicalBuffer& buffer) { auto* instruction = buffer.instruction(); auto* computation = instruction->parent(); auto* module = computation->parent(); return instruction->opcode() == HloOpcode::kParameter && - computation == module->entry_computation(); + computation == module->entry_computation() && + !module->input_output_alias_config().ParameterHasAlias( + instruction->parameter_number(), buffer.index()); }; std::vector set_can_be_merged(colocated_buffer_sets.size(), true); @@ -1452,7 +1451,7 @@ BufferAssigner::MergeColocatedBufferSets( for (int64 i = 0; i < colocated_buffer_sets.size(); ++i) { for (auto& buffer : colocated_buffer_sets[i]) { if (buffer_liveness.MaybeLiveOut(*buffer) || - is_entry_parameter(*buffer) || + is_readonly_entry_parameter(*buffer) || buffer->instruction()->opcode() == HloOpcode::kConstant) { set_can_be_merged[i] = false; break; @@ -1539,15 +1538,16 @@ void BufferAssigner::BuildColocatedBufferSets( VLOG(4) << "Input/Output Alias Config: "; VLOG(4) << module->input_output_alias_config(); module->input_output_alias_config().ForEachAlias( - [&](const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& param_index) { + [&](const ShapeIndex& output_index, + const HloInputOutputAliasConfig::Alias& alias) { std::vector colocated_set; AddBufferToColocatedSet(module->entry_computation()->root_instruction(), output_index, points_to_analysis, &colocated_set); AddBufferToColocatedSet( - module->entry_computation()->parameter_instruction(param_number), - param_index, points_to_analysis, &colocated_set); + module->entry_computation()->parameter_instruction( + alias.parameter_number), + alias.parameter_index, points_to_analysis, &colocated_set); AddSetToColocatedBufferSets(colocated_set, colocated_buffer_sets); }); @@ -1741,10 +1741,6 @@ void BufferAssigner::AssignColocatedBufferSets( // module-level scope, we can allow buffers to be shared across // computations (in some cases). allocation = assignment->NewAllocation(*buffer, buffer_size); - if (entry_parameter_number >= 0) { - allocation->set_entry_computation_parameter( - entry_parameter_number, *entry_parameter_shape_idx); - } if (is_constant) { allocation->set_constant(true); } @@ -1758,6 +1754,16 @@ void BufferAssigner::AssignColocatedBufferSets( } colocated_buffers->insert(buffer); } + + // If an allocation contains a parameter, set corresponding fields. + if (entry_parameter_number >= 0) { + bool parameter_has_alias = + assignment->module().input_output_alias_config().ParameterHasAlias( + entry_parameter_number, *entry_parameter_shape_idx); + allocation->set_entry_computation_parameter(entry_parameter_number, + *entry_parameter_shape_idx, + parameter_has_alias); + } } } diff --git a/tensorflow/compiler/xla/service/buffer_assignment.h b/tensorflow/compiler/xla/service/buffer_assignment.h index 0a9fdede803e84ca42472259084615c031b206eb..448dec3b1aa0c0f85e1060a70e965fcf3952c320 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment.h +++ b/tensorflow/compiler/xla/service/buffer_assignment.h @@ -96,7 +96,11 @@ class BufferAllocation { // Whether this allocation is readonly i.e. backed by memory we cannot write // to. bool is_readonly() const { - return is_entry_computation_parameter() || is_constant(); + // Entry parameters are generally readonly, except when they are aliased + // with any output. + return (is_entry_computation_parameter() && + !is_parameter_aliased_with_output_) || + is_constant(); } bool is_tuple() const { return is_tuple_; } @@ -186,9 +190,10 @@ class BufferAllocation { end > other.offset_; } - struct Hasher { - size_t operator()(Slice s) const; - }; + template + friend H AbslHashValue(H h, const Slice& s) { + return H::combine(std::move(h), s.index(), s.offset(), s.size()); + } string ToString() const; @@ -273,8 +278,10 @@ class BufferAllocation { void AddAssignment(const LogicalBuffer& buffer, int64 offset, int64 size); void set_entry_computation_parameter(int64 parameter_number, - ShapeIndex param_shape_index) { + ShapeIndex param_shape_index, + bool parameter_aliased_with_output) { is_entry_computation_parameter_ = true; + is_parameter_aliased_with_output_ = parameter_aliased_with_output; parameter_number_ = parameter_number; param_shape_index_ = std::move(param_shape_index); } @@ -304,6 +311,9 @@ class BufferAllocation { // outlast the computation. bool is_entry_computation_parameter_ = false; + // Whether this entry computation parameter is aliased with output. + bool is_parameter_aliased_with_output_ = false; + // If this allocation holds an entry computation parameter, this field // indicates the index (starting from 0) of the parameter. int64 parameter_number_ = 0; diff --git a/tensorflow/compiler/xla/service/buffer_assignment_test.cc b/tensorflow/compiler/xla/service/buffer_assignment_test.cc index 8f482e6ba8c3e71c9980be5e6947ea61f3b4ef29..580bc2f43384006eab8711490689a200fc887d37 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment_test.cc +++ b/tensorflow/compiler/xla/service/buffer_assignment_test.cc @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_set.h" #include "absl/memory/memory.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/service/buffer_value.h" @@ -309,7 +310,7 @@ class BufferAssignmentTest : public HloTestBase { static bool BuffersDistinct(const std::vector& a, const std::vector& b, const BufferAssignment& assignment) { - std::set a_slices; + absl::flat_hash_set a_slices; for (const HloInstruction* instruction : a) { if (assignment.HasTopLevelAllocation(instruction)) { a_slices.insert( @@ -319,8 +320,8 @@ static bool BuffersDistinct(const std::vector& a, for (const HloInstruction* instruction : b) { if (assignment.HasTopLevelAllocation(instruction)) { - if (a_slices.count(assignment.GetUniqueTopLevelSlice(instruction) - .ConsumeValueOrDie())) { + if (a_slices.contains(assignment.GetUniqueTopLevelSlice(instruction) + .ConsumeValueOrDie())) { return false; } } @@ -464,6 +465,40 @@ TEST_F(BufferAssignmentTest, Basic) { GetAssignedOutputAllocation(*buffers, sub); } +TEST_F(BufferAssignmentTest, AliasedParamCanBeReused) { + // If an input buffer and output buffer aliases, the input buffer can be + // reused for other intermediate results. + // + // param0[100] ----- (neg1) -- (neg2) + // | | + // + -------- Aliased ---------+ + + auto builder = HloComputation::Builder(TestName()); + + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32vec100_, "p0")); + auto neg_1 = builder.AddInstruction( + HloInstruction::CreateUnary(f32vec100_, HloOpcode::kNegate, param)); + auto neg_2 = builder.AddInstruction( + HloInstruction::CreateUnary(f32vec100_, HloOpcode::kNegate, neg_1)); + + auto module = CreateNewVerifiedModule(); + module->AddEntryComputation(builder.Build()); + + TF_ASSERT_OK(module->input_output_alias_config().SetUpAlias( + {}, 0, {}, HloInputOutputAliasConfig::kUserAlias)); + + auto buffers = RunBufferAssignment(module.get()); + + BufferAllocation param_buffer = GetAssignedInputAllocation(*buffers, param); + BufferAllocation neg_1_buffer = GetAllocation(*buffers, neg_1, {}); + BufferAllocation neg_2_buffer = GetAllocation(*buffers, neg_2, {}); + + // Everything use one buffer. + EXPECT_EQ(param_buffer.index(), neg_1_buffer.index()); + EXPECT_EQ(neg_2_buffer.index(), neg_1_buffer.index()); +} + TEST_F(BufferAssignmentTest, AddCannotReuse) { // Pass in a special rule to indicate that "add" cannot reuse any buffer. // @@ -2485,9 +2520,9 @@ while_body { get-tuple-element.3 = s32[] get-tuple-element(state), index=0 constant.2 = s32[] constant(128) add.5 = s32[] add(get-tuple-element.3, constant.2) - constant.3 = s32[3]{0} constant({0, 0, 0}) - dynamic-update-slice.5 = f32[1280,1,128]{2,1,0} dynamic-update-slice(get-tuple-element.4, broadcast.6, constant.3) - dynamic-update-slice.9 = f32[1280,1,128]{2,1,0} dynamic-update-slice(dynamic-update-slice.5, broadcast.6, constant.3) + constant.3 = s32[] constant(0) + dynamic-update-slice.5 = f32[1280,1,128]{2,1,0} dynamic-update-slice(get-tuple-element.4, broadcast.6, constant.3, constant.3, constant.3) + dynamic-update-slice.9 = f32[1280,1,128]{2,1,0} dynamic-update-slice(dynamic-update-slice.5, broadcast.6, constant.3, constant.3, constant.3) ROOT tuple.85 = (s32[], s32[], s32[2]{0}, f32[1280,1,128]{2,1,0}) tuple(add.5, dynamic-update-slice.9) } diff --git a/tensorflow/compiler/xla/service/buffer_liveness_test.cc b/tensorflow/compiler/xla/service/buffer_liveness_test.cc index 46218e6499ac5b71c75a1de2588918db0be572fe..23b9af0281b0d5ee1ef6ca2315f0cc1042285609 100644 --- a/tensorflow/compiler/xla/service/buffer_liveness_test.cc +++ b/tensorflow/compiler/xla/service/buffer_liveness_test.cc @@ -638,10 +638,10 @@ class FusedDynamicUpdateSliceLivenessTest : public BufferLivenessTest { } // Create a DynamicUpdateSlice instruction of tuple element 1 with 'update'. auto starts = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({2}))); + HloInstruction::CreateConstant(LiteralUtil::CreateR0(2))); auto dynamic_update_slice = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - data_shape, gte1, update, starts)); + data_shape, gte1, update, {starts})); // Create output tuple. builder.AddInstruction( HloInstruction::CreateTuple({gte0, dynamic_update_slice})); @@ -794,10 +794,10 @@ class DynamicUpdateSliceLivenessTest : public BufferLivenessTest { } // Create a DynamicUpdateSlice instruction of tuple element 1 with 'update'. auto starts = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({2}))); + HloInstruction::CreateConstant(LiteralUtil::CreateR0(2))); auto dynamic_update_slice = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - data_shape, gte1, update, starts)); + data_shape, gte1, update, {starts})); // Create output tuple. auto tuple_root = builder.AddInstruction( HloInstruction::CreateTuple({gte0, dynamic_update_slice})); diff --git a/tensorflow/compiler/xla/service/call_graph.cc b/tensorflow/compiler/xla/service/call_graph.cc index 94af788c54f6c722997311bec50da3ed93aa3cee..98304757cae91d22466ed25f8c6e36ce90a848db 100644 --- a/tensorflow/compiler/xla/service/call_graph.cc +++ b/tensorflow/compiler/xla/service/call_graph.cc @@ -64,6 +64,7 @@ CallContext GetInstructionCallContext(HloOpcode opcode) { case HloOpcode::kReduceWindow: case HloOpcode::kScatter: case HloOpcode::kSelectAndScatter: + case HloOpcode::kSort: case HloOpcode::kFusion: return CallContext::kParallel; default: diff --git a/tensorflow/compiler/xla/service/call_graph.h b/tensorflow/compiler/xla/service/call_graph.h index f41367914c0cb5fe66b1dbbc5ec6f8b7a67d592c..c02ffda575278905f6549b362e5e7d94f5713b36 100644 --- a/tensorflow/compiler/xla/service/call_graph.h +++ b/tensorflow/compiler/xla/service/call_graph.h @@ -256,6 +256,10 @@ class CallGraph { private: CallGraph(const HloModule* module); + // Not copyable. + CallGraph(const CallGraph&) = delete; + CallGraph& operator=(const CallGraph&) = delete; + // Sets the call contexts for every node in the graph. void SetCallContexts(); diff --git a/tensorflow/compiler/xla/service/call_graph_test.cc b/tensorflow/compiler/xla/service/call_graph_test.cc index 79241342005d3dc59051ea638798062ff56f11a9..5de724f8924b78008ba4c56603b61bf93fbc5e7c 100644 --- a/tensorflow/compiler/xla/service/call_graph_test.cc +++ b/tensorflow/compiler/xla/service/call_graph_test.cc @@ -384,7 +384,7 @@ TEST_F(CallGraphTest, ComplexGraph) { // Verify visitation order of some computations in the graph. auto index_of = [&visited](const HloComputation* comp) { - auto it = std::find(visited.begin(), visited.end(), comp); + auto it = absl::c_find(visited, comp); EXPECT_NE(it, visited.end()); return std::distance(visited.begin(), it); }; diff --git a/tensorflow/compiler/xla/service/channel_tracker.cc b/tensorflow/compiler/xla/service/channel_tracker.cc index 3c2d1ae6d82ebc6c10d52194fd1cec5e291025f7..b517495f2ea0c75679685c67f757ff586f8c79e3 100644 --- a/tensorflow/compiler/xla/service/channel_tracker.cc +++ b/tensorflow/compiler/xla/service/channel_tracker.cc @@ -72,7 +72,7 @@ ChannelHandle ChannelTracker::AllocateHandle(ChannelHandle::ChannelType type) { } Status ChannelTracker::RegisterSendInternal(const ChannelHandle& handle) { - if (opaque_to_channel_.count(handle.handle()) == 0) { + if (!opaque_to_channel_.contains(handle.handle())) { return NotFound("channel handle not found: %d", handle.handle()); } Channel& channel = opaque_to_channel_[handle.handle()]; @@ -94,7 +94,7 @@ Status ChannelTracker::RegisterSendInternal(const ChannelHandle& handle) { } Status ChannelTracker::RegisterRecvInternal(const ChannelHandle& handle) { - if (opaque_to_channel_.count(handle.handle()) == 0) { + if (!opaque_to_channel_.contains(handle.handle())) { return NotFound("channel handle not found: %d", handle.handle()); } Channel& channel = opaque_to_channel_[handle.handle()]; diff --git a/tensorflow/compiler/xla/service/channel_tracker.h b/tensorflow/compiler/xla/service/channel_tracker.h index 52037bf9b52556c6aa2e66dd3209e25cf085cfe3..89e17eba36f23077ce4cf0704e7455b76bee68d1 100644 --- a/tensorflow/compiler/xla/service/channel_tracker.h +++ b/tensorflow/compiler/xla/service/channel_tracker.h @@ -18,6 +18,7 @@ limitations under the License. #include +#include "absl/container/flat_hash_map.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/status.h" @@ -83,7 +84,8 @@ class ChannelTracker { // Mapping from ChannelHandle value to the corresponding registered // Channel object. - std::map opaque_to_channel_ GUARDED_BY(channel_mutex_); + absl::flat_hash_map opaque_to_channel_ + GUARDED_BY(channel_mutex_); TF_DISALLOW_COPY_AND_ASSIGN(ChannelTracker); }; diff --git a/tensorflow/compiler/xla/service/compiler.cc b/tensorflow/compiler/xla/service/compiler.cc index 8f08c244908efb823b3870c19bdc3491fa87d44f..653f4555a77cc82e91fb1cd26206b93826375732 100644 --- a/tensorflow/compiler/xla/service/compiler.cc +++ b/tensorflow/compiler/xla/service/compiler.cc @@ -98,10 +98,17 @@ Compiler::GetPlatformCompilers() { auto* factories = GetPlatformCompilerFactories(); auto it = factories->find(platform->id()); if (it == factories->end()) { + string hint; + if (platform->Name() == "Host") { + hint = " (hint: try linking in tensorflow/compiler/jit:xla_cpu_jit)"; + } else if (platform->Name() == "CUDA") { + hint = " (hint: try linking in tensorflow/compiler/jit:xla_gpu_jit)"; + } + return NotFound( "could not find registered compiler for platform %s -- check " - "target linkage", - platform->Name()); + "target linkage%s", + platform->Name(), hint); } // And then we invoke the factory, placing the result into the mapping. diff --git a/tensorflow/compiler/xla/service/computation_layout.cc b/tensorflow/compiler/xla/service/computation_layout.cc index efc893818d03a20d6bd65b7dc1da72ea5da5ceb0..92d1ca4ba5da802a5f1c544017ac52dda38e9b1d 100644 --- a/tensorflow/compiler/xla/service/computation_layout.cc +++ b/tensorflow/compiler/xla/service/computation_layout.cc @@ -42,8 +42,8 @@ void ComputationLayout::SetToDefaultLayout() { } bool ComputationLayout::LayoutIsSet() const { - return std::all_of(parameter_layouts_.begin(), parameter_layouts_.end(), - [](const ShapeLayout& s) { return s.LayoutIsSet(); }) && + return absl::c_all_of(parameter_layouts_, + [](const ShapeLayout& s) { return s.LayoutIsSet(); }) && result_layout_.LayoutIsSet(); } diff --git a/tensorflow/compiler/xla/service/convolution_group_converter.cc b/tensorflow/compiler/xla/service/convolution_group_converter.cc index 7a24faec17f0c4f0a57406328b1c21cd73506d82..f11f9e5fc2949a92f83ff66506a9b162ffda1c92 100644 --- a/tensorflow/compiler/xla/service/convolution_group_converter.cc +++ b/tensorflow/compiler/xla/service/convolution_group_converter.cc @@ -207,88 +207,23 @@ Status ConvolutionVisitor::HandleBatchGroupCount(HloInstruction* convolution) { return Status::OK(); } - VLOG(2) << "Dealing with batch_group_count " << batch_group_count << "\n"; + VLOG(2) << "Dealing with batch_group_count " << batch_group_count + << " for convolution " << convolution->ToString() << "\n"; auto add = [&](std::unique_ptr inst) { return computation_->AddInstruction(std::move(inst)); }; int64 input_batch_dimension = dim_numbers.input_batch_dimension(); - int64 input_feature_dimension = dim_numbers.input_feature_dimension(); int64 output_batch_dimension = dim_numbers.output_batch_dimension(); int64 output_feature_dimension = dim_numbers.output_feature_dimension(); - int64 kernel_input_feature_dimension = - dim_numbers.kernel_input_feature_dimension(); int64 input_batch = activation->shape().dimensions(input_batch_dimension); // We are not yet supporting batch_group of sizes greater than 1. TF_RET_CHECK(input_batch == batch_group_count); - if (is_cost_viable_(convolution)) { - // Add a dimension to the activation, and reshape. - Shape reshaped_activation_shape = activation->shape(); - ShapeUtil::AppendMajorDimension(1, &reshaped_activation_shape); - - activation = add( - HloInstruction::CreateReshape(reshaped_activation_shape, activation)); - - // Add a dimension to the filter, and reshape. - Shape reshaped_filter_shape = filter->shape(); - ShapeUtil::AppendMajorDimension(1, &reshaped_filter_shape); - - filter = add(HloInstruction::CreateReshape(reshaped_filter_shape, filter)); - - int64 new_spatial_dim = reshaped_activation_shape.dimensions().size() - 1; - - Shape new_output_shape = convolution->shape(); - ShapeUtil::AppendMajorDimension(1, &new_output_shape); - - int64 input_feature = - activation->shape().dimensions(input_feature_dimension); - - // The code below edits convolution dimension numbers. Please refer to - // conv_op_helpers.cc to find how the dimensions were set up originally. - - // Effectively, the new input batch becomes 1, and so does the kernel - // input feature. The original input batch now becomes a spatial dimension. - // The output batch (remember that the output is the new kernel for in - // backprop) becomes a spatial dimension too. - - dim_numbers.set_input_batch_dimension(new_spatial_dim); - dim_numbers.set_input_feature_dimension(input_batch_dimension); - dim_numbers.set_kernel_input_feature_dimension(new_spatial_dim); - - dim_numbers.add_input_spatial_dimensions(input_feature_dimension); - dim_numbers.add_kernel_spatial_dimensions(kernel_input_feature_dimension); - - dim_numbers.add_output_spatial_dimensions(output_batch_dimension); - dim_numbers.set_output_batch_dimension(new_spatial_dim); - - // Add window for the new spatial dimension. - Window new_window = convolution->window(); - auto* dim = new_window.add_dimensions(); - dim->set_window_dilation(1); - dim->set_base_dilation(1); - dim->set_stride(1); - dim->set_size(input_feature); - - auto new_convolution = add(HloInstruction::CreateConvolve( - new_output_shape, activation, filter, - /*feature_group_count=*/batch_group_count, /*batch_group_count=*/1, - new_window, dim_numbers, convolution->precision_config())); - - // Delete the extra spatial dimension, and reshape. - Shape reshaped_convolution_shape = ShapeUtil::DeleteDimension( - new_spatial_dim - 1, new_convolution->shape()); - auto reshaped_convolution = HloInstruction::CreateReshape( - reshaped_convolution_shape, new_convolution); - - TF_RETURN_IF_ERROR(computation_->ReplaceWithNewInstruction( - convolution, std::move(reshaped_convolution))); - - changed_ = true; - } else { + if (!is_cost_viable_(convolution) || filter_expansion_) { // We first obtain the expanded the filter (which is the convolution // output). The batch dimension is the expanded one (which originally // represents kernel input feature dimension). We mask the filter to zero @@ -315,14 +250,27 @@ Status ConvolutionVisitor::HandleBatchGroupCount(HloInstruction* convolution) { expanded_filter_shape, HloOpcode::kSelect, filter_mask, new_convolution, zero_filter)); - auto zero_literal = LiteralUtil::CreateR0(0.0f); - TF_ASSIGN_OR_RETURN(zero_literal, zero_literal.Convert(F32)); + PrimitiveType reduce_type = new_filter->shape().element_type(); + auto reduce_window_shape = new_convolution->shape(); + reduce_window_shape.set_dimensions(output_batch_dimension, 1); + + // Ensure that data input to reduce window uses at least 32 bits. + if (primitive_util::BitWidth(reduce_type) < primitive_util::BitWidth(F32)) { + reduce_type = F32; + reduce_window_shape.set_element_type(F32); + Shape convert_shape = new_filter->shape(); + convert_shape.set_element_type(F32); + new_filter = + add(HloInstruction::CreateConvert(convert_shape, new_filter)); + } + + auto zero_literal = LiteralUtil::Zero(reduce_type); auto zero_scalar = add(HloInstruction::CreateConstant(std::move(zero_literal))); auto reduce_function = [&]() -> HloComputation* { HloComputation::Builder b("add_computation"); - Shape shape = ShapeUtil::MakeShape(F32, {}); + Shape shape = ShapeUtil::MakeShape(reduce_type, {}); auto lhs = b.AddInstruction(HloInstruction::CreateParameter(0, shape, "lhs")); auto rhs = @@ -332,18 +280,6 @@ Status ConvolutionVisitor::HandleBatchGroupCount(HloInstruction* convolution) { return computation_->parent()->AddEmbeddedComputation(b.Build(scalar_op)); }; - // Ensure that data input to reduce window is of type F32. - if (primitive_util::BitWidth(new_filter->shape().element_type()) < - primitive_util::BitWidth(F32)) { - Shape convert_shape = new_filter->shape(); - convert_shape.set_element_type(F32); - new_filter = - add(HloInstruction::CreateBitcastConvert(convert_shape, new_filter)); - } - - auto reduce_window_shape = new_convolution->shape(); - reduce_window_shape.set_dimensions(output_batch_dimension, 1); - // Create the reduce window. Window window; for (int64 i = 0; i < new_convolution->shape().dimensions_size(); ++i) { @@ -369,10 +305,11 @@ Status ConvolutionVisitor::HandleBatchGroupCount(HloInstruction* convolution) { // Convert reduced data back to the original data type. auto reduce_window_converted = - HloInstruction::CreateBitcastConvert(convert_back_shape, reduce_window); + HloInstruction::CreateConvert(convert_back_shape, reduce_window); TF_RETURN_IF_ERROR(computation_->ReplaceWithNewInstruction( convolution, std::move(reduce_window_converted))); + changed_ = true; } return Status::OK(); diff --git a/tensorflow/compiler/xla/service/convolution_feature_group_converter_test.cc b/tensorflow/compiler/xla/service/convolution_group_converter_test.cc similarity index 76% rename from tensorflow/compiler/xla/service/convolution_feature_group_converter_test.cc rename to tensorflow/compiler/xla/service/convolution_group_converter_test.cc index d58f157242f5fb9690f7fda3e7d8f71ca6c8db84..9cee3eda95252d6c7d725fbb03030bd58f52e71f 100644 --- a/tensorflow/compiler/xla/service/convolution_feature_group_converter_test.cc +++ b/tensorflow/compiler/xla/service/convolution_group_converter_test.cc @@ -94,5 +94,32 @@ ENTRY %Convolve1D1Window_0.v3 (input: f32[1,2,4], filter: f32[1,2,2]) -> f32[1,2 EXPECT_EQ(root->operand(1)->feature_group_count(), 1); } +TEST_F(ConvolutionGroupConverterTest, + ConvertBatchGroupCountEqualToInputBatchDim) { + string hlo_string = R"(HloModule Convolve1D1Window_0_module + +ENTRY %Convolve1D1Window_0.v3 (input: f32[16,19,19,512]{3,2,1,0}, filter: f32[16,19,19,512]{3,2,1,0}) -> f32[3,3,512,1]{3,2,1,0} { + %input = f32[16,19,19,512]{3,2,1,0} parameter(0) + %filter = f32[16,19,19,512]{3,2,1,0} parameter(1) + ROOT %convolution = f32[3,3,512,1]{3,2,1,0} convolution(f32[16,19,19,512]{3,2,1,0} %input, f32[16,19,19,512]{3,2,1,0} %filter), window={size=19x19 pad=1_1x1_1}, dim_labels=f01b_i01o->01fb, batch_group_count=512 + })"; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(hlo_string)); + + auto computation = module->entry_computation(); + HloInstruction* root = computation->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kConvolution); + auto cost_model = [](HloInstruction* conv) { return false; }; + ConvolutionGroupConverter converter(cost_model, /*convert_batch_groups_only=*/ + true); + ASSERT_TRUE(converter.Run(module.get()).ValueOrDie()); + root = computation->root_instruction(); + + // Verify that the convolution is replaced by a convert. + EXPECT_EQ(root->opcode(), HloOpcode::kConvert); + // Make sure the convert is being fed by a reduce window. + EXPECT_EQ(root->operand(0)->opcode(), HloOpcode::kReduceWindow); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/copy_insertion.cc b/tensorflow/compiler/xla/service/copy_insertion.cc index e0165d557deb390636b42df12f099792042f1b65..5e26a63cebfa9b2e50f4b13335c10c246999d4df 100644 --- a/tensorflow/compiler/xla/service/copy_insertion.cc +++ b/tensorflow/compiler/xla/service/copy_insertion.cc @@ -349,11 +349,12 @@ Status AddCopiesForAliasedInputOutputs(HloModule* module) { ShapeTree param_indices_to_copy(param->shape()); module->input_output_alias_config().ForEachAlias( - [&](const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& param_index) { - if (param_number == param->parameter_number()) { + [&](const ShapeIndex& output_index, + const HloInputOutputAliasConfig::Alias& alias) { + if (alias.parameter_number == param->parameter_number()) { param_has_alias = true; - *(param_indices_to_copy.mutable_element(param_index)) = true; + *(param_indices_to_copy.mutable_element(alias.parameter_index)) = + true; *(output_indices_to_copy.mutable_element(output_index)) = true; } }); @@ -395,13 +396,14 @@ Status AddCopiesForAliasedInputOutputs(HloModule* module) { // Add control dependencies between the input/output copies. TF_RETURN_IF_ERROR(module->input_output_alias_config().ForEachAliasWithStatus( - [&](const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& input_index) -> Status { - if (!copied_parameters[param_number]) { + [&](const ShapeIndex& output_index, + const HloInputOutputAliasConfig::Alias& alias) -> Status { + if (!copied_parameters[alias.parameter_number]) { return Status::OK(); } HloInstruction* from = - copied_parameters[param_number]->element(input_index); + copied_parameters[alias.parameter_number]->element( + alias.parameter_index); HloInstruction* to = output_copy_tree.element(output_index); TF_RET_CHECK(from != nullptr); @@ -539,10 +541,9 @@ class CopyRemover { } std::vector values = buffer.values(); - std::sort(values.begin(), values.end(), - [this](const HloValue* a, const HloValue* b) { - return ordering_.IsDefinedBefore(*a, *b); - }); + absl::c_sort(values, [this](const HloValue* a, const HloValue* b) { + return ordering_.IsDefinedBefore(*a, *b); + }); // Create a list containing all of the values in the buffer. AddValueList(values, &value_to_node); @@ -842,12 +843,11 @@ class CopyRemover { copy_value_node->next->prev = operand_node; // Patch up uses. Remove use of copy from operand_node uses. - auto it = - std::find_if(operand_node->uses.begin(), operand_node->uses.end(), - [copy_value_node](const HloUse* use) { - return use->instruction == - copy_value_node->value->defining_instruction(); - }); + auto it = absl::c_find_if( + operand_node->uses, [copy_value_node](const HloUse* use) { + return use->instruction == + copy_value_node->value->defining_instruction(); + }); CHECK(it != operand_node->uses.end()); operand_node->uses.erase(it); diff --git a/tensorflow/compiler/xla/service/copy_insertion_test.cc b/tensorflow/compiler/xla/service/copy_insertion_test.cc index e4e9d7ba05c115be9dd0eb53ebd7de208d514efb..4391bdcba532661a0fde789e2c4ed324c40bcd32 100644 --- a/tensorflow/compiler/xla/service/copy_insertion_test.cc +++ b/tensorflow/compiler/xla/service/copy_insertion_test.cc @@ -1376,9 +1376,11 @@ TEST_F(CopyInsertionTest, CrossingParameters) { builder.AddInstruction(HloInstruction::CreateTuple({gte1, gte0})); module->AddEntryComputation(builder.Build()); ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); + /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1})); + /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); InsertCopies(module.get()); EXPECT_EQ(CountCopies(*module), 4); @@ -1409,9 +1411,11 @@ TEST_F(CopyInsertionTest, ParametersAliasing) { builder.AddInstruction(HloInstruction::CreateTuple({gte0, gte1})); module->AddEntryComputation(builder.Build()); ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); + /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1})); + /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); InsertCopies(module.get()); EXPECT_EQ(CountCopies(*module), 0); @@ -1475,7 +1479,8 @@ TEST_F(CopyInsertionTest, ParameterWithPartialAliasing) { builder.AddInstruction(HloInstruction::CreateTuple({gte0, gte1})); module->AddEntryComputation(builder.Build()); ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); + /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); InsertCopies(module.get()); EXPECT_THAT(module->entry_computation()->root_instruction(), @@ -1516,7 +1521,8 @@ TEST_F(CopyInsertionTest, ParameterAndParallelOpsWithPartialAliasing) { builder.AddInstruction(HloInstruction::CreateTuple({negate0, negate1})); module->AddEntryComputation(builder.Build()); ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); + /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); InsertCopies(module.get()); EXPECT_EQ(CountCopies(*module), 0); @@ -1557,7 +1563,8 @@ TEST_F(CopyInsertionTest, ParameterAndOpsWithPartialAliasing) { builder.AddInstruction(HloInstruction::CreateTuple({add, negate1})); module->AddEntryComputation(builder.Build()); ASSERT_IS_OK(module->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); + /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); InsertCopies(module.get()); EXPECT_EQ(CountCopies(*module), 0); @@ -1848,8 +1855,7 @@ ENTRY %TokensShouldNotBeCopied () -> s32[] { } )"; TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - HloRunner::CreateModuleFromString( - module_string, GetDebugOptionsForTest())); + ParseAndReturnVerifiedModule(module_string)); InsertCopies(module.get()); // There should be no copies added because tokens should not be copied. @@ -2112,8 +2118,7 @@ ENTRY TestComputation { ROOT while.3 = (s32[], s32[], s32[], s32[], s32[]) while(arg_tuple.6), condition=cond_wrapper.v3.2, body=_functionalize_body_2__.v25 } )"; - auto module_or_status = - HloRunner::CreateModuleFromString(hlo_string, GetDebugOptionsForTest()); + auto module_or_status = ParseAndReturnVerifiedModule(hlo_string); auto module = module_or_status.ConsumeValueOrDie(); InsertCopies(module.get()); } @@ -2213,8 +2218,7 @@ ENTRY TestComputation { ROOT while.3 = (s32[], s32[], s32[], s32[], s32[]) while(arg_tuple.6), condition=cond_wrapper.v3.2, body=_functionalize_body_2__.v25 } )"; - auto module_or_status = - HloRunner::CreateModuleFromString(hlo_string, GetDebugOptionsForTest()); + auto module_or_status = ParseAndReturnVerifiedModule(hlo_string); auto module = module_or_status.ConsumeValueOrDie(); InsertCopies(module.get()); } @@ -2231,7 +2235,7 @@ cond.inner { body.inner { param.body.inner = pred[] parameter(0) - ROOT neg = pred[] negate(param.body.inner) + ROOT not = pred[] not(param.body.inner) } cond.outer { @@ -2248,9 +2252,8 @@ ENTRY TestComputation { ROOT while = pred[] while(entry_param), condition=cond.outer, body=body.outer } )"; - TF_ASSERT_OK_AND_ASSIGN( - std::unique_ptr module, - HloRunner::CreateModuleFromString(hlo_string, GetDebugOptionsForTest())); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_string)); InsertCopies(module.get()); // There should only be a single copy inserted, and it's in the entry diff --git a/tensorflow/compiler/xla/service/cpu/BUILD b/tensorflow/compiler/xla/service/cpu/BUILD index f49b5110be5c4bab63b423e5ed2e67bc1828f6e3..42672bc3875af2d732d80691df6bf85b9d8080cd 100644 --- a/tensorflow/compiler/xla/service/cpu/BUILD +++ b/tensorflow/compiler/xla/service/cpu/BUILD @@ -1,6 +1,14 @@ # Description: # LLVM-based CPU backend for XLA. +load("//tensorflow/compiler/xla:xla.bzl", "ORC_JIT_MEMORY_MAPPER_TARGETS") +load( + "//third_party/mkl:build_defs.bzl", + "mkl_deps", +) +load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test") +load(":build_defs.bzl", "runtime_copts") + licenses(["notice"]) # Apache 2.0 package( @@ -14,15 +22,6 @@ package_group( ], ) -load(":build_defs.bzl", "runtime_copts") -load("//tensorflow:tensorflow.bzl", "tf_cc_test") -load("//tensorflow:tensorflow.bzl", "tf_cc_binary") -load("//tensorflow/compiler/xla:xla.bzl", "ORC_JIT_MEMORY_MAPPER_TARGETS") -load( - "//third_party/mkl:build_defs.bzl", - "mkl_deps", -) - # Filegroup used to collect source files for dependency checking. filegroup( name = "c_srcs", @@ -95,6 +94,7 @@ cc_library( ":target_machine_features", "@com_google_absl//absl/types:span", "//tensorflow/compiler/tf2xla:cpu_function_runtime", + "//tensorflow/compiler/xla/service:hlo_casting_utils", "//tensorflow/compiler/xla/service:map_inliner", "//tensorflow/compiler/xla/service:hlo_get_dimension_size_rewriter", "//tensorflow/compiler/xla/service:scatter_expander", @@ -114,6 +114,7 @@ cc_library( "//tensorflow/compiler/xla/service:conditional_simplifier", "//tensorflow/compiler/xla/service:convolution_group_converter", "//tensorflow/compiler/xla/service:dot_decomposer", + "//tensorflow/compiler/xla/service:dynamic_index_splitter", "//tensorflow/compiler/xla/service:executable", "//tensorflow/compiler/xla/service:flatten_call_graph", "//tensorflow/compiler/xla/service:hlo", @@ -133,7 +134,9 @@ cc_library( "//tensorflow/compiler/xla/service:llvm_compiler", "//tensorflow/compiler/xla/service:reduce_precision_insertion", "//tensorflow/compiler/xla/service:reshape_mover", + "//tensorflow/compiler/xla/service:sort_simplifier", "//tensorflow/compiler/xla/service:transpose_folding", + "//tensorflow/compiler/xla/service:triangular_solve_expander", "//tensorflow/compiler/xla/service:tuple_simplifier", "//tensorflow/compiler/xla/service:while_loop_constant_sinking", "//tensorflow/compiler/xla/service:while_loop_invariant_code_motion", @@ -241,6 +244,7 @@ cc_library( "//tensorflow/compiler/xla/service:tuple_points_to_analysis", "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", + "//tensorflow/stream_executor/host:host_stream", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/types:span", @@ -364,15 +368,33 @@ cc_library( ], ) +cc_library( + name = "tiled_dot_emitter", + srcs = ["tiled_dot_emitter.cc"], + hdrs = ["tiled_dot_emitter.h"], + deps = [ + ":vector_support_library", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service/llvm_ir:kernel_support_library", + "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", + "//tensorflow/core:lib", + "@llvm//:core", + ], +) + cc_library( name = "dot_op_emitter", srcs = ["dot_op_emitter.cc"], - hdrs = ["dot_op_emitter.h"], + hdrs = [ + "dot_op_emitter.h", + ], deps = [ ":cpu_options", ":cpu_runtime", ":ir_emission_utils", ":target_machine_features", + ":tiled_dot_emitter", ":vector_support_library", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -380,6 +402,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_casting_utils", "//tensorflow/compiler/xla/service:hlo_module_config", "//tensorflow/compiler/xla/service/llvm_ir:ir_array", "//tensorflow/compiler/xla/service/llvm_ir:kernel_support_library", @@ -631,6 +654,7 @@ cc_library( deps = [ ":runtime_matvec", "//tensorflow/core:framework_lite", + "//tensorflow/core/kernels:eigen_contraction_kernel", "//third_party/eigen3", ], ) @@ -767,8 +791,6 @@ cc_library( ":target_machine_features", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/service:computation_layout", - "//tensorflow/compiler/xla/service:hlo", - "//tensorflow/compiler/xla/service:hlo_casting_utils", "//tensorflow/compiler/xla/service:layout_assignment", "//tensorflow/core:lib", "@com_google_absl//absl/container:flat_hash_map", @@ -1008,7 +1030,6 @@ tf_cc_test( size = "small", srcs = ["cpu_eigen_tensor_alignment_test.cc"], deps = [ - ":dot_op_emitter", ":ir_emission_utils", ":target_machine_features_fake", "//tensorflow/compiler/xla:test", diff --git a/tensorflow/compiler/xla/service/cpu/compiler_functor.cc b/tensorflow/compiler/xla/service/cpu/compiler_functor.cc index 796a7cf94d02b0ad42366387a9d3f8d589b8840a..414eacddfc7ba3c295c027c64c445a2046235d36 100644 --- a/tensorflow/compiler/xla/service/cpu/compiler_functor.cc +++ b/tensorflow/compiler/xla/service/cpu/compiler_functor.cc @@ -66,9 +66,14 @@ class FilteredPassManager : public llvm::legacy::PassManager { explicit FilteredPassManager(bool disable_expensive_passes) : disable_expensive_passes_(disable_expensive_passes) {} void add(llvm::Pass* p) override { + llvm::StringRef PassName = p->getPassName(); + if (PassName.contains("Warn about non-applied transformations")) { + delete p; + return; + } if (disable_expensive_passes_) { - llvm::StringRef PassName = p->getPassName(); if (PassName.contains("Unroll loops")) { + delete p; return; } } diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc index dbab839aee95b249bd993e7a0a63fd3af38f8537..3676de56a309c6ea100e5fbf27b2bcdc8673d15b 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc @@ -69,6 +69,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/simple_orc_jit.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" #include "tensorflow/compiler/xla/service/dot_decomposer.h" +#include "tensorflow/compiler/xla/service/dynamic_index_splitter.h" #include "tensorflow/compiler/xla/service/flatten_call_graph.h" #include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" @@ -92,7 +93,9 @@ limitations under the License. #include "tensorflow/compiler/xla/service/reduce_precision_insertion.h" #include "tensorflow/compiler/xla/service/reshape_mover.h" #include "tensorflow/compiler/xla/service/scatter_expander.h" +#include "tensorflow/compiler/xla/service/sort_simplifier.h" #include "tensorflow/compiler/xla/service/transpose_folding.h" +#include "tensorflow/compiler/xla/service/triangular_solve_expander.h" #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_invariant_code_motion.h" @@ -244,6 +247,7 @@ Status CpuCompiler::RunHloPassesThroughLayoutAssn( HloPassPipeline pipeline("HLO passes through layout assignment"); pipeline.AddInvariantChecker(/*layout_sensitive=*/false, /*allow_mixed_precision=*/false); + pipeline.AddPass(); pipeline.AddPass(); ReducePrecisionInsertion::AddPasses( @@ -252,11 +256,13 @@ Status CpuCompiler::RunHloPassesThroughLayoutAssn( pipeline.AddPass(); + pipeline.AddPass(); + // TODO(b/65775800): Fix wrong output bug in Call and remove the CallInliner // pass. pipeline.AddPass(); pipeline.AddPass(); - pipeline.AddPass(); + pipeline.AddPass(/*decompose_batch_dot=*/false); auto cost_model = [](HloInstruction* conv) { // We need a cost model for CPUs. Currently, do nothing. return false; @@ -279,10 +285,10 @@ Status CpuCompiler::RunHloPassesThroughLayoutAssn( /*rewrite_inference_op=*/true, /*rewrite_grad_op=*/true); pipeline.AddPass(); - AlgebraicSimplifierOptions options( - [](const Shape&, const Shape&) { return false; }); + AlgebraicSimplifierOptions options; options.set_enable_dot_strength_reduction(false); pass.AddPass(options); + pass.AddPass(); pass.AddPass(); // BatchNormExpander can create zero-sized ops, so zero-sized HLO @@ -302,7 +308,8 @@ Status CpuCompiler::RunHloPassesThroughLayoutAssn( pipeline.AddPass( [&](const HloInstruction& dot, const TransposeFolding::OperandIndices& candidate_operands) { - return PotentiallyImplementedAsEigenDot(dot, *target_machine_features) + return DotImplementationCanHandleTranspose(dot, + *target_machine_features) ? candidate_operands : TransposeFolding::OperandIndices{}; }, @@ -345,8 +352,7 @@ Status CpuCompiler::RunHloPassesAfterLayoutAssn( pass.AddInvariantChecker( /*layout_sensitive=*/true, /*allow_mixed_precision=*/false); - AlgebraicSimplifierOptions options( - [](const Shape&, const Shape&) { return true; }); + AlgebraicSimplifierOptions options; options.set_is_layout_sensitive(true); options.set_enable_dot_strength_reduction(false); pass.AddPass>(options); diff --git a/tensorflow/compiler/xla/service/cpu/cpu_eigen_tensor_alignment_test.cc b/tensorflow/compiler/xla/service/cpu/cpu_eigen_tensor_alignment_test.cc index 8727c72b6e42517b1859e98ecadb41bbceed761c..485769a373acf5ae70c471b1a5dfcfb20ff772ef 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_eigen_tensor_alignment_test.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_eigen_tensor_alignment_test.cc @@ -13,7 +13,6 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/xla/service/cpu/dot_op_emitter.h" #include "tensorflow/compiler/xla/service/cpu/ir_emission_utils.h" #include "tensorflow/compiler/xla/service/cpu/target_machine_features_fake.h" #include "tensorflow/compiler/xla/service/hlo_parser.h" @@ -28,37 +27,6 @@ namespace { class CpuEigenTensorAlignmentTest : public ::testing::Test {}; -TEST_F(CpuEigenTensorAlignmentTest, EigenDotAlignment) { - string hlo_string = R"( -HloModule DotOperation - -ENTRY DotOperation { - arg0 = f32[5,256] parameter(0) - arg1 = f32[256,1024] parameter(1) - ROOT dot = f32[5,1024] dot(arg0, arg1), lhs_contracting_dims={1}, rhs_contracting_dims={0} -} -)"; - - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseHloString(hlo_string)); - - HloInstruction* dot = module->entry_computation()->root_instruction(); - - TargetMachineFeaturesWithFakeAlignmentLogic target_machine_with_no_alignment( - [](int64 size) { return 1; }); - - EXPECT_FALSE( - PotentiallyImplementedAsEigenDot(*dot, target_machine_with_no_alignment)); - - TargetMachineFeaturesWithFakeAlignmentLogic - target_machine_with_full_alignment([](int64 size) { - return TargetMachineFeatures::kEigenExpectedTensorAlignment; - }); - - EXPECT_TRUE(PotentiallyImplementedAsEigenDot( - *dot, target_machine_with_full_alignment)); -} - TEST_F(CpuEigenTensorAlignmentTest, EigenConvAlignment) { string hlo_string = R"( HloModule ConvOperation diff --git a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc index 412c2715b9ce0a7c454f4fd1a1a354175d440fd8..23d0af34233858515af21df5e92346742a5b5dc3 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc @@ -213,6 +213,8 @@ StatusOr CpuExecutable::CreateResultShapedBuffer( /*on_host_shape=*/result_shape(), /*on_device_shape=*/result_shape(), run_options->allocator(), stream->parent()->device_ordinal()); + const HloInputOutputAliasConfig& input_output_alias = + module().input_output_alias_config(); // Move OwningDeviceMemory values which contain the array(s) of the result // into the respective location in ScopedShapedBuffer which is returned to the @@ -232,12 +234,31 @@ StatusOr CpuExecutable::CreateResultShapedBuffer( TF_ASSIGN_OR_RETURN( const BufferAllocation::Slice slice, this->assignment_->GetUniqueSlice(src, buffer_source->index())); - CHECK(!slice.allocation()->is_entry_computation_parameter()); - const BufferAllocation::Index buffer_index = slice.index(); OwningDeviceMemory& buffer = buffers[buffer_index]; - CHECK(!buffer.is_null() || buffer.size() == 0); - *device_memory = buffer.Forget(); + if (!slice.allocation()->is_entry_computation_parameter()) { + // If the buffer coming out of the result is from a parameter, the + // owning buffer will be null, and that means the caller aliased some + // parameter buffer to an output one (via the + // HloInputOutputAliasConfig API). If that is the case, the caller + // will receive a partially complete scoped shaped buffer, which they + // will have to fill up on return. Unfortunately the interface to the + // execute APIs are ShapedBuffer pointer based, which assumes caller + // ownership, and hence a buffer coming from there cannot be part of + // the new ScopedShapedBuffer we create for the result (which assumes + // ownership). + *device_memory = buffer.Forget(); + } else { + auto output_alias = input_output_alias.GetAliasedOutput( + slice.allocation()->parameter_number(), + slice.allocation()->param_shape_index()); + CHECK(output_alias) + << "Ouput buffer is coming from parameter " + << slice.allocation()->parameter_number() << " at index " + << slice.allocation()->param_shape_index() + << ", but no alias exists"; + CHECK_EQ(*output_alias, index); + } return Status::OK(); })); return std::move(result_buffer); diff --git a/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker.cc b/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker.cc index 7fbe0fa157c57eb0c274662a1de95cf5328ccfa8..4ac61f44d9f38425da2d1fc6b9495cb4deba5047 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker.cc @@ -17,7 +17,6 @@ limitations under the License. #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/errors.h" namespace xla { diff --git a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc index 527df0bd1c23bba74f32226e5622fed32f7dcf84..c4bde837e57e82584c2a007858ed8d55608acd3c 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc @@ -332,7 +332,7 @@ TEST_F(OpcodeFusionTest, Exponential_Reshape_Negate) { TEST_F(OpcodeFusionTest, Broadcast_Reshape_DynamicSlice_Tanh) { HloComputation::Builder builder(TestName()); Shape param_shape = ShapeUtil::MakeShape(F32, {8}); - Shape starts_shape = ShapeUtil::MakeShape(F32, {2}); + Shape starts_shape = ShapeUtil::MakeShape(F32, {}); Shape broadcast_shape = ShapeUtil::MakeShape(F32, {1, 8, 8}); Shape reshape_shape = ShapeUtil::MakeShape(F32, {8, 8}); Shape dynamic_slice_shape = ShapeUtil::MakeShape(F32, {4, 4}); @@ -340,13 +340,15 @@ TEST_F(OpcodeFusionTest, Broadcast_Reshape_DynamicSlice_Tanh) { HloInstruction::CreateParameter(0, param_shape, "param")); HloInstruction* param1 = builder.AddInstruction( HloInstruction::CreateParameter(1, starts_shape, "starts")); + HloInstruction* param2 = builder.AddInstruction( + HloInstruction::CreateParameter(2, starts_shape, "starts")); HloInstruction* broadcast2 = builder.AddInstruction( HloInstruction::CreateBroadcast(broadcast_shape, param0, {1})); HloInstruction* reshape3 = builder.AddInstruction( HloInstruction::CreateReshape(reshape_shape, broadcast2)); HloInstruction* dynamic_slice4 = builder.AddInstruction(HloInstruction::CreateDynamicSlice( - dynamic_slice_shape, reshape3, param1, {4, 4})); + dynamic_slice_shape, reshape3, {param1, param2}, {4, 4})); builder.AddInstruction(HloInstruction::CreateUnary( dynamic_slice_shape, HloOpcode::kTanh, dynamic_slice4)); @@ -356,7 +358,8 @@ TEST_F(OpcodeFusionTest, Broadcast_Reshape_DynamicSlice_Tanh) { RunFusionAndCheckOpcodesWereFused( module.get(), {HloOpcode::kTanh, HloOpcode::kDynamicSlice, HloOpcode::kReshape, - HloOpcode::kBroadcast, HloOpcode::kParameter, HloOpcode::kParameter}); + HloOpcode::kBroadcast, HloOpcode::kParameter, HloOpcode::kParameter, + HloOpcode::kParameter}); } TEST_F(OpcodeFusionTest, Broadcast_Negate) { @@ -381,14 +384,14 @@ TEST_F(OpcodeFusionTest, Broadcast_Negate) { TEST_F(OpcodeFusionTest, DynamicSlice_Negate) { HloComputation::Builder builder(TestName()); Shape param_shape = ShapeUtil::MakeShape(F32, {4}); - Shape slice_shape = ShapeUtil::MakeShape(F32, {1}); + Shape slice_shape = ShapeUtil::MakeShape(F32, {}); Shape result_shape = ShapeUtil::MakeShape(F32, {2}); HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, param_shape, "param")); HloInstruction* param1 = builder.AddInstruction( HloInstruction::CreateParameter(1, slice_shape, "starts")); HloInstruction* dynamic_slice2 = builder.AddInstruction( - HloInstruction::CreateDynamicSlice(result_shape, param0, param1, {2})); + HloInstruction::CreateDynamicSlice(result_shape, param0, {param1}, {2})); builder.AddInstruction(HloInstruction::CreateUnary( result_shape, HloOpcode::kNegate, dynamic_slice2)); @@ -548,28 +551,36 @@ TEST_F(OpcodeFusionTest, DynamicSliceWithDynamicUpdateSlice) { Shape full_shape = ShapeUtil::MakeShape(F32, {10, 100, 1000}); Shape slice_shape = ShapeUtil::MakeShape(F32, {10, 1, 1000}); + std::vector slice_indices, update_indices; + for (int i = 0; i < 3; ++i) { + slice_indices.push_back( + builder.AddInstruction(HloInstruction::CreateParameter( + 1 + i, ShapeUtil::MakeShape(U32, {}), "slice_indices"))); + update_indices.push_back( + builder.AddInstruction(HloInstruction::CreateParameter( + 5 + i, ShapeUtil::MakeShape(U32, {}), "update_indices"))); + } HloInstruction* slice = builder.AddInstruction(HloInstruction::CreateDynamicSlice( slice_shape, builder.AddInstruction( HloInstruction::CreateParameter(0, full_shape, "slice_from")), - builder.AddInstruction(HloInstruction::CreateParameter( - 1, ShapeUtil::MakeShape(U32, {3}), "slice_indices")), + slice_indices, /*slice_sizes=*/{10, 1, 1000})); builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( full_shape, builder.AddInstruction( - HloInstruction::CreateParameter(2, full_shape, "to_update")), - slice, - builder.AddInstruction(HloInstruction::CreateParameter( - 3, ShapeUtil::MakeShape(U32, {3}), "update_indices")))); + HloInstruction::CreateParameter(4, full_shape, "to_update")), + slice, update_indices)); module->AddEntryComputation(builder.Build()); RunFusionAndCheckOpcodesWereFused( - module.get(), {HloOpcode::kDynamicSlice, HloOpcode::kDynamicUpdateSlice, - HloOpcode::kParameter, HloOpcode::kParameter, - HloOpcode::kParameter, HloOpcode::kParameter}); + module.get(), + {HloOpcode::kDynamicSlice, HloOpcode::kDynamicUpdateSlice, + HloOpcode::kParameter, HloOpcode::kParameter, HloOpcode::kParameter, + HloOpcode::kParameter, HloOpcode::kParameter, HloOpcode::kParameter, + HloOpcode::kParameter, HloOpcode::kParameter}); } TEST_F(OpcodeFusionTest, MessOfFusibleNodes) { @@ -578,49 +589,40 @@ TEST_F(OpcodeFusionTest, MessOfFusibleNodes) { Shape full_shape = ShapeUtil::MakeShape(F32, {4, 100, 10, 100, 50}); - auto loop_idx = builder.AddInstruction(HloInstruction::CreateReshape( - ShapeUtil::MakeShape(S32, {1}), - builder.AddInstruction(HloInstruction::CreateParameter( - 0, ShapeUtil::MakeShape(S32, {}), "param0")))); - + auto loop_idx = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(S32, {}), "param0")); auto param1 = builder.AddInstruction(HloInstruction::CreateParameter( - 1, ShapeUtil::MakeShape(S32, {1}), "param1")); - auto concat = builder.AddInstruction(HloInstruction::CreateConcatenate( - ShapeUtil::MakeShape(S32, {5}), - {loop_idx, param1, param1, param1, param1}, /*dimension=*/0)); + 1, ShapeUtil::MakeShape(S32, {}), "param1")); - auto idx_choice = builder.AddInstruction(HloInstruction::CreateDynamicSlice( - ShapeUtil::MakeShape(S32, {1}), - builder.AddInstruction(HloInstruction::CreateParameter( - 2, ShapeUtil::MakeShape(S32, {4}), "param2")), - loop_idx, - /*slice_sizes=*/{1})); - - PaddingConfig padding_config; - padding_config.add_dimensions()->set_edge_padding_high(4); - auto pad = builder.AddInstruction(HloInstruction::CreatePad( - ShapeUtil::MakeShape(S32, {5}), idx_choice, - builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))), - padding_config)); + auto idx_choice = builder.AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(S32, {}), + builder.AddInstruction(HloInstruction::CreateDynamicSlice( + ShapeUtil::MakeShape(S32, {1}), + builder.AddInstruction(HloInstruction::CreateParameter( + 2, ShapeUtil::MakeShape(S32, {4}), "param2")), + {loop_idx}, + /*slice_sizes=*/{1})))); + auto zero = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); auto slice = builder.AddInstruction(HloInstruction::CreateDynamicSlice( ShapeUtil::MakeShape(F32, {1, 100, 10, 100, 50}), builder.AddInstruction(HloInstruction::CreateParameter( 3, ShapeUtil::MakeShape(F32, {100, 100, 10, 100, 50}), "param3")), - pad, /*slice_sizes=*/{1, 100, 10, 100, 50})); + {idx_choice, zero, zero, zero, zero}, + /*slice_sizes=*/{1, 100, 10, 100, 50})); builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( full_shape, builder.AddInstruction( HloInstruction::CreateParameter(4, full_shape, "param4")), - slice, concat)); + slice, {loop_idx, param1, param1, param1, param1})); module->AddEntryComputation(builder.Build()); RunFusionAndCheckOpcodesWereFused( module.get(), - {HloOpcode::kConcatenate, HloOpcode::kPad, HloOpcode::kDynamicSlice, - HloOpcode::kDynamicSlice, HloOpcode::kDynamicUpdateSlice, + {HloOpcode::kDynamicSlice, HloOpcode::kDynamicSlice, + HloOpcode::kDynamicUpdateSlice, HloOpcode::kReshape, HloOpcode::kParameter, HloOpcode::kParameter, HloOpcode::kParameter, HloOpcode::kParameter, HloOpcode::kParameter, HloOpcode::kParameter}); } @@ -930,9 +932,10 @@ ENTRY main { return result; } -INSTANTIATE_TEST_CASE_P(GatherLoopFusionTestInstantiation, GatherLoopFusionTest, - ::testing::ValuesIn(GetGatherLoopFusionTestSpecs()), - GatherLoopFusionTestSpec::Name); +INSTANTIATE_TEST_SUITE_P(GatherLoopFusionTestInstantiation, + GatherLoopFusionTest, + ::testing::ValuesIn(GetGatherLoopFusionTestSpecs()), + GatherLoopFusionTestSpec::Name); } // namespace } // namespace cpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc index 1cb154e9ca3bbfd9742df96316272e0fc653da36..95b8025f873c56bea063ff258d4abd6614257d85 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc @@ -46,8 +46,7 @@ static bool ShouldMakeAllUsersColMajor(const HloInstruction* instruction) { for (auto* user : instruction->users()) { optional operand_idx = ProfitableToMakeDotOperandColumnMajor(*user); if (!operand_idx || user->operand(*operand_idx) != instruction || - std::count(user->operands().begin(), user->operands().end(), - instruction) != 1) { + absl::c_count(user->operands(), instruction) != 1) { return false; } } @@ -94,60 +93,38 @@ static Shape ColMajorShape(const Shape& old_shape) { return new_shape; } +static bool OperandsAndResultMustHaveRowMajorLayout( + const HloInstruction& instr, + const TargetMachineFeatures& target_machine_features) { + if (instr.opcode() == HloOpcode::kConvolution) { + return PotentiallyImplementedAsEigenConvolution(instr, + target_machine_features); + } else if (instr.opcode() == HloOpcode::kDot) { + return DotOperandsAndResultMustHaveRowMajorLayout(instr, + target_machine_features); + } + return false; +} + Status CpuLayoutAssignment::AddBackendConstraints( LayoutConstraints* constraints) { ShouldMakeOperandColMajorCache cache; const HloComputation* computation = constraints->computation(); for (auto* instruction : computation->instructions()) { - if (instruction->opcode() == HloOpcode::kConvolution && - PotentiallyImplementedAsEigenConvolution(*instruction, - target_machine_features_)) { - const HloInstruction* convolution = instruction; - const HloInstruction* lhs_instruction = convolution->operand(0); - const HloInstruction* rhs_instruction = convolution->operand(1); - - // In order to implement `convolution` with Eigen convolution, the layouts - // of the input, filter, and output need to be row-major. - // - // These constraints are not hard constraints. Ideally, we should decide - // which layouts to choose according to some cost model. - Shape output_shape(RowMajorShape(convolution->shape())); - Shape input_shape(RowMajorShape(lhs_instruction->shape())); - Shape filter_shape(RowMajorShape(rhs_instruction->shape())); - - // Set layouts of the instructions' shapes. - TF_RETURN_IF_ERROR( - constraints->SetOperandLayout(input_shape, convolution, 0)); - TF_RETURN_IF_ERROR( - constraints->SetOperandLayout(filter_shape, convolution, 1)); - TF_RETURN_IF_ERROR( - constraints->SetInstructionLayout(output_shape, convolution)); + if (OperandsAndResultMustHaveRowMajorLayout(*instruction, + target_machine_features_)) { + TF_RETURN_IF_ERROR(constraints->SetInstructionLayout( + RowMajorShape(instruction->shape()), instruction)); + for (int i = 0; i < instruction->operand_count(); i++) { + TF_RETURN_IF_ERROR(constraints->SetOperandLayout( + RowMajorShape(instruction->operand(i)->shape()), instruction, i)); + } } else if (optional op_idx = ShouldMakeOperandColumnMajor(&cache, *instruction)) { const HloInstruction* op = instruction->operand(*op_idx); TF_RETURN_IF_ERROR(constraints->SetOperandLayout( ColMajorShape(op->shape()), instruction, *op_idx)); - } else if (PotentiallyImplementedAsEigenDot(*instruction, - target_machine_features_)) { - const HloInstruction* dot = instruction; - // In order to implement `dot` with Eigen dot, the layouts of the lhs, - // rhs, and output need to be row-major. - // - // These constraints are not hard constraints. Ideally, we should decide - // which layouts to choose according to some cost model. - Shape output_shape(RowMajorShape(dot->shape())); - - const HloInstruction* lhs_instruction = dot->operand(0); - Shape lhs_shape(RowMajorShape(lhs_instruction->shape())); - TF_RETURN_IF_ERROR(constraints->SetOperandLayout(lhs_shape, dot, 0)); - - const HloInstruction* rhs_instruction = dot->operand(1); - Shape rhs_shape(RowMajorShape(rhs_instruction->shape())); - TF_RETURN_IF_ERROR(constraints->SetOperandLayout(rhs_shape, dot, 1)); - - // Set layouts of the instructions' shapes. - TF_RETURN_IF_ERROR(constraints->SetInstructionLayout(output_shape, dot)); } else { for (int64 operand_no = 0; operand_no < instruction->operand_count(); ++operand_no) { diff --git a/tensorflow/compiler/xla/service/cpu/cpu_options.cc b/tensorflow/compiler/xla/service/cpu/cpu_options.cc index 92debb83e33b1400a59e5eef0f90971392ab7b22..ff654c83d61e7cc09ac7839feccaf2bc9cb3c63c 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_options.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_options.cc @@ -23,8 +23,8 @@ namespace { const char* const kXlaOptimizeForSizeCpuOption = "xla_cpu_optimize_for_size"; const char* const kLlvmIrDotTilingFactor = "xla_llvm_dot_tiling_factor"; -const char* const kXlaEnableExperimentalLlvmIrGemm = - "xla_enable_experimental_llvm_ir_gemm"; +const char* const kXlaForceEnableExperimentalLlvmIrGemm = + "xla_force_enable_experimental_llvm_ir_gemm"; const char* const kLlvmIrGemmTileSize = "xla_llvm_ir_gemm_tile_size"; } // namespace @@ -57,10 +57,10 @@ absl::optional LlvmIrGemvTilingFactor(const HloModuleConfig& config) { return absl::nullopt; } -bool EnableExperimentalLlvmIrGemm(const HloModuleConfig& config) { +bool ForceEnableExperimentalLlvmIrGemm(const HloModuleConfig& config) { const auto& extra_options_map = config.debug_options().xla_backend_extra_options(); - return extra_options_map.count(kXlaEnableExperimentalLlvmIrGemm) > 0; + return extra_options_map.count(kXlaForceEnableExperimentalLlvmIrGemm) > 0; } static absl::string_view RemoveSuffix(absl::string_view str, diff --git a/tensorflow/compiler/xla/service/cpu/cpu_options.h b/tensorflow/compiler/xla/service/cpu/cpu_options.h index 47c7eb13b6e4cc05a23f82b8d2a25249f4b82ac0..99e6702d14aed8ffb148adec2bdd02dbc7c3c7e3 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_options.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_options.h @@ -26,7 +26,7 @@ namespace options { bool OptimizeForSizeRequested(const HloModuleConfig& config); bool VectorizedReduceDisabled(const HloModuleConfig& config); -bool EnableExperimentalLlvmIrGemm(const HloModuleConfig& config); +bool ForceEnableExperimentalLlvmIrGemm(const HloModuleConfig& config); absl::optional LlvmIrGemvTilingFactor(const HloModuleConfig& config); absl::optional> LlvmIrGemmTileSize( const HloModuleConfig& config); diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc index a9febe891b5e9d1eb9e6b297952b50d1d26a3396..d8878e622c0500fc5328aa6c295a9e24a3a037f7 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc @@ -84,31 +84,8 @@ extern const char* const kReleaseOutfeedBufferAfterPopulationSymbolName = "__xla_cpu_runtime_ReleaseOutfeedBufferAfterPopulation"; extern const char* const kParallelForkJoinSymbolName = "__xla_cpu_runtime_ParallelForkJoin"; -extern const char* const kKeyValueSortPREDSymbolName = - "__xla_cpu_runtime_KeyValueSortPRED"; -extern const char* const kKeyValueSortS8SymbolName = - "__xla_cpu_runtime_KeyValueSortS8"; -extern const char* const kKeyValueSortU8SymbolName = - "__xla_cpu_runtime_KeyValueSortU8"; -extern const char* const kKeyValueSortS16SymbolName = - "__xla_cpu_runtime_KeyValueSortS16"; -extern const char* const kKeyValueSortU16SymbolName = - "__xla_cpu_runtime_KeyValueSortU16"; -extern const char* const kKeyValueSortF16SymbolName = - "__xla_cpu_runtime_KeyValueSortF16"; -extern const char* const kKeyValueSortS32SymbolName = - "__xla_cpu_runtime_KeyValueSortS32"; -extern const char* const kKeyValueSortU32SymbolName = - "__xla_cpu_runtime_KeyValueSortU32"; -extern const char* const kKeyValueSortF32SymbolName = - "__xla_cpu_runtime_KeyValueSortF32"; -extern const char* const kKeyValueSortS64SymbolName = - "__xla_cpu_runtime_KeyValueSortS64"; -extern const char* const kKeyValueSortU64SymbolName = - "__xla_cpu_runtime_KeyValueSortU64"; -extern const char* const kKeyValueSortF64SymbolName = - "__xla_cpu_runtime_KeyValueSortF64"; - +extern const char* const kKeyValueSortSymbolName = + "__xla_cpu_runtime_KeyValueSort"; extern const char* const kXlaCpuRuntimeSymbolNamePrefix = "__xla_cpu_runtime_"; } // namespace runtime } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime.h b/tensorflow/compiler/xla/service/cpu/cpu_runtime.h index b2e760a224ad8eaa61dae57b0f9cece04a7e54ae..3a2b44d8c1a80128d3577c374e751e73a89e9d59 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime.h @@ -64,18 +64,7 @@ extern const char* const kReleaseInfeedBufferAfterDequeueSymbolName; extern const char* const kAcquireOutfeedBufferForPopulationSymbolName; extern const char* const kReleaseOutfeedBufferAfterPopulationSymbolName; extern const char* const kParallelForkJoinSymbolName; -extern const char* const kKeyValueSortPREDSymbolName; -extern const char* const kKeyValueSortS8SymbolName; -extern const char* const kKeyValueSortU8SymbolName; -extern const char* const kKeyValueSortS16SymbolName; -extern const char* const kKeyValueSortU16SymbolName; -extern const char* const kKeyValueSortF16SymbolName; -extern const char* const kKeyValueSortS32SymbolName; -extern const char* const kKeyValueSortU32SymbolName; -extern const char* const kKeyValueSortF32SymbolName; -extern const char* const kKeyValueSortS64SymbolName; -extern const char* const kKeyValueSortU64SymbolName; -extern const char* const kKeyValueSortF64SymbolName; +extern const char* const kKeyValueSortSymbolName; // All symbol names for XLA CPU runtime functions need to start with this // prefix. diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime_test.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime_test.cc index 1ae3aa57111e3a3b7ac18b4907c5c282edf89b7e..4e8c98678309fa4d573f1aac1290c9afc87643a4 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_test.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime_test.cc @@ -162,11 +162,12 @@ TEST_P(EigenMatMulTest, DoIt) { CheckMatrixMultiply(*a, *b, *c); } -INSTANTIATE_TEST_CASE_P(EigenMatMulTestInstantiaion, EigenMatMulTest, - ::testing::Combine(::testing::ValuesIn(MatMulShapes), - ::testing::Bool(), ::testing::Bool(), - ::testing::Bool()), - EigenMatMulTest::Name); +INSTANTIATE_TEST_SUITE_P(EigenMatMulTestInstantiaion, EigenMatMulTest, + ::testing::Combine(::testing::ValuesIn(MatMulShapes), + ::testing::Bool(), + ::testing::Bool(), + ::testing::Bool()), + EigenMatMulTest::Name); #ifdef INTEL_MKL class MKLMatMulTest : public CpuRuntimeTest, diff --git a/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc b/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc index 3361a5973f5e8c91802b26d68477347b196d3cac..fae9670051a654f38f09856368ffb700b0c7a085 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc @@ -29,7 +29,6 @@ limitations under the License. #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/gtl/cleanup.h" #include "tensorflow/core/platform/logging.h" diff --git a/tensorflow/compiler/xla/service/cpu/disassembler.cc b/tensorflow/compiler/xla/service/cpu/disassembler.cc index 3ae64142cd7e32d3aa8d50870efaf94698c06440..c3c6847b7b77e2fb0470630815de9f5d7a6c5b9c 100644 --- a/tensorflow/compiler/xla/service/cpu/disassembler.cc +++ b/tensorflow/compiler/xla/service/cpu/disassembler.cc @@ -77,17 +77,16 @@ StatusOr Disassembler::DisassembleObjectFile( } // Sort the symbols in increasing address order. - std::sort( - symbols.begin(), symbols.end(), - [](const llvm::object::SymbolRef& a, const llvm::object::SymbolRef& b) { - // getAddress returns a Expected object. Assert there is no error - // before extracting the address. - llvm::Expected a_address_or_error = a.getAddress(); - CHECK(a_address_or_error); - llvm::Expected b_address_or_error = b.getAddress(); - CHECK(b_address_or_error); - return a_address_or_error.get() < b_address_or_error.get(); - }); + absl::c_sort(symbols, [](const llvm::object::SymbolRef& a, + const llvm::object::SymbolRef& b) { + // getAddress returns a Expected object. Assert there is no error + // before extracting the address. + llvm::Expected a_address_or_error = a.getAddress(); + CHECK(a_address_or_error); + llvm::Expected b_address_or_error = b.getAddress(); + CHECK(b_address_or_error); + return a_address_or_error.get() < b_address_or_error.get(); + }); // Construct ArrayRef pointing to section contents. llvm::StringRef section_content_string; diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc index 1525a33af7a490f99762733a677d7913c480200d..0fecbaf391bc3122646af30b508fc1a88b6641e9 100644 --- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc @@ -26,7 +26,10 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" #include "tensorflow/compiler/xla/service/cpu/ir_emission_utils.h" #include "tensorflow/compiler/xla/service/cpu/target_machine_features.h" +#include "tensorflow/compiler/xla/service/cpu/tiled_dot_emitter.h" #include "tensorflow/compiler/xla/service/cpu/vector_support_library.h" +#include "tensorflow/compiler/xla/service/hlo_casting_utils.h" +#include "tensorflow/compiler/xla/service/hlo_instructions.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" @@ -41,931 +44,165 @@ namespace xla { using llvm_ir::SetToFirstInsertPoint; namespace cpu { - namespace { -// Provides tiled access to an in-memory rank 2 array. -class MemoryTile { - public: - // Constructs a MemoryTile that can operate on tiles consisting of - // `tile_size_along_major_dim` vectors from the matrix `matrix`, starting at - // `major_dim_offset` in the major dimension. The tile size along the minor - // dimension is the vector size, and that is implicitly determined by `vsl`. - MemoryTile(VectorSupportLibrary* vsl, llvm::IRBuilder<>* b, - llvm::Value* matrix, int64 matrix_size_along_minor_dim, - llvm::Value* major_dim_offset, int64 tile_size_along_major_dim) - : vsl_(vsl), b_(b) { - pointers_.reserve(tile_size_along_major_dim); - for (int64 i = 0; i < tile_size_along_major_dim; i++) { - llvm::Value* total_offset = - b->CreateMul(b->getInt64(matrix_size_along_minor_dim), - b->CreateAdd(b->getInt64(i), major_dim_offset)); - pointers_.push_back(vsl_->ComputeOffsetPointer(matrix, total_offset)); - } - } - - // Load a tile consisting of `tile_size_along_major_dim` vectors from position - // {major: `major_dim_offset`, minor: `minor_dim_offset`}. - // - // Note: `major_dim_offset` is a parameter to the constructor. - std::vector LoadTile(llvm::Value* minor_dim_offset) const { - std::vector result; - result.reserve(pointers_.size()); - for (const auto& pointer : pointers_) { - result.push_back(vsl_->LoadVector(pointer, minor_dim_offset)); - } - return result; - } - - // Stores `tile` to position {major: `major_dim_offset`, minor: - // `minor_dim_offset`}. - // - // Note: `major_dim_offset` is a parameter to the constructor. - void StoreTile(absl::Span tile, - llvm::Value* minor_dim_offset) const { - CHECK_EQ(tile.size(), pointers_.size()); - for (int64 i = 0; i < pointers_.size(); i++) { - vsl_->StoreVector(tile[i], pointers_[i], minor_dim_offset); - } - } - - // Loads a tile of size [`tile_size_along_major_dim`, - // `tile_size_along_middle_dim`] from position {major: `major_dim_offset`, - // minor: `minor_dim_offset`} and then broadcasts each element into a vector - // of size vsl_.vector_size(). The (i,j)'th element of the return value is - // the (i,j)'th element in the tile broadcasted into an LLVM vector. - // - // Note: `major_dim_offset` is a parameter to the constructor. - std::vector> LoadBroadcastTile( - llvm::Value* minor_dim_offset, int64 tile_size_along_middle_dim) const { - std::vector> result; - result.resize(pointers_.size()); - for (int64 i = 0; i < pointers_.size(); i++) { - for (int64 j = 0; j < tile_size_along_middle_dim; j++) { - result[i].push_back(vsl_->LoadBroadcast( - pointers_[i], b_->CreateAdd(minor_dim_offset, b_->getInt64(j)))); - } - } - return result; - } - - private: - VectorSupportLibrary* vsl_; - llvm::IRBuilder<>* b_; - std::vector pointers_; -}; - -// The base class for the classes representing the GEMV emitter configurations. -// -// The IR emitted (modulo the LLVM values representing the input and output -// buffers) by the row major and column major GEMV emitters should be a function -// of their configuration. This is important because their configuration is -// used as a key to cache the generated IR. -class GemvConfig { - public: - // Mixin for convenience. - template - struct User { - public: - PrimitiveType scalar_type() const { - return derived().config().scalar_type(); - } - int64 tile_rows() const { return derived().config().tile_rows(); } - int64 tile_cols() const { return derived().config().tile_cols(); } - int64 m() const { return derived().config().m(); } - int64 k() const { return derived().config().k(); } - int64 has_addend() const { return derived().config().has_addend(); } - - private: - const T& derived() const { return *static_cast(this); } - }; +// Returns true if we should call into multi-threaded Eigen routines. +bool ShouldUseMultiThreadedEigen(const HloModuleConfig& config) { + return config.debug_options().xla_cpu_multi_thread_eigen(); +} - PrimitiveType scalar_type() const { return scalar_type_; } - int64 tile_rows() const { return tile_rows_; } - int64 tile_cols() const { return tile_cols_; } - int64 m() const { return m_; } - int64 k() const { return k_; } - bool has_addend() const { return has_addend_; } - - string GetCacheKey() const { - return absl::StrCat(name_, "_", PrimitiveType_Name(scalar_type()), "_", - tile_rows(), "_", tile_cols(), "_", m(), "_", k(), - has_addend() ? "_with_addend" : ""); +// Represents a dot operation. We use this in lieu of an `HloInstruction` +// because we want to be able to create this for the "inner" dot operation in a +// batch dot, for which there is no separate HLO instruction. +struct DotInfo { + Shape lhs_shape; + Shape rhs_shape; + Shape result_shape; + DotDimensionNumbers dim_nums; + + DotInfo() = default; + + explicit DotInfo(const HloInstruction& instr) { + CHECK_EQ(instr.opcode(), HloOpcode::kDot); + lhs_shape = instr.operand(0)->shape(); + rhs_shape = instr.operand(1)->shape(); + result_shape = instr.shape(); + dim_nums = instr.dot_dimension_numbers(); } - - protected: - explicit GemvConfig(string name, PrimitiveType scalar_type, int64 tile_rows, - int64 tile_cols, int64 m, int64 k, bool has_addend) - : name_(std::move(name)), - scalar_type_(scalar_type), - tile_rows_(tile_rows), - tile_cols_(tile_cols), - m_(m), - k_(k), - has_addend_(has_addend) {} - - private: - string name_; - PrimitiveType scalar_type_; - int64 tile_rows_; - int64 tile_cols_; - int64 m_; - int64 k_; - bool has_addend_; }; -// Computes a dot product between "[M,K]{0,1} lhs" with a [K,1] vector (the -// layout of the vector does not matter). This implementation uses a tiling -// scheme to improve performance. -// -// We logically separate the LHS matrix into four segments: -// -// +----------------------+---+ -// | | | -// | | | -// | A | B | -// | | | -// | | | -// | | | -// +----------------------+---+ -// | C | D | -// +----------------------+---+ -// -// where A is the largest submatrix of the LHS that can be evenly dividied into -// tiles. For each tile in A, assuming tile_rows_ == tile_cols_ == 4, we have: -// -// +---+---+---+---+ +--+--+--+--+ -// |M00|M10|M20|M30| |V0|V1|V2|V3| -// +---+---+---+---+ +--+--+--+--+ -// |M01|M11|M21|M31| and |V0|V1|V2|V3| -// +---+---+---+---+ +--+--+--+--+ -// |M02|M12|M22|M32| |V0|V1|V2|V3| -// +---+---+---+---+ +--+--+--+--+ -// |M03|M13|M23|M33| |V0|V1|V2|V3| -// +---+---+---+---+ +--+--+--+--+ -// -// (Legend: rows are horizontal and columns are vertical; and each column is one -// llvm::Value of a vector type) -// -// where: -// -// a. The left tile is from the column major left matrix. -// b. The right tile is an elementwise broadcast of a [V0, V1, V2, V3] -// vector loaded from the RHS vector. -// -// As we iterate through the column dimension, we compute the change to the -// result vector by an elementwise multiplication between the two tiles above -// followed by a reduction along the major dimension: -// -// +-----------------------------------+ -// | M00*V0 + M10*V1 + M20*V2 + M30*V3 | -// +-----------------------------------+ -// | M01*V0 + M11*V1 + M21*V2 + M31*V3 | -// Result[R:R+4] += +-----------------------------------+ -// | M02*V0 + M12*V1 + M22*V2 + M32*V3 | -// +-----------------------------------+ -// | M03*V0 + M13*V1 + M23*V2 + M33*V3 | -// +-----------------------------------+ -// -// Where R is the starting row for the tile. -// -// We have an inner epilogue loop to deal with the "C" submatrix and an outer -// epilogue loop to deal with the B,D submarix. -// -// TODO(sanjoy): We should investigate if using gather loads and scatter stores -// can be used here have the same inner loop for both column-major and row-major -// matrix-vector products. -class ColumnMajorMatrixVectorProductEmitter - : public GemvConfig::User { - public: - class Config : public GemvConfig { - public: - explicit Config(PrimitiveType scalar_type, int64 tile_rows, int64 tile_cols, - int64 m, int64 k, bool has_addend) - : GemvConfig(/*name=*/"col_major_gemv", scalar_type, - /*tile_rows=*/tile_rows, /*tile_cols=*/tile_cols, /*m=*/m, - /*k=*/k, /*has_addend=*/has_addend) {} - }; - - ColumnMajorMatrixVectorProductEmitter(const Config& config, llvm::Value* lhs, - llvm::Value* rhs, llvm::Value* addend, - llvm::Value* result, - llvm::IRBuilder<>* b) - : config_(config), - lhs_(lhs), - rhs_(rhs), - addend_(addend), - result_(result), - b_(b), - ksl_(b_), - vsl_(config.scalar_type(), /*vector_size=*/config.tile_rows(), b_, "") { - CHECK(tile_rows() > 0 && IsPowerOfTwo(static_cast(tile_rows()))); - CHECK(!has_addend() || addend != nullptr); - } - - void Emit(); - - const Config& config() const { return config_; } - - private: - void EmitOuterLoopBody(llvm::Value* column, int64 column_count, - bool is_first_column); - - MemoryTile GetLhsMemoryTile(llvm::Value* column_start, int64 column_count) { - return MemoryTile(&vsl_, b_, /*matrix=*/lhs_, - /*matrix_size_along_minor_dim=*/m(), - /*major_dim_offset=*/column_start, - /*tile_size_along_major_dim=*/column_count); - } - - // Load a tile of values from the RHS. For the RHS a "tile" is a contiguous - // sequence of `count` values, each one broadcasted to the vector width. - std::vector LoadRhsTile(llvm::Value* offset, int64 count) { - llvm::Value* base_pointer = vsl_.ComputeOffsetPointer(rhs_, offset); - std::vector result; - result.reserve(count); - for (int64 i = 0; i < count; i++) { - result.push_back(vsl_.LoadBroadcast(base_pointer, i)); - } - return result; - } - - void EmitInnerLoopTiled(MemoryTile* lhs_memory_tile, - const std::vector& rhs_tile, - int64 columns, bool is_first_column); - - void EmitInnerLoopEpilogue(llvm::Value* current_tile_col, int64 columns, - bool is_first_tiled_column); - - Config config_; - llvm::Value* lhs_; - llvm::Value* rhs_; - llvm::Value* addend_; - llvm::Value* result_; - llvm::IRBuilder<>* b_; - KernelSupportLibrary ksl_; - VectorSupportLibrary vsl_; +// Dictates how a dot operation is implemented. +enum class DotImplementationStrategy { + // The dot operation is lowered into LLVM IR that implements a naive nested + // loop that computes the result one element at a time. This is our + // "fallback"; we don't really want this to kick in for any non-trival dot + // operation. + kNaiveLlvmIr, + + // The dot operation is lowered into LLVM IR that implements a tiled + // Matrix*Vector operation. This strategy also allows fusing in a bias add + // into the dot. The matrix can be row major or column major, both are + // supported. + kTiledLlvmIrGemv, + + // The dot operation is lowered into LLVM IR that implemetns a tiled + // Matrix*Matrix operation. No fusions are supported. The two inputs + // and the output have to be row major. + kTiledLlvmIrGemm, + + // The dot operation is lowered into a call into an Eigen routine. No fusions + // are supported today. The two inputs and the output have to be row major. + // However, we do allow transposing either the LHS or the RHS as part of the + // GEMM -- we expose this flexibility as flexibility in the contraction + // dimensions, but we can also see this as flexibility in the input layouts. + kEigen, }; -void ColumnMajorMatrixVectorProductEmitter::EmitOuterLoopBody( - llvm::Value* column, int64 column_count, bool is_first_column) { - MemoryTile lhs_memory_tile = GetLhsMemoryTile(/*column_start=*/column, - /*column_count=*/column_count); - - std::vector rhs_tile = - LoadRhsTile(column, /*count=*/column_count); - EmitInnerLoopTiled(&lhs_memory_tile, rhs_tile, - /*columns=*/column_count, is_first_column); - EmitInnerLoopEpilogue(column, /*columns=*/column_count, is_first_column); -} - -void ColumnMajorMatrixVectorProductEmitter::Emit() { - // See the comment on the class declaration for the algorithm used here. - int64 column_remainder = k() % tile_cols(); - int64 column_limit = k() - column_remainder; - - ksl_.For("dot.outer.tiled", - /*start=*/0, /*end=*/column_limit, /*step=*/tile_cols(), - [&](llvm::Value* column, bool is_first_column) { - EmitOuterLoopBody(column, tile_cols(), is_first_column); - }); - - if (column_remainder != 0) { - EmitOuterLoopBody(b_->getInt64(column_limit), column_remainder, - column_limit == 0); - } -} - -void ColumnMajorMatrixVectorProductEmitter::EmitInnerLoopTiled( - MemoryTile* lhs_memory_tile, const std::vector& rhs_tile, - int64 columns, bool is_first_column) { - int64 row_limit = m() - (m() % tile_rows()); - - ksl_.For( - "dot.inner.tiled", /*start=*/0, /*end=*/row_limit, - /*step=*/tile_rows(), [&](llvm::Value* row) { - std::vector lhs_tile = - lhs_memory_tile->LoadTile(/*minor_dim_offset=*/row); - llvm::Value* accumulator = - is_first_column ? (addend_ ? vsl_.LoadVector(addend_, row) - : vsl_.GetZeroVector()) - : vsl_.LoadVector(result_, row); - for (int i = 0; i < columns; i++) { - accumulator = vsl_.MulAdd(lhs_tile[i], rhs_tile[i], accumulator); - } - vsl_.StoreVector(accumulator, result_, row); - }); -} - -void ColumnMajorMatrixVectorProductEmitter::EmitInnerLoopEpilogue( - llvm::Value* current_tile_col, int64 columns, bool is_first_tiled_column) { - int64 row_start = m() - (m() % tile_rows()); - if (row_start == m()) { - return; - } - - llvm::Value* columns_llvm = b_->getInt64(columns); - - // for (col = current_tile_col; col < (columns + current_tile_col); col++) - // for (row = row_start, row < m_; row++) { - // result[row] += lhs[row, col] * rhs[col] - // // Also take into account that if col is 0 then result[row] is not - // // initialized. - // } - - ksl_.For( - "dot.inner.epilg.outer", /*start=*/current_tile_col, - /*end=*/b_->CreateAdd(columns_llvm, current_tile_col), - /*step=*/1, /*peel_first_iteration=*/false, - [&](llvm::Value* col, llvm::Value* is_first_scalar_col) { - llvm::Value* rhs_element = vsl_.LoadScalar(rhs_, col); - llvm::Value* total_offset = b_->CreateMul(col, b_->getInt64(m())); - llvm::Value* lhs_base_pointer = - vsl_.ComputeOffsetPointer(lhs_, total_offset); - ksl_.For( - "dot.inner.epilg.inner", /*start=*/row_start, /*end=*/m(), - /*step=*/1, [&](llvm::Value* scalar_row) { - llvm::Value* product = vsl_.Mul( - vsl_.LoadScalar(lhs_base_pointer, scalar_row), rhs_element); - llvm::Value* setting_result_first_time = b_->CreateAnd( - is_first_scalar_col, b_->getInt1(is_first_tiled_column)); - ksl_.If( - setting_result_first_time, - /*true_block_generator=*/ - [&]() { - if (addend_) { - vsl_.StoreScalar( - vsl_.Add(vsl_.LoadScalar(addend_, scalar_row), - product), - result_, scalar_row); - } else { - vsl_.StoreScalar(product, result_, scalar_row); - } - }, - /*false_block_generator=*/ - [&]() { - vsl_.StoreScalar( - vsl_.Add(vsl_.LoadScalar(result_, scalar_row), product), - result_, scalar_row); - }); - }); - }); -} +// Returns the implementation strategy for a dot with the configuration +// `dot_info`. +DotImplementationStrategy GetDotImplementationStrategy( + const HloModuleConfig& config, const DotInfo& dot_info, + const TargetMachineFeatures& target_machine_features); -// Computes a dot product between "[M,K]{1,0} lhs" with a [K,1] vector (the -// layout of the vector does not matter). This implementation uses a tiling -// scheme to improve performance. -// -// We logically separate the LHS matrix into four segments: -// -// +----------------------+---+ -// | | | -// | | | -// | A | B | -// | | | -// | | | -// | | | -// +----------------------+---+ -// | C | D | -// +----------------------+---+ -// -// where A is the largest submatrix of the LHS that can be evenly dividied into -// tiles. For each tile in A, assuming tile_rows_ == tile_cols_ == 4, we have: -// -// +---+---+---+---+ -// |M00|M10|M20|M30| -// +---+---+---+---+ +--+--+--+--+ -// |M01|M11|M21|M31| and |V0|V1|V2|V3| -// +---+---+---+---+ +--+--+--+--+ -// |M02|M12|M22|M32| -// +---+---+---+---+ -// |M03|M13|M23|M33| -// +---+---+---+---+ -// -// (Legend: rows are horizontal and columns are vertical; and each row is one -// llvm::Value of a vector type) -// -// where: -// -// a. The left tile is loaded from the row major left matrix. -// b. The right vector is loaded from the RHS vector. -// -// We keep 4 vector accumulators accumulating the following four vector -// expressions as we iterate over the row dimension: -// -// +------+------+------+------+ -// |M0I*V0|M1I*V1|M2I*V2|M3I*V3| for I in [0,4) -// +------+------+------+------+ -// -// In the end we do a horizontal reduction over these 4 vector accumulators to -// get 4 values in the result vector. -// -// We have an inner epilogue loop to deal with the "B" sub-matrix and an outer -// epilogue loop to deal with the C,D submatrix. -class RowMajorMatrixVectorProductEmitter - : public GemvConfig::User { +// Helper class for emitting LLVM IR to perform the dot operation. +class DotOpEmitter { public: - class Config : public GemvConfig { - public: - explicit Config(PrimitiveType scalar_type, int64 tile_rows, int64 tile_cols, - int64 m, int64 k, bool has_addend) - : GemvConfig(/*name=*/"row_major_gemv", scalar_type, - /*tile_rows=*/tile_rows, /*tile_cols=*/tile_cols, /*m=*/m, - /*k=*/k, /*has_addend=*/has_addend) {} - }; - - RowMajorMatrixVectorProductEmitter(const Config& config, llvm::Value* lhs, - llvm::Value* rhs, llvm::Value* addend, - llvm::Value* result, llvm::IRBuilder<>* b) - : config_(config), - lhs_(lhs), - rhs_(rhs), - addend_(addend), - result_(result), - b_(b), - ksl_(b_), - vsl_(scalar_type(), /*vector_size=*/tile_cols(), b_, "") { - CHECK(tile_cols() > 0 && IsPowerOfTwo(static_cast(tile_cols()))); - CHECK(!has_addend() || addend != nullptr); - } - - void Emit(); - - const Config& config() const { return config_; } + explicit DotOpEmitter(DotInfo dot_info, string dot_hlo_name, + const llvm_ir::IrArray& target_array, + const llvm_ir::IrArray& lhs_array, + const llvm_ir::IrArray& rhs_array, + const llvm_ir::IrArray* addend_array, + llvm::Value* executable_run_options_value, + llvm::IRBuilder<>* b, + const HloModuleConfig& hlo_module_config, + const TargetMachineFeatures& target_machine_features); + + // Emits the IR to perform the dot operation. + Status Emit(); private: - MemoryTile GetLhsMemoryTile(llvm::Value* row_start, int64 row_count) { - return MemoryTile(&vsl_, b_, /*matrix=*/lhs_, - /*matrix_size_along_minor_dim=*/k(), - /*major_dim_offset=*/row_start, - /*tile_size_along_major_dim=*/row_count); - } - - void EmitOuterLoopBody(llvm::Value* row, int64 row_count); - - void EmitInnerLoopTiled(MemoryTile* lhs_memory_tile, int64 rows, - std::vector* vector_accumulators); - - void EmitInnerLoopEpilogue(llvm::Value* current_tile_row, int64 rows, - std::vector* scalar_accumulators); - - Config config_; - llvm::Value* lhs_; - llvm::Value* rhs_; - llvm::Value* addend_; - llvm::Value* result_; - llvm::IRBuilder<>* b_; - KernelSupportLibrary ksl_; - VectorSupportLibrary vsl_; -}; - -void RowMajorMatrixVectorProductEmitter::EmitOuterLoopBody(llvm::Value* row, - int64 row_count) { - MemoryTile lhs_memory_tile = GetLhsMemoryTile(/*row_start=*/row, - /*row_count=*/row_count); - std::vector vector_accumulators; - std::vector scalar_accumulators; - for (int i = 0; i < row_count; i++) { - vector_accumulators.emplace_back(&vsl_, vsl_.GetZeroVector()); - scalar_accumulators.emplace_back(&vsl_, vsl_.GetZeroScalar()); - } - EmitInnerLoopTiled(&lhs_memory_tile, /*rows=*/row_count, - &vector_accumulators); - EmitInnerLoopEpilogue(/*current_tile_row=*/row, /*rows=*/row_count, - &scalar_accumulators); - - std::vector accumulator_values; - std::transform( - vector_accumulators.begin(), vector_accumulators.end(), - std::back_inserter(accumulator_values), - [](const VectorVariable& vector_var) { return vector_var.Get(); }); - - std::vector horizontal_sums; - if (row_count == vsl_.vector_size()) { - if (addend_) { - horizontal_sums = vsl_.ComputeHorizontalSums( - std::move(accumulator_values), vsl_.LoadVector(addend_, row)); - } else { - horizontal_sums = - vsl_.ComputeHorizontalSums(std::move(accumulator_values)); - } - } else { - horizontal_sums = vsl_.ComputeHorizontalSums(std::move(accumulator_values)); - } - - for (int i = 0; i < row_count; i++) { - llvm::Value* result_value = - vsl_.Add(horizontal_sums[i], scalar_accumulators[i].Get()); - llvm::Value* offset = b_->CreateAdd(b_->getInt64(i), row); - if (addend_ && row_count != vsl_.vector_size()) { - result_value = vsl_.Add(vsl_.LoadScalar(addend_, offset), result_value); - } - vsl_.StoreScalar(result_value, result_, offset); - } -} + // Emits instructions to perform a scalar dot product (a multiply of the + // LHS and RHS) and store the results in the target. + Status EmitScalarDot(); -void RowMajorMatrixVectorProductEmitter::Emit() { - // See the comment on the class declaration for the algorithm used here. - int64 row_remainder = m() % tile_rows(); - int64 row_limit = m() - row_remainder; + // Emits a call to the CPU runtime to perform the matrix multiply. + Status EmitCallToRuntime(); - ksl_.For("dot.outer.tiled", - /*start=*/0, /*end=*/row_limit, /*step=*/tile_rows(), - [&](llvm::Value* row) { EmitOuterLoopBody(row, tile_rows()); }); - - if (row_remainder != 0) { - EmitOuterLoopBody(b_->getInt64(row_limit), row_remainder); - } -} + // Represents the dimensions of a matrix-matrix multiply operation. + struct MatMultDims { + // The number of rows in the LHS. + int64 m; -void RowMajorMatrixVectorProductEmitter::EmitInnerLoopTiled( - MemoryTile* lhs_memory_tile, int64 rows, - std::vector* vector_accumulators) { - int64 column_limit = k() - (k() % tile_cols()); - - ksl_.For("dot.inner.tiled", /*start=*/0, /*end=*/column_limit, - /*step=*/tile_cols(), [&](llvm::Value* col) { - std::vector lhs_tile = - lhs_memory_tile->LoadTile(/*minor_dim_offset=*/col); - llvm::Value* rhs_value = vsl_.LoadVector(rhs_, col); - for (int i = 0; i < rows; i++) { - llvm::Value* old_sum = (*vector_accumulators)[i].Get(); - (*vector_accumulators)[i].Set( - vsl_.Add(old_sum, vsl_.Mul(rhs_value, lhs_tile[i]))); - } - }); -} + // The number of columns in the LHS, which is also must be equal to the + // number of rows in the RHS. + int64 k; -void RowMajorMatrixVectorProductEmitter::EmitInnerLoopEpilogue( - llvm::Value* current_tile_row, int64 rows, - std::vector* scalar_accumulators) { - int64 column_start = k() - (k() % tile_cols()); - if (column_start == k()) { - return; - } + // The number of columns on the RHS. + int64 n; - for (int r = 0; r < rows; r++) { - llvm::Value* total_offset = b_->CreateMul( - b_->CreateAdd(b_->getInt64(r), current_tile_row), b_->getInt64(k())); - llvm::Value* lhs_base_pointer = - vsl_.ComputeOffsetPointer(lhs_, total_offset); - ksl_.For( - "dot.inner.epilg.inner", /*start=*/column_start, /*end=*/k(), - /*step=*/1, [&](llvm::Value* scalar_col) { - llvm::Value* product = - vsl_.Mul(vsl_.LoadScalar(lhs_base_pointer, scalar_col), - vsl_.LoadScalar(rhs_, scalar_col)); - llvm::Value* old_value = (*scalar_accumulators)[r].Get(); - (*scalar_accumulators)[r].Set(vsl_.Add(old_value, product)); - }); - } -} + // True if the LHS matrix is column major. + bool lhs_column_major; -// This class implements a tiled matrix multiplication algorithm, intended for -// multiplying small matrices that don't need cache tiling. -// -// In the future this can be used as the innermost GEBP loop in a GEMM kernel as -// described in "Goto, Kazushige, and Robert A. Geijn. "Anatomy of -// high-performance matrix multiplication." ACM Transactions on Mathematical -// Software (TOMS) 34.3 (2008): 12.". -// -// This only supports canonical dot operations (i.e. where the lhs contraction -// dimension is 1 and the rhs contraction dimension is 0) over row major -// matrices. -class TiledSmallGemmEmitter { - public: - // Describe the dimensions of the kernel. - class Dimensions { - public: - explicit Dimensions(int64 m, int64 k, int64 n) : m_(m), k_(k), n_(n) {} + // True if the LHS contraction dimension is not 1. + bool lhs_non_canonical; - int64 m() const { return m_; } - int64 k() const { return k_; } - int64 n() const { return n_; } + // True if the RHS matrix is column major. + bool rhs_column_major; - string ToString() const { return absl::StrCat(m(), "x", k(), "x", n()); } + // True if the RHS contraction dimension is not 0. + bool rhs_non_canonical; - private: - const int64 m_; - const int64 k_; - const int64 n_; + // True if the result matrix is column major. + bool target_column_major; }; - // Represents the configuration of the emitter. The LLVM IR emitted by the - // emitter, modulo the LLVM values holding the input and output buffers, must - // be a function of the instance of `Config` passed to it. - // - // `dims` holds the matrix multiplication dimensions. - // - // `max_vectorization_width` is the maximum vector width (i.e. the width of - // the largest vector register we will use). This can be larger than the - // largest vector register supported by the machine -- LLVM will legalize - // these large vector widths into legally sized vectors. - // - // `max_vector_count` is the maximum number of vectors of size - // `max_vectorization_width` that we will attempt to process at once. - // - // `min_vectorization_width` is the smallest vector width the emitter will use - // -- below that it will devolve to using a scalar loop. - // - // The innermost reduction loop executes the matrix multiply in tiles of size - // [`tile_size_m`, `tile_size_k`] from the LHS and [`tile_size_k`, - // ] in the RHS. - class Config { - public: - explicit Config(PrimitiveType scalar_type, Dimensions dims, - int64 max_vectorization_width, int64 max_vector_count, - int64 min_vectorization_width, int64 tile_size_m, - int64 tile_size_k) - : scalar_type_(scalar_type), - dims_(dims), - max_vectorization_width_(max_vectorization_width), - max_vector_count_(max_vector_count), - min_vectorization_width_(min_vectorization_width), - tile_size_m_(tile_size_m), - tile_size_k_(tile_size_k) {} - - string GetCacheKey() const { - return absl::StrCat("gemm_", PrimitiveType_Name(scalar_type()), "_", - dims().ToString(), "_", max_vectorization_width(), - "_", min_vectorization_width(), "_", tile_size_m(), - "_", tile_size_k()); - } + // Get the MatMultDims instance for the dot product this DotOpEmitter + // represents. Precondition: the dot is of rank 2 (and thus its operands are + // of rank 2 as well). + MatMultDims GetMatMultDims() const; - PrimitiveType scalar_type() const { return scalar_type_; } - Dimensions dims() const { return dims_; } - int64 max_vectorization_width() const { return max_vectorization_width_; } - int64 max_vector_count() const { return max_vector_count_; } - int64 min_vectorization_width() const { return min_vectorization_width_; } - - int64 tile_size_m() const { return tile_size_m_; } - int64 tile_size_k() const { return tile_size_k_; } - - private: - PrimitiveType scalar_type_; - Dimensions dims_; - int64 max_vectorization_width_; - int64 max_vector_count_; - int64 min_vectorization_width_; - int64 tile_size_m_; - int64 tile_size_k_; - }; + // Lowers the dot operation as a tiled Matrix*Vector loop. + void EmitTiledLlvmIrGemv(); - // Creates an instance of TiledSmallGemmEmitter that matrix-multiplies - // `lhs` with `rhs` and stores the result in `result`. - explicit TiledSmallGemmEmitter(Config config, llvm::Value* lhs, - llvm::Value* rhs, llvm::Value* result, - llvm::IRBuilder<>* b) - : lhs_(lhs), - rhs_(rhs), - result_(result), - config_(config), - b_(b), - ksl_(b_) { - CHECK(max_vectorization_width() > 0 && - IsPowerOfTwo(static_cast(max_vectorization_width()))); - CHECK_GT(max_vector_count(), 0); - CHECK(min_vectorization_width() > 0 && - IsPowerOfTwo(static_cast(min_vectorization_width()))); - CHECK_GE(max_vectorization_width(), min_vectorization_width()); - CHECK_GT(tile_size_k(), 0); - } + // Lowers the dot operation as a tiled Matrix*Matrix loop. + void EmitTiledLlvmIrGemm(); - void Emit(); + // Lowers the dot operation as a naive nested loop that computes the result + // one element at a time. + void EmitNaiveLlvmIrGemm(); - private: - // The HandleResiduesOnX helpers split the iteration space for dimension X - // into a multiple of the tile size on dimension X and an epilogue. These - // helpers ultimately call into `EmitTiledGemm` for emitting the - // tiled GEMM kernel. - - void HandleResiduesOnN(); - void HandleResiduesOnK(VectorSupportLibrary* vsl, llvm::Value* n_start, - llvm::Value* n_end); - void HandleResiduesOnM(VectorSupportLibrary* vsl, int64 tile_size_k, - llvm::Value* k_start, llvm::Value* k_end, - llvm::Value* n_start, llvm::Value* n_end); - - // This emits a tiled GEMM kernel. For a detailed description see the comment - // on the implementation. - void EmitTiledGemm(VectorSupportLibrary* vsl, int64 tile_size_k, - llvm::Value* k_start, llvm::Value* k_end, - llvm::Value* n_start, llvm::Value* n_end, - int64 tile_size_m, llvm::Value* m_start, - llvm::Value* m_end); - - llvm::Value* GetInt64(int64 value) { return b_->getInt64(value); } - - Config config() const { return config_; } - Dimensions dims() const { return config().dims(); } - - int64 max_vectorization_width() const { - return config().max_vectorization_width(); + // When doing a tiled GEMV in LLVM IR, a "tile" consists of this many vector + // registers. + int64 GetGemvTilingFactor() const { + const int64 kDefaultTilingFactor = 8; + return options::LlvmIrGemvTilingFactor(hlo_module_config_) + .value_or(kDefaultTilingFactor); } - int64 max_vector_count() const { return config().max_vector_count(); } - int64 min_vectorization_width() const { - return config().min_vectorization_width(); - } - int64 tile_size_m() const { return config().tile_size_m(); } - int64 tile_size_k() const { return config().tile_size_k(); } - PrimitiveType scalar_type() const { return config().scalar_type(); } - llvm::Value* lhs_; - llvm::Value* rhs_; - llvm::Value* result_; - Config config_; + std::tuple GetGemmTileSize() const { + // Tuned for broadwell - Intel(R) Xeon(R) CPU E5-2690 v4 @ 2.60GHz + // + // TODO(b/80093688): Tune for other architectures and centralize this + // information in one place. + const std::tuple kDefaultTileSize = + std::tuple(11, 9, 1); + return options::LlvmIrGemmTileSize(hlo_module_config_) + .value_or(kDefaultTileSize); + } + DotInfo dot_info_; + string dot_hlo_name_; + const llvm_ir::IrArray& target_array_; + const llvm_ir::IrArray& lhs_array_; + const llvm_ir::IrArray& rhs_array_; + const llvm_ir::IrArray* addend_array_; + llvm::Value* executable_run_options_value_; llvm::IRBuilder<>* b_; - KernelSupportLibrary ksl_; + const HloModuleConfig& hlo_module_config_; + const TargetMachineFeatures& target_machine_features_; }; - -void TiledSmallGemmEmitter::Emit() { HandleResiduesOnN(); } - -void TiledSmallGemmEmitter::HandleResiduesOnN() { - // We can only iterate the `n` dimension for an extent that is divisible by - // the vectorization width. So we emit an outer loop that first processes the - // largest extent in `n` that is divisible by max_vectorization_width, then - // the largest remaining extent that is divisible by max_vectorization_width / - // 2 etc. - - int64 current_vectorization_width = - max_vector_count() * max_vectorization_width(); - int64 current_vector_count = max_vector_count(); - - int64 n_start = 0; - while (n_start != dims().n() && - current_vectorization_width >= min_vectorization_width()) { - int64 n_end = dims().n() - (dims().n() % current_vectorization_width); - if (n_start != n_end) { - VectorSupportLibrary vsl(scalar_type(), current_vectorization_width, b_, - "gemm"); - HandleResiduesOnK(&vsl, GetInt64(n_start), GetInt64(n_end)); - n_start = n_end; - } - if (current_vector_count == 1) { - current_vectorization_width /= 2; - } else { - current_vector_count--; - current_vectorization_width = - current_vector_count * max_vectorization_width(); - } - } - - if (n_start != dims().n()) { - VectorSupportLibrary vsl(scalar_type(), 1, b_, "gemm"); - ksl_.For("epi.n", n_start, dims().n(), 1, [&](llvm::Value* n_i) { - llvm::Value* n_i_next = b_->CreateAdd(n_i, b_->getInt64(1)); - HandleResiduesOnK(&vsl, n_i, n_i_next); - }); - } -} - -void TiledSmallGemmEmitter::HandleResiduesOnK(VectorSupportLibrary* vsl, - llvm::Value* n_start, - llvm::Value* n_end) { - int64 k_start = 0; - int64 k_end = dims().k() - (dims().k() % tile_size_k()); - if (k_end != k_start) { - HandleResiduesOnM(vsl, tile_size_k(), GetInt64(k_start), GetInt64(k_end), - n_start, n_end); - k_start = k_end; - } - - if (k_start != dims().k()) { - HandleResiduesOnM(vsl, dims().k() - k_start, GetInt64(k_start), - GetInt64(dims().k()), n_start, n_end); - } -} - -void TiledSmallGemmEmitter::HandleResiduesOnM( - VectorSupportLibrary* vsl, int64 tile_size_k, llvm::Value* k_start, - llvm::Value* k_end, llvm::Value* n_start, llvm::Value* n_end) { - const int64 m_end = dims().m() - dims().m() % tile_size_m(); - EmitTiledGemm(vsl, tile_size_k, k_start, k_end, n_start, n_end, tile_size_m(), - GetInt64(0), GetInt64(m_end)); - - if (m_end != dims().m()) { - EmitTiledGemm(vsl, tile_size_k, k_start, k_end, n_start, n_end, - dims().m() - m_end, GetInt64(m_end), GetInt64(dims().m())); - } -} - -// The loop structure is: -// -// Iterate over dimension M as m: -// Iterate over dimension N as n: -// Iterate over dimension K as k: -// OutputTile[m,n] += Dot(LhsTile[m,k], RhsTile[k,n]) -// -// I.e. a just a tiled version of a "naive" GEMM. -// -// The tiling scheme is as follows: -// -// Let the LHS be: -// -// +----+----+----+ -// | a0 | b0 | c0 | . -// +----+----+----+ . -// | a1 | b1 | c1 | . -// +----+----+----+ -// .. .. -// -// and the RHS be: -// -// +----+----+----+----+ -// | p0 | p1 | p2 | p3 | . -// +----+----+----+----+ . -// | q0 | q1 | q2 | q3 | . -// +----+----+----+----+ -// | r0 | r1 | r2 | r3 | . -// +----+----+----+----+ . -// ...... ...... -// -// and let tile_size_m=2, tile_size_k=3 and the vector width (implicitly denoted -// by `vsl`) be 4. Then we want to matrix multiply this tile to get a [2,4] -// matrix that we can increment the result matrix by. -// -// First broadcast the rows row in LHS to 3 vectors of width 4, giving us a rank -// 3 array, L, of dimension [2,3,4]: -// -// L[0,_,_] * L[1,_,_] -// * -// +----+----+----+----+ * +----+----+----+----+ -// | a0 | a0 | a0 | a0 | * | a1 | a1 | a1 | a1 | -// +----+----+----+----+ * +----+----+----+----+ -// | b0 | b0 | b0 | b0 | * | b1 | b1 | b1 | b1 | -// +----+----+----+----+ * +----+----+----+----+ -// | c0 | c0 | c0 | c0 | * | c1 | c1 | c1 | c1 | -// +----+----+----+----+ * +----+----+----+----+ -// -// -// Then we FMA L[0,_,_] with the RHS to get the first row of the result and -// L[1,_,_] with the RHS to get the second row of the result. For example, -// L[0,_,_] is computed as: -// -// +----+----+----+----+ +----+----+----+----+ -// | a0 | a0 | a0 | a0 | * | p0 | p1 | p2 | p3 | + -// +----+----+----+----+ +----+----+----+----+ -// -// +----+----+----+----+ +----+----+----+----+ -// | b0 | b0 | b0 | b0 | * | q0 | q1 | q2 | q3 | + -// +----+----+----+----+ +----+----+----+----+ -// -// +----+----+----+----+ +----+----+----+----+ -// | c0 | c0 | c0 | c0 | * | r0 | r1 | r2 | r3 | -// +----+----+----+----+ +----+----+----+----+ -// -// to get: -// -// +-------------------+-------------------+-------------------+--------- -// | a0*p0+b0*q0+c0*r0 | a0*p1+b0*q1+c0*r1 | a0*p2+b0*q2+c0*r2 | ... -// +-------------------+-------------------+-------------------+--------- -void TiledSmallGemmEmitter::EmitTiledGemm( - VectorSupportLibrary* vsl, int64 tile_size_k, llvm::Value* k_start, - llvm::Value* k_end, llvm::Value* n_start, llvm::Value* n_end, - int64 tile_size_m, llvm::Value* m_start, llvm::Value* m_end) { - ksl_.For( - "dot.m", m_start, m_end, tile_size_m, [&](llvm::Value* m_i) { - MemoryTile result_memory_tile( - vsl, b_, /*matrix=*/result_, - /*matrix_size_along_minor_dim=*/dims().n(), - /*major_dim_offset=*/m_i, - /*tile_size_along_major_dim=*/tile_size_m); - MemoryTile lhs_memory_tile(vsl, b_, /*matrix=*/lhs_, - /*matrix_size_along_minor_dim=*/dims().k(), - /*major_dim_offset=*/m_i, - /*tile_size_along_major_dim=*/tile_size_m); - ksl_.For( - "dot.n", n_start, n_end, vsl->vector_size(), [&](llvm::Value* n_i) { - TileVariable result_tile_var(vsl, - result_memory_tile.LoadTile(n_i)); - ksl_.For( - "dot.k", k_start, k_end, tile_size_k, [&](llvm::Value* k_i) { - MemoryTile rhs_memory_tile(vsl, b_, rhs_, dims().n(), k_i, - tile_size_k); - std::vector> lhs_tile = - lhs_memory_tile.LoadBroadcastTile(k_i, tile_size_k); - std::vector rhs_tile = - rhs_memory_tile.LoadTile(n_i); - std::vector result_tile = - result_tile_var.Get(); - for (int64 r_m_i = 0; r_m_i < tile_size_m; r_m_i++) { - for (int64 r_k_i = 0; r_k_i < tile_size_k; r_k_i++) { - result_tile[r_m_i] = - vsl->MulAdd(lhs_tile[r_m_i][r_k_i], rhs_tile[r_k_i], - result_tile[r_m_i]); - } - } - result_tile_var.Set(result_tile); - }); - - result_memory_tile.StoreTile(result_tile_var.Get(), n_i); - }); - }); -} - } // namespace -DotOpEmitter::DotOpEmitter(const HloInstruction& dot, +DotOpEmitter::DotOpEmitter(DotInfo dot_info, string dot_hlo_name, const llvm_ir::IrArray& target_array, const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array, @@ -974,7 +211,8 @@ DotOpEmitter::DotOpEmitter(const HloInstruction& dot, llvm::IRBuilder<>* b, const HloModuleConfig& hlo_module_config, const TargetMachineFeatures& target_machine_features) - : dot_(dot), + : dot_info_(std::move(dot_info)), + dot_hlo_name_(std::move(dot_hlo_name)), target_array_(target_array), lhs_array_(lhs_array), rhs_array_(rhs_array), @@ -984,58 +222,9 @@ DotOpEmitter::DotOpEmitter(const HloInstruction& dot, hlo_module_config_(hlo_module_config), target_machine_features_(target_machine_features) {} -/* static */ Status DotOpEmitter::EmitDotOperation( - const HloInstruction& dot, const llvm_ir::IrArray& target_array, - const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array, - const llvm_ir::IrArray* addend_array, - llvm::Value* executable_run_options_value, llvm::IRBuilder<>* b, - const HloModuleConfig& hlo_module_config, - const TargetMachineFeatures& target_machine_features) { - PrimitiveType type = target_array.GetShape().element_type(); - TF_RET_CHECK(F16 == type || F32 == type || F64 == type || C64 == type); - DotOpEmitter dot_emitter(dot, target_array, lhs_array, rhs_array, - addend_array, executable_run_options_value, b, - hlo_module_config, target_machine_features); - return dot_emitter.Emit(); -} - -bool DotOpEmitter::EmitSmallGemmIfProfitable( - const DotOpEmitter::MatMultDims& mat_mult_dims) { - if (ShouldUseMultiThreadedEigen()) { - return false; - } - - if (!EnableExperimentalLlvmIrGemm()) { - // TODO(sanjoy): We should make these numbers micro-arch specific. - bool small_gemm = mat_mult_dims.k <= 128 && - ((mat_mult_dims.m <= 32 && mat_mult_dims.n <= 128) || - (mat_mult_dims.m <= 128 && mat_mult_dims.n <= 32)); - if (!small_gemm) { - return false; - } - } - - if (mat_mult_dims.lhs_non_canonical || mat_mult_dims.rhs_non_canonical) { - return false; - } - - PrimitiveType primitive_type = dot_.shape().element_type(); - - switch (primitive_type) { - default: - return false; - - case F32: - case F64: - case S32: - case S64: - break; - } - - if (!(mat_mult_dims.lhs_column_major == mat_mult_dims.rhs_column_major && - mat_mult_dims.rhs_column_major == mat_mult_dims.target_column_major)) { - return false; - } +void DotOpEmitter::EmitTiledLlvmIrGemm() { + PrimitiveType primitive_type = dot_info_.result_shape.element_type(); + MatMultDims mat_mult_dims = GetMatMultDims(); llvm::Value* lhs = lhs_array_.GetBasePointer(); llvm::Value* rhs = rhs_array_.GetBasePointer(); @@ -1050,9 +239,8 @@ bool DotOpEmitter::EmitSmallGemmIfProfitable( } int64 size_bytes = m * n * ShapeUtil::ByteSizeOfPrimitiveType(primitive_type); - b_->CreateMemSet( - target, b_->getInt8(0), size_bytes, - target_machine_features_.minimum_alignment_for_allocation(size_bytes)); + b_->CreateMemSet(target, b_->getInt8(0), /*Size=*/size_bytes, + /*Align=*/1); int64 max_target_vector_width = target_machine_features_.vector_register_num_elements( @@ -1062,47 +250,28 @@ bool DotOpEmitter::EmitSmallGemmIfProfitable( std::tie(tile_size_m, tile_size_k, tile_size_n_in_vector_width) = GetGemmTileSize(); - TiledSmallGemmEmitter::Config config( - /*scalar_type=*/primitive_type, - TiledSmallGemmEmitter::Dimensions{/*m=*/m, /*k=*/k, /*n=*/n}, - /*max_vectorization_width=*/max_target_vector_width, - /*max_vector_count=*/tile_size_n_in_vector_width, - /*min_vectorization_width=*/std::min(4, max_target_vector_width), - /*tile_size_m=*/tile_size_m, /*tile_size_k=*/tile_size_k); - - VLOG(2) << "Emitting GEMM kernel in LLVM IR with config " - << config.GetCacheKey(); - const bool enable_fast_math = hlo_module_config_.debug_options().xla_cpu_enable_fast_math(); const bool optimize_for_size = options::OptimizeForSizeRequested(hlo_module_config_); - KernelSupportLibrary::EmitAndCallOutlinedKernel( + EmitSmallGemm( + /*scalar_type=*/primitive_type, + /*m=*/m, /*k=*/k, /*n=*/n, + /*max_vectorization_width=*/max_target_vector_width, + /*max_vector_count=*/tile_size_n_in_vector_width, + /*min_vectorization_width=*/std::min(4, max_target_vector_width), + /*tile_size_m=*/tile_size_m, /*tile_size_k=*/tile_size_k, /*lhs=*/lhs, + /*rhs=*/rhs, /*result=*/target, b_, /*enable_fast_math=*/enable_fast_math, - /*optimize_for_size=*/optimize_for_size, b_, config.GetCacheKey(), lhs, - rhs, target, - [this, config](llvm::Value* lhs, llvm::Value* rhs, llvm::Value* target) { - TiledSmallGemmEmitter small_gemm_emitter(config, /*lhs=*/lhs, - /*rhs=*/rhs, - /*result=*/target, b_); - small_gemm_emitter.Emit(); - }); - - return true; + /*optimize_for_size=*/optimize_for_size); } -bool DotOpEmitter::EmitLlvmIrDotIfProfitable() { - if (dot_.shape().dimensions_size() != 2) { - return false; - } - - PrimitiveType primitive_type = dot_.shape().element_type(); +void DotOpEmitter::EmitTiledLlvmIrGemv() { + PrimitiveType primitive_type = dot_info_.result_shape.element_type(); - if (!primitive_util::IsFloatingPointType(primitive_type) && - !primitive_util::IsIntegralType(primitive_type)) { - return false; - } + CHECK(primitive_util::IsFloatingPointType(primitive_type) || + primitive_util::IsIntegralType(primitive_type)); MatMultDims mat_mult_dims = GetMatMultDims(); bool is_column_major_matrix_vector = false; @@ -1143,9 +312,7 @@ bool DotOpEmitter::EmitLlvmIrDotIfProfitable() { } } - if (!is_column_major_matrix_vector && !is_row_major_matrix_vector) { - return EmitSmallGemmIfProfitable(mat_mult_dims); - } + CHECK(is_column_major_matrix_vector || is_row_major_matrix_vector); int64 tiling_factor = GetGemvTilingFactor(); CHECK_GT(tiling_factor, 0); @@ -1177,44 +344,27 @@ bool DotOpEmitter::EmitLlvmIrDotIfProfitable() { if (is_column_major_matrix_vector) { VLOG(2) << "Emitting column major matrix-vector multiply with m = " << m << " and k = " << k; - ColumnMajorMatrixVectorProductEmitter::Config config( + EmitColumnMajorGemv( /*scalar_type=*/primitive_type, /*tile_rows=*/vector_register_element_size, /*tile_cols=*/tiling_factor, - /*m=*/m, /*k=*/k, /*has_addend=*/addend_array_ != nullptr); - - KernelSupportLibrary::EmitAndCallOutlinedKernel( + /*m=*/m, /*k=*/k, /*lhs=*/lhs_op, /*rhs=*/rhs_op, + /*addend=*/addend_array_ ? addend_array_->GetBasePointer() : nullptr, + /*result=*/result_op, b_, /*enable_fast_math=*/enable_fast_math, - /*optimize_for_size=*/optimize_for_size, b_, config.GetCacheKey(), - lhs_op, rhs_op, - addend_array_ ? addend_array_->GetBasePointer() : nullptr, result_op, - [this, config](llvm::Value* lhs_op, llvm::Value* rhs_op, - llvm::Value* addend_op, llvm::Value* result_op) { - ColumnMajorMatrixVectorProductEmitter emitter( - config, lhs_op, rhs_op, addend_op, result_op, b_); - emitter.Emit(); - }); + /*optimize_for_size=*/optimize_for_size); } else { VLOG(2) << "Emitting row major matrix-vector multiply with m = " << m << " and k = " << k; - RowMajorMatrixVectorProductEmitter::Config config( + EmitRowMajorGemv( /*scalar_type=*/primitive_type, - /*tile_rows=*/tiling_factor, /*tile_cols=*/vector_register_element_size, - /*m=*/m, /*k=*/k, /*has_addend=*/addend_array_ != nullptr); - - KernelSupportLibrary::EmitAndCallOutlinedKernel( + /*tile_rows=*/tiling_factor, + /*tile_cols=*/vector_register_element_size, + /*m=*/m, /*k=*/k, /*lhs=*/lhs_op, /*rhs=*/rhs_op, + /*addend=*/addend_array_ ? addend_array_->GetBasePointer() : nullptr, + /*result=*/result_op, b_, /*enable_fast_math=*/enable_fast_math, - /*optimize_for_size=*/optimize_for_size, b_, config.GetCacheKey(), - lhs_op, rhs_op, - addend_array_ ? addend_array_->GetBasePointer() : nullptr, result_op, - [this, config](llvm::Value* lhs_op, llvm::Value* rhs_op, - llvm::Value* addend_op, llvm::Value* result_op) { - RowMajorMatrixVectorProductEmitter emitter(config, lhs_op, rhs_op, - addend_op, result_op, b_); - emitter.Emit(); - }); + /*optimize_for_size=*/optimize_for_size); } - - return true; } Status DotOpEmitter::Emit() { @@ -1240,11 +390,6 @@ Status DotOpEmitter::Emit() { // which performs the sum-of-products (the reduction loop) before storing // the result in the output buffer. - // This routine assumes that the dot operation is not in a parallelized - // enclosing computation. - CHECK( - dot_.parent()->root_instruction()->outer_dimension_partitions().empty()); - const Shape& lhs_shape = lhs_array_.GetShape(); const Shape& rhs_shape = rhs_array_.GetShape(); @@ -1255,27 +400,41 @@ Status DotOpEmitter::Emit() { return EmitScalarDot(); } - if (EmitLlvmIrDotIfProfitable()) { - return Status::OK(); + switch (GetDotImplementationStrategy(hlo_module_config_, dot_info_, + target_machine_features_)) { + case DotImplementationStrategy::kNaiveLlvmIr: + EmitNaiveLlvmIrGemm(); + return Status::OK(); + + case DotImplementationStrategy::kTiledLlvmIrGemv: + EmitTiledLlvmIrGemv(); + return Status::OK(); + + case DotImplementationStrategy::kTiledLlvmIrGemm: + EmitTiledLlvmIrGemm(); + return Status::OK(); + + case DotImplementationStrategy::kEigen: + return EmitCallToRuntime(); } +} +void DotOpEmitter::EmitNaiveLlvmIrGemm() { CHECK_EQ(addend_array_, nullptr); - if (PotentiallyImplementedAsEigenDot(dot_, target_machine_features_)) { - return EmitCallToRuntime(); - } + const Shape& lhs_shape = lhs_array_.GetShape(); + const Shape& rhs_shape = rhs_array_.GetShape(); + const DotDimensionNumbers& dim_nums = dot_info_.dim_nums; // Reduce along dimension 0 of the LHS and 1 of the RHS. Vectors are a special // case where the reduction dimension is 0 for both LHS and RHS. This results // in a vector dot product producing a scalar. - int64 lhs_reduction_dimension = - dot_.dot_dimension_numbers().lhs_contracting_dimensions(0); - int64 rhs_reduction_dimension = - dot_.dot_dimension_numbers().rhs_contracting_dimensions(0); + int64 lhs_reduction_dimension = dim_nums.lhs_contracting_dimensions(0); + int64 rhs_reduction_dimension = dim_nums.rhs_contracting_dimensions(0); // Verify the reduction dimension in the two operands are the same size. - TF_RET_CHECK(lhs_shape.dimensions(lhs_reduction_dimension) == - rhs_shape.dimensions(rhs_reduction_dimension)); + CHECK_EQ(lhs_shape.dimensions(lhs_reduction_dimension), + rhs_shape.dimensions(rhs_reduction_dimension)); bool lhs_reduction_along_minor_dimension = lhs_reduction_dimension == LayoutUtil::Minor(lhs_shape.layout(), 0); @@ -1285,7 +444,7 @@ Status DotOpEmitter::Emit() { // Create loop nests which loop through the LHS operand dimensions and the RHS // operand dimensions. The reduction dimension of the LHS and RHS are handled // in a separate innermost loop which performs the sum of products. - llvm_ir::ForLoopNest loop_nest(llvm_ir::IrName(&dot_), b_); + llvm_ir::ForLoopNest loop_nest(llvm_ir::IrName(dot_hlo_name_), b_); llvm_ir::IrArray::Index lhs_index = loop_nest.EmitOperandArrayLoopNest( lhs_array_, /*dimension_to_skip=*/lhs_reduction_dimension, "lhs"); llvm_ir::IrArray::Index rhs_index = loop_nest.EmitOperandArrayLoopNest( @@ -1390,8 +549,6 @@ Status DotOpEmitter::Emit() { // Set the IR builder insert point to the exit basic block of the outer most // loop. b_->SetInsertPoint(loop_nest.GetOuterLoopExitBasicBlock()); - - return Status::OK(); } Status DotOpEmitter::EmitScalarDot() { @@ -1438,7 +595,7 @@ Status DotOpEmitter::EmitCallToRuntime() { // The two transpose_... parameters are actually booleans, but we use int32 // to avoid target-dependent calling convention details. - bool multi_threaded = ShouldUseMultiThreadedEigen(); + bool multi_threaded = ShouldUseMultiThreadedEigen(hlo_module_config_); bool use_mkl_dnn = hlo_module_config_.debug_options().xla_cpu_use_mkl_dnn(); PrimitiveType type = target_array_.GetShape().element_type(); llvm::Type* float_type; @@ -1486,11 +643,13 @@ Status DotOpEmitter::EmitCallToRuntime() { llvm::Function* function = b_->GetInsertBlock()->getParent(); llvm::Module* module = function->getParent(); - llvm::Function* matmul_func = llvm::cast( - module->getOrInsertFunction(fn_name, matmul_type)); - matmul_func->setCallingConv(llvm::CallingConv::C); - matmul_func->setDoesNotThrow(); - matmul_func->setOnlyAccessesArgMemory(); + llvm::FunctionCallee matmul_func = + module->getOrInsertFunction(fn_name, matmul_type); + if (auto* fn = llvm::dyn_cast(matmul_func.getCallee())) { + fn->setCallingConv(llvm::CallingConv::C); + fn->setDoesNotThrow(); + fn->setOnlyAccessesArgMemory(); + } // The Eigen runtime function expects column-major layout. If the matrices are // row major, then use the following identity to compute the product: @@ -1531,11 +690,11 @@ Status DotOpEmitter::EmitCallToRuntime() { } DotOpEmitter::MatMultDims DotOpEmitter::GetMatMultDims() const { - CHECK_EQ(dot_.shape().dimensions_size(), 2); + CHECK_EQ(dot_info_.result_shape.dimensions_size(), 2); const Shape& lhs_shape = lhs_array_.GetShape(); const Shape& rhs_shape = rhs_array_.GetShape(); - const DotDimensionNumbers& dim_nums = dot_.dot_dimension_numbers(); + const DotDimensionNumbers& dim_nums = dot_info_.dim_nums; return { /*m=*/lhs_shape.dimensions(1 - dim_nums.lhs_contracting_dimensions(0)), @@ -1549,74 +708,6 @@ DotOpEmitter::MatMultDims DotOpEmitter::GetMatMultDims() const { LayoutUtil::Minor(target_array_.GetShape().layout(), 0) == 0}; } -// Return whether the given shape is rank 2. -static bool IsRank2(const Shape& shape) { return shape.rank() == 2; } - -// In a gemm operation where output = lhs * rhs, check whether the given shapes -// are valid for the operation. -static bool AreValidGemmShapes( - const Shape& lhs_shape, const Shape& rhs_shape, const Shape& output_shape, - const TargetMachineFeatures& target_machine_features) { - // The inputs and the output must - // 1) be matrices with no padding, and - // 2) have an allowed element type. - PrimitiveType output_primitive_type = output_shape.element_type(); - if (!(output_primitive_type == F64 || output_primitive_type == F32 || - output_primitive_type == F16)) { - return false; - } - - if (!(IsRank2(lhs_shape) && IsRank2(rhs_shape) && IsRank2(output_shape))) { - return false; - } - - auto is_aligned = [&](const Shape& shape) { - return GetMinimumAlignmentForArray(shape, target_machine_features) >= - TargetMachineFeatures::kEigenExpectedTensorAlignment; - }; - - if (!is_aligned(lhs_shape) || !is_aligned(rhs_shape) || - !is_aligned(output_shape)) { - return false; - } - - return true; -} - -bool PotentiallyImplementedAsEigenDot( - const HloInstruction& hlo, - const TargetMachineFeatures& target_machine_features) { - // For certain types of Dot, we can call Eigen - if (hlo.opcode() == HloOpcode::kDot) { - const Shape& lhs_shape = hlo.operand(0)->shape(); - const Shape& rhs_shape = hlo.operand(1)->shape(); - - if (ShapeUtil::IsZeroElementArray(lhs_shape) || - ShapeUtil::IsZeroElementArray(rhs_shape)) { - return false; - } - - if (ProfitableToImplementDotInTiledLlvmIr(hlo)) { - return false; - } - - // If gemm can accept the operand shapes, use it rather than a custom - // kernel. - if (AreValidGemmShapes(lhs_shape, rhs_shape, hlo.shape(), - target_machine_features)) { - const DotDimensionNumbers& dim_numbers = hlo.dot_dimension_numbers(); - // The size of the reduction dimension should match. The shape inference - // guarantees this invariant, so the check here is for programming - // errors. - CHECK_EQ(lhs_shape.dimensions(dim_numbers.lhs_contracting_dimensions(0)), - rhs_shape.dimensions(dim_numbers.rhs_contracting_dimensions(0))); - return true; - } - } - - return false; -} - // For vector-matrix dot products, it is always profitable to make the Rhs // column major. absl::optional ProfitableToMakeDotOperandColumnMajor( @@ -1655,16 +746,319 @@ absl::optional ProfitableToMakeDotOperandColumnMajor( return {}; } -bool ProfitableToImplementDotInTiledLlvmIr(const HloInstruction& dot) { +namespace { +// Return whether the given shape is rank 2. +bool IsRank2(const Shape& shape) { return shape.rank() == 2; } + +bool IsSimpleLayout(const Layout& layout) { + return layout.tiles().empty() && layout.format() == DENSE; +} + +// In a gemm operation where output = lhs * rhs, check whether the given shapes +// are valid for the operation. +bool AreGemmShapes(const Shape& lhs_shape, const Shape& rhs_shape, + const Shape& output_shape, + const TargetMachineFeatures& target_machine_features) { + CHECK(!lhs_shape.has_layout() || IsSimpleLayout(lhs_shape.layout())) + << lhs_shape.DebugString(); + CHECK(!rhs_shape.has_layout() || IsSimpleLayout(rhs_shape.layout())) + << rhs_shape.DebugString(); + CHECK(!output_shape.has_layout() || IsSimpleLayout(output_shape.layout())) + << output_shape.DebugString(); + + switch (output_shape.element_type()) { + case F64: + case F32: + case F16: + return IsRank2(lhs_shape) && IsRank2(rhs_shape) && IsRank2(output_shape); + default: + return false; + } +} + +bool IsAlignedGemm(const DotInfo& dot_info, + const TargetMachineFeatures& target_machine_features) { + if (ShapeUtil::IsZeroElementArray(dot_info.lhs_shape) || + ShapeUtil::IsZeroElementArray(dot_info.rhs_shape)) { + return false; + } + + return AreGemmShapes(dot_info.lhs_shape, dot_info.rhs_shape, + dot_info.result_shape, target_machine_features); +} + +bool CanEmitTiledLlvmIrGemm( + const HloModuleConfig& config, const DotInfo& dot_info, + const TargetMachineFeatures& target_machine_features) { + CHECK(IsAlignedGemm(dot_info, target_machine_features)); + + if (ShouldUseMultiThreadedEigen(config)) { + return false; + } + + int m = dot_info.result_shape.dimensions(0); + int k = dot_info.lhs_shape.dimensions( + dot_info.dim_nums.lhs_contracting_dimensions(0)); + int n = dot_info.result_shape.dimensions(1); + + if (!options::ForceEnableExperimentalLlvmIrGemm(config)) { + // TODO(sanjoy): We should make these numbers micro-arch specific. + bool small_gemm = + k <= 128 && ((m <= 32 && n <= 128) || (m <= 128 && n <= 32)); + if (!small_gemm) { + return false; + } + } + + bool lhs_non_canonical = dot_info.dim_nums.lhs_contracting_dimensions(0) == 0; + bool rhs_non_canonical = dot_info.dim_nums.rhs_contracting_dimensions(0) == 1; + + if (lhs_non_canonical || rhs_non_canonical) { + return false; + } + + if (dot_info.result_shape.element_type() == F16) { + // TODO(sanjoy): This is probably easy to fix, but I want to keep the CL + // adding this comment NFC. + return false; + } + + return true; +} + +DotImplementationStrategy GetDotImplementationStrategy( + const HloModuleConfig& config, const DotInfo& dot_info, + const TargetMachineFeatures& target_machine_features) { + PrimitiveType element_type = dot_info.result_shape.element_type(); // Any Matrix-Vector product of floating point or integral type, or // a transpose-dot fusion of the same can be lowered to a tiled LLVM // IR implementation. - const Shape& shape = dot.shape(); - return shape.dimensions_size() == 2 && - (shape.dimensions(0) == 1 || shape.dimensions(1) == 1) && - (primitive_util::IsFloatingPointType(shape.element_type()) || - primitive_util::IsIntegralType(shape.element_type())); + if (dot_info.result_shape.dimensions_size() == 2 && + (dot_info.result_shape.dimensions(0) == 1 || + dot_info.result_shape.dimensions(1) == 1) && + (primitive_util::IsFloatingPointType(element_type) || + primitive_util::IsIntegralType(element_type))) { + return DotImplementationStrategy::kTiledLlvmIrGemv; + } + + if (IsAlignedGemm(dot_info, target_machine_features)) { + return CanEmitTiledLlvmIrGemm(config, dot_info, target_machine_features) + ? DotImplementationStrategy::kTiledLlvmIrGemm + : DotImplementationStrategy::kEigen; + } + + return DotImplementationStrategy::kNaiveLlvmIr; +} + +Status EmitNonBatchDotOperation( + DotInfo dot_info, string hlo_name, const llvm_ir::IrArray& target_array, + const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array, + const llvm_ir::IrArray* addend_array, + llvm::Value* executable_run_options_value, llvm::IRBuilder<>* b, + const HloModuleConfig& hlo_module_config, + const TargetMachineFeatures& target_machine_features) { + PrimitiveType type = target_array.GetShape().element_type(); + TF_RET_CHECK(F16 == type || F32 == type || F64 == type || C64 == type || + C128 == type); + DotOpEmitter dot_emitter(std::move(dot_info), std::move(hlo_name), + target_array, lhs_array, rhs_array, addend_array, + executable_run_options_value, b, hlo_module_config, + target_machine_features); + return dot_emitter.Emit(); +} + +Shape DropFirstDim(const Shape& shape) { + absl::Span array_shape_dims(shape.dimensions()); + array_shape_dims.remove_prefix(1); + return ShapeUtil::MakeShapeWithDescendingLayout(shape.element_type(), + array_shape_dims); +} + +Shape CollapseFirstNDims(const Shape& shape, int64 n) { + absl::Span input_shape_dims(shape.dimensions()); + int64 prefix_dim = + std::accumulate(input_shape_dims.begin(), input_shape_dims.begin() + n, + 1ll, std::multiplies()); + DimensionVector result_dims; + result_dims.push_back(prefix_dim); + std::copy(input_shape_dims.begin() + n, input_shape_dims.end(), + std::back_inserter(result_dims)); + return ShapeUtil::MakeShapeWithDescendingLayout(shape.element_type(), + result_dims); +} + +llvm_ir::IrArray CollapseFirstNDims(llvm::IRBuilder<>* b, + const llvm_ir::IrArray& array, int64 n) { + llvm::Module* module = b->GetInsertBlock()->getParent()->getParent(); + const Shape& shape = array.GetShape(); + CHECK(shape.has_layout() && + LayoutUtil::IsMonotonicWithDim0Major(shape.layout())); + CHECK_GE(shape.dimensions_size(), n); + Shape new_shape = CollapseFirstNDims(shape, n); + llvm::Value* new_value = b->CreateBitCast( + array.GetBasePointer(), + llvm_ir::ShapeToIrType(new_shape, module)->getPointerTo()); + return llvm_ir::IrArray(new_value, std::move(new_shape)); +} + +Status ValidateDotDimensionNumbers(const DotDimensionNumbers& dim_numbers) { + // Checks some invariants that do not hold in general, but DotDecomposer + // should have established for us. This is just a debugging aid. + TF_RET_CHECK(dim_numbers.lhs_contracting_dimensions_size() == 1); + std::vector batch_dim_numbers(dim_numbers.lhs_batch_dimensions_size()); + absl::c_iota(batch_dim_numbers, 0); + TF_RET_CHECK( + absl::c_equal(batch_dim_numbers, dim_numbers.lhs_batch_dimensions())); + TF_RET_CHECK( + absl::c_equal(batch_dim_numbers, dim_numbers.rhs_batch_dimensions())); + return Status::OK(); +} + +// Slice out the inner array at batch index `batch_index` from `outer_array`. +llvm_ir::IrArray SliceOutInnerArray(llvm_ir::IrArray outer_array, + llvm::Value* batch_index, + llvm::IRBuilder<>* b) { + llvm::Module* module = b->GetInsertBlock()->getParent()->getParent(); + + Shape inner_shape = DropFirstDim(outer_array.GetShape()); + llvm_ir::IrArray::Index slice_index(b->getInt64Ty()); + slice_index.push_back(batch_index); + slice_index.InsertAt( + /*index=*/1, outer_array.GetShape().dimensions_size() - 1, + b->getInt64(0)); + llvm::Value* slice_ptr = outer_array.EmitArrayElementAddress(slice_index, b); + llvm::Type* slice_ptr_type = + llvm_ir::ShapeToIrType(inner_shape, module)->getPointerTo(); + return llvm_ir::IrArray(b->CreateBitCast(slice_ptr, slice_ptr_type), + std::move(inner_shape)); +} + +Status EmitBatchDotOperation( + const HloInstruction& dot, const llvm_ir::IrArray& target_array, + const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array, + llvm::Value* executable_run_options_value, llvm::IRBuilder<>* b, + const HloModuleConfig& hlo_module_config, + const TargetMachineFeatures& target_machine_features) { + TF_RETURN_IF_ERROR(ValidateDotDimensionNumbers(dot.dot_dimension_numbers())); + + // Lower a batch dot into a sequence of non-batch dot operations. + + int64 num_batch_dims = + dot.dot_dimension_numbers().lhs_batch_dimensions_size(); + + // First reshape the inputs to make sure we only have one batch dimension. + // This is a no-op bitcast because the operands have to be in row-major layout + // (enforced in CpuLayoutAssignment), and the batch dimensions are the leading + // dimensions (established by DotDecomposer and checked by + // ValidateDotDimensionNumbers above). + llvm_ir::IrArray lhs_array_reshaped = + CollapseFirstNDims(b, lhs_array, num_batch_dims); + llvm_ir::IrArray rhs_array_reshaped = + CollapseFirstNDims(b, rhs_array, num_batch_dims); + llvm_ir::IrArray target_array_reshaped = + CollapseFirstNDims(b, target_array, num_batch_dims); + + int64 batch_count = lhs_array_reshaped.GetShape().dimensions(0); + + KernelSupportLibrary ksl(b); + + return ksl.ForWithStatus( + "bdot", /*start=*/0, /*end=*/batch_count, /*step=*/1, + [&](llvm::Value* indvar) { + DotDimensionNumbers adjusted_dim_numbers = dot.dot_dimension_numbers(); + adjusted_dim_numbers.clear_lhs_batch_dimensions(); + adjusted_dim_numbers.clear_rhs_batch_dimensions(); + + // Create a DotInfo representing the "inner" non-batch dot operation. + DotInfo dot_info; + dot_info.lhs_shape = DropFirstDim(lhs_array_reshaped.GetShape()); + dot_info.rhs_shape = DropFirstDim(rhs_array_reshaped.GetShape()); + dot_info.result_shape = DropFirstDim(target_array_reshaped.GetShape()); + dot_info.dim_nums = dot.dot_dimension_numbers(); + dot_info.dim_nums.clear_lhs_batch_dimensions(); + dot_info.dim_nums.clear_rhs_batch_dimensions(); + + dot_info.dim_nums.set_lhs_contracting_dimensions( + 0, + dot_info.dim_nums.lhs_contracting_dimensions(0) - num_batch_dims); + dot_info.dim_nums.set_rhs_contracting_dimensions( + 0, + dot_info.dim_nums.rhs_contracting_dimensions(0) - num_batch_dims); + + llvm_ir::IrArray lhs_slice = + SliceOutInnerArray(lhs_array_reshaped, /*batch_index=*/indvar, b); + llvm_ir::IrArray rhs_slice = + SliceOutInnerArray(rhs_array_reshaped, /*batch_index=*/indvar, b); + llvm_ir::IrArray target_slice = SliceOutInnerArray( + target_array_reshaped, /*batch_index=*/indvar, b); + + // Emit the inner non-batch dot operation. + return EmitNonBatchDotOperation( + dot_info, dot.name(), target_slice, lhs_slice, rhs_slice, nullptr, + executable_run_options_value, b, hlo_module_config, + target_machine_features); + }); +} + +bool IsBatchDot(const HloInstruction& instr) { + if (auto* dot_instr = DynCast(&instr)) { + return dot_instr->dot_dimension_numbers().lhs_batch_dimensions_size() > 0; + } + + return false; +} +} // namespace + +bool DotImplementationCanHandleTranspose( + const HloInstruction& dot_instr, + const TargetMachineFeatures& target_machine_features) { + DotImplementationStrategy impl_strategy = + GetDotImplementationStrategy(dot_instr.parent()->parent()->config(), + DotInfo(dot_instr), target_machine_features); + + // TODO(sanjoy): This is not quite right, it should be `impl_strategy == + // kEigen || impl_strategy == kTiledLlvmIrGemv || impl_strategy == + // kNaiveLlvmIr` but I'll fix this in a later CL in the interest of keeping + // the CL adding this comment NFC. + return impl_strategy == DotImplementationStrategy::kTiledLlvmIrGemm || + impl_strategy == DotImplementationStrategy::kEigen; } +bool DotOperandsAndResultMustHaveRowMajorLayout( + const HloInstruction& dot_instr, + const TargetMachineFeatures& target_machine_features) { + DotImplementationStrategy impl_strategy = + GetDotImplementationStrategy(dot_instr.parent()->parent()->config(), + DotInfo(dot_instr), target_machine_features); + + return impl_strategy == DotImplementationStrategy::kTiledLlvmIrGemm || + impl_strategy == DotImplementationStrategy::kEigen; +} + +Status EmitDotOperation(const HloInstruction& dot, + const llvm_ir::IrArray& target_array, + const llvm_ir::IrArray& lhs_array, + const llvm_ir::IrArray& rhs_array, + const llvm_ir::IrArray* addend_array, + llvm::Value* executable_run_options_value, + llvm::IRBuilder<>* b, + const HloModuleConfig& hlo_module_config, + const TargetMachineFeatures& target_machine_features) { + // This routine assumes that the dot operation is not in a parallelized + // enclosing computation. + CHECK(dot.parent()->root_instruction()->outer_dimension_partitions().empty()); + + if (IsBatchDot(dot)) { + TF_RET_CHECK(addend_array == nullptr); + return EmitBatchDotOperation(dot, target_array, lhs_array, rhs_array, + executable_run_options_value, b, + hlo_module_config, target_machine_features); + } + + return EmitNonBatchDotOperation(DotInfo(dot), dot.name(), target_array, + lhs_array, rhs_array, addend_array, + executable_run_options_value, b, + hlo_module_config, target_machine_features); +} } // namespace cpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h index 4c2041b556aa8bf8fe8fb8e0674c0f4f04f0acae..105bd3005c86d87443b2528eba7b0106ad70590e 100644 --- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h +++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h @@ -30,9 +30,16 @@ limitations under the License. namespace xla { namespace cpu { +// Returns true if the two operands and the output of `dot_instr` must have row +// major layout. +bool DotOperandsAndResultMustHaveRowMajorLayout( + const HloInstruction& dot_instr, + const TargetMachineFeatures& target_machine_features); -bool PotentiallyImplementedAsEigenDot( - const HloInstruction& hlo, +// Returns true our lowering strategy for `dot_instr` can fold in transposes to +// the either of the inputs. +bool DotImplementationCanHandleTranspose( + const HloInstruction& dot_instr, const TargetMachineFeatures& target_machine_features); // Returns the index for an operand to `hlo` that should ideally be column @@ -41,129 +48,24 @@ bool PotentiallyImplementedAsEigenDot( absl::optional ProfitableToMakeDotOperandColumnMajor( const HloInstruction& hlo); -// Returns true to indicate that we can generate a tiled LLVM IR implementation -// for |dot|. -bool ProfitableToImplementDotInTiledLlvmIr(const HloInstruction& dot); - -// Helper class for emitting LLVM IR to perform the dot operation. -class DotOpEmitter { - public: - // Emit LLVM IR to perform the dot operation on lhs_array and rhs_array and - // place the result in target_array. IR is emitted at current insert point of - // the builder. Upon completion of the method, the insert point is set to the - // end of all instructions emitted for this operation. - // - // If `addend_array` is not nullptr then it must be an array of the same - // dimensions as the result, and the result is computed as `addend_array` + - // dot(`lhs_array`, `rhs_array`). A non-null `addend_array` is only supported - // for Matrix-vector products. - static Status EmitDotOperation( - const HloInstruction& dot, const llvm_ir::IrArray& target_array, - const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array, - const llvm_ir::IrArray* addend_array, - llvm::Value* executable_run_options_value, llvm::IRBuilder<>* b, - const HloModuleConfig& hlo_module_config, - const TargetMachineFeatures& target_machine_features); - - private: - DotOpEmitter(const HloInstruction& dot, const llvm_ir::IrArray& target_array, - const llvm_ir::IrArray& lhs_array, - const llvm_ir::IrArray& rhs_array, - const llvm_ir::IrArray* addend_array, - llvm::Value* executable_run_options_value, llvm::IRBuilder<>* b, - const HloModuleConfig& hlo_module_config, - const TargetMachineFeatures& target_machine_features); - - // Emits the IR to perform the dot operation. - Status Emit(); - - // Emits instructions to perform a scalar dot product (a multiply of the - // LHS and RHS) and store the results in the target. - Status EmitScalarDot(); - - // Emit an LLVM IR implementation of the dot operation if we can. Returns - // true if an LLVM IR implementation was emitted. - bool EmitLlvmIrDotIfProfitable(); - - // Emits a call to the CPU runtime to perform the matrix multiply. - Status EmitCallToRuntime(); - - // Represents the dimensions of a matrix-matrix multiply operation. - struct MatMultDims { - // The number of rows in the LHS. - int64 m; - - // The number of columns in the LHS, which is also must be equal to the - // number of rows in the RHS. - int64 k; - - // The number of columns on the RHS. - int64 n; - - // True if the LHS matrix is column major. - bool lhs_column_major; - - // True if the LHS contraction dimension is not 1. - bool lhs_non_canonical; - - // True if the RHS matrix is column major. - bool rhs_column_major; - - // True if the RHS contraction dimension is not 0. - bool rhs_non_canonical; - - // True if the result matrix is column major. - bool target_column_major; - }; - - // Get the MatMultDims instance for the dot product this DotOpEmitter - // represents. Precondition: the dot is of rank 2 (and thus its operands are - // of rank 2 as well). - MatMultDims GetMatMultDims() const; - - bool EmitSmallGemmIfProfitable(const MatMultDims& mat_mult_dims); - - // When doing a tiled GEMV in LLVM IR, a "tile" consists of this many vector - // registers. - int64 GetGemvTilingFactor() const { - const int64 kDefaultTilingFactor = 8; - return options::LlvmIrGemvTilingFactor(hlo_module_config_) - .value_or(kDefaultTilingFactor); - } - - std::tuple GetGemmTileSize() const { - // Tuned for broadwell - Intel(R) Xeon(R) CPU E5-2690 v4 @ 2.60GHz - // - // TODO(b/80093688): Tune for other architectures and centralize this - // information in one place. - const std::tuple kDefaultTileSize = - std::tuple(11, 9, 1); - return options::LlvmIrGemmTileSize(hlo_module_config_) - .value_or(kDefaultTileSize); - } - - // Returns true if we should use an experimental implementation of GEMM - // (general matrix matrix multiplication) if possible. - bool EnableExperimentalLlvmIrGemm() const { - return options::EnableExperimentalLlvmIrGemm(hlo_module_config_); - } - - // Returns true if we should call into multi-threaded Eigen routines. - bool ShouldUseMultiThreadedEigen() { - return hlo_module_config_.debug_options().xla_cpu_multi_thread_eigen(); - } - - const HloInstruction& dot_; - const llvm_ir::IrArray& target_array_; - const llvm_ir::IrArray& lhs_array_; - const llvm_ir::IrArray& rhs_array_; - const llvm_ir::IrArray* addend_array_; - llvm::Value* executable_run_options_value_; - llvm::IRBuilder<>* b_; - const HloModuleConfig& hlo_module_config_; - const TargetMachineFeatures& target_machine_features_; -}; - +// Emit LLVM IR to perform the dot operation on lhs_array and rhs_array and +// place the result in target_array. IR is emitted at current insert point of +// the builder. Upon completion of the method, the insert point is set to the +// end of all instructions emitted for this operation. +// +// If `addend_array` is not nullptr then it must be an array of the same +// dimensions as the result, and the result is computed as `addend_array` + +// dot(`lhs_array`, `rhs_array`). A non-null `addend_array` is only supported +// for Matrix-vector products. +Status EmitDotOperation(const HloInstruction& dot, + const llvm_ir::IrArray& target_array, + const llvm_ir::IrArray& lhs_array, + const llvm_ir::IrArray& rhs_array, + const llvm_ir::IrArray* addend_array, + llvm::Value* executable_run_options_value, + llvm::IRBuilder<>* b, + const HloModuleConfig& hlo_module_config, + const TargetMachineFeatures& target_machine_features); } // namespace cpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter_internal.h b/tensorflow/compiler/xla/service/cpu/dot_op_emitter_internal.h new file mode 100644 index 0000000000000000000000000000000000000000..cc28918ed60a8086135846e2b9b1b9d75ec31ef6 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter_internal.h @@ -0,0 +1,88 @@ +/* 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_CPU_DOT_OP_EMITTER_INTERNAL_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_DOT_OP_EMITTER_INTERNAL_H_ + +#include "tensorflow/compiler/xla/service/cpu/target_machine_features.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" + +// ----------------------------------------------------------------------------- +// INTERNAL HEADER. +// +// This file exposes internal implementation details from dot_op_emitter.cc for +// unit tests. Please do not depend on this! +// +// ----------------------------------------------------------------------------- + +namespace xla { +namespace cpu { +namespace internal { + +// Represents a dot operation. We use this in lieu of an `HloInstruction` +// because we want to be able to create this for the "inner" dot operation in a +// batch dot, for which there is no separate HLO instruction. +struct DotInfo { + Shape lhs_shape; + Shape rhs_shape; + Shape result_shape; + DotDimensionNumbers dim_nums; + + explicit DotInfo(const HloInstruction& instr) { + CHECK_EQ(instr.opcode(), HloOpcode::kDot); + lhs_shape = instr.operand(0)->shape(); + rhs_shape = instr.operand(1)->shape(); + result_shape = instr.shape(); + dim_nums = instr.dot_dimension_numbers(); + } +}; + +// Dictates how a dot operation is implemented. +enum class DotImplementationStrategy { + // The dot operation is lowered into LLVM IR that implements a naive nested + // loop that computes the result one element at a time. This is our + // "fallback"; we don't really want this to kick in for any non-trival dot + // operation. + kNaiveLlvmIr, + + // The dot operation is lowered into LLVM IR that implements a tiled + // Matrix*Vector operation. This strategy also allows fusing in a bias add + // into the dot. The matrix can be row major or column major, both are + // supported. + kTiledLlvmIrGemv, + + // The dot operation is lowered into LLVM IR that implemetns a tiled + // Matrix*Matrix operation. No fusions are supported. The two inputs + // and the output have to be row major. + kTiledLlvmIrGemm, + + // The dot operation is lowered into a call into an Eigen routine. No fusions + // are supported today. The two inputs and the output have to be row major. + // However, we do allow transposing either the LHS or the RHS as part of the + // GEMM -- we expose this flexibility as flexibility in the contraction + // dimensions, but we can also see this as flexibility in the input layouts. + kEigen, +}; + +// Returns the implementation strategy for a dot with the configuration +// `dot_info`. +DotImplementationStrategy GetDotImplementationStrategy( + const HloModuleConfig& config, const DotInfo& dot_info, + const TargetMachineFeatures& target_machine_features); +} // namespace internal +} // namespace cpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_DOT_OP_EMITTER_INTERNAL_H_ diff --git a/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc index c8312d80bd5012e5bcb42a410db18a7fa77a2eb6..0028fbaed895becad8da496aa8acdf7dc173a2a0 100644 --- a/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc @@ -51,10 +51,11 @@ StatusOr CpuElementalIrEmitter::EmitAtan2(PrimitiveType prim_type, return Unimplemented("atan2"); } // Create a function declaration. - llvm::Function* function = - llvm::cast(module_->getOrInsertFunction( - llvm_ir::AsStringRef(function_name), lhs->getType(), lhs->getType(), - rhs->getType())); + llvm::Function* function = llvm::dyn_cast( + module_ + ->getOrInsertFunction(llvm_ir::AsStringRef(function_name), + lhs->getType(), lhs->getType(), rhs->getType()) + .getCallee()); function->setCallingConv(llvm::CallingConv::C); function->setDoesNotThrow(); function->setDoesNotAccessMemory(); @@ -85,9 +86,11 @@ StatusOr CpuElementalIrEmitter::EmitTanh(PrimitiveType prim_type, return Unimplemented("tanh"); } // Create a function declaration. - llvm::Function* function = llvm::cast( - module_->getOrInsertFunction(llvm_ir::AsStringRef(function_name), - value->getType(), value->getType())); + llvm::Function* function = llvm::dyn_cast( + module_ + ->getOrInsertFunction(llvm_ir::AsStringRef(function_name), + value->getType(), value->getType()) + .getCallee()); function->setCallingConv(llvm::CallingConv::C); function->setDoesNotThrow(); function->setDoesNotAccessMemory(); diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc index 169d628923b05c43590782bcbebb0e5e539d83e3..5abb3eb38725a3d8e2b761abff1b66f35e92c130 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc @@ -24,11 +24,9 @@ limitations under the License. #include #include +// IWYU pragma: no_include "llvm/IR/Intrinsics.gen.inc" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" -#include "tensorflow/core/lib/math/math_util.h" -#include "tensorflow/core/platform/logging.h" -// IWYU pragma: no_include "llvm/IR/Intrinsics.gen.inc" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/types/span.h" @@ -70,6 +68,8 @@ limitations under the License. #include "tensorflow/compiler/xla/window_util.h" #include "tensorflow/core/lib/core/bits.h" #include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/math/math_util.h" +#include "tensorflow/core/platform/logging.h" namespace xla { @@ -77,7 +77,6 @@ namespace { using llvm_ir::AsStringRef; using llvm_ir::IrName; using llvm_ir::SetToFirstInsertPoint; -namespace gtl = tensorflow::gtl; } // namespace namespace cpu { @@ -239,10 +238,12 @@ Status IrEmitter::HandleCopy(HloInstruction* copy) { int IrEmitter::MinimumAlignmentForPrimitiveType(PrimitiveType primitive_type) { int64 byte_size = ShapeUtil::ByteSizeOfPrimitiveType(primitive_type); DCHECK_GE(byte_size, 0); - // Largest scalar is a complex64 so we don't need to worry about the + // Largest scalar is a complex128 so we don't need to worry about the // int64->int truncation here. - DCHECK_LE(byte_size, 8); - return byte_size; + DCHECK_LE(byte_size, 16); + + // Allocations may be 8-byte aligned if part of a small block. + return std::min(8LL, byte_size); } int64 IrEmitter::ByteSizeOf(const Shape& shape) const { @@ -410,11 +411,18 @@ Status IrEmitter::EmitXfeedTransfer(XfeedKind kind, const Shape& shape, llvm::Function* acquire_func; if (kind == XfeedKind::kInfeed) { - acquire_func = llvm::cast(module_->getOrInsertFunction( - runtime::kAcquireInfeedBufferForDequeueSymbolName, acquire_type)); + acquire_func = llvm::dyn_cast( + module_ + ->getOrInsertFunction( + runtime::kAcquireInfeedBufferForDequeueSymbolName, acquire_type) + .getCallee()); } else { - acquire_func = llvm::cast(module_->getOrInsertFunction( - runtime::kAcquireOutfeedBufferForPopulationSymbolName, acquire_type)); + acquire_func = llvm::dyn_cast( + module_ + ->getOrInsertFunction( + runtime::kAcquireOutfeedBufferForPopulationSymbolName, + acquire_type) + .getCallee()); } acquire_func->setCallingConv(llvm::CallingConv::C); @@ -427,11 +435,19 @@ Status IrEmitter::EmitXfeedTransfer(XfeedKind kind, const Shape& shape, llvm::Function* release_func; if (kind == XfeedKind::kInfeed) { - release_func = llvm::cast(module_->getOrInsertFunction( - runtime::kReleaseInfeedBufferAfterDequeueSymbolName, release_type)); + release_func = llvm::dyn_cast( + module_ + ->getOrInsertFunction( + runtime::kReleaseInfeedBufferAfterDequeueSymbolName, + release_type) + .getCallee()); } else { - release_func = llvm::cast(module_->getOrInsertFunction( - runtime::kReleaseOutfeedBufferAfterPopulationSymbolName, release_type)); + release_func = llvm::dyn_cast( + module_ + ->getOrInsertFunction( + runtime::kReleaseOutfeedBufferAfterPopulationSymbolName, + release_type) + .getCallee()); } release_func->setCallingConv(llvm::CallingConv::C); @@ -493,6 +509,27 @@ Status IrEmitter::HandleSort(HloInstruction* hlo) { const HloSortInstruction* sort = Cast(hlo); TF_RETURN_IF_ERROR(EmitTargetAddressForOp(sort)); Shape keys_shape = sort->keys()->shape(); + PrimitiveType keys_type = keys_shape.element_type(); + switch (keys_type) { + case PRED: + case S8: + case U8: + case S16: + case U16: + case BF16: + case F16: + case S32: + case U32: + case F32: + case S64: + case U64: + case F64: + break; + default: + return Unimplemented( + "Element type %s not supported in the Sort op on CPU.", + PrimitiveType_Name(keys_type)); + } std::vector destination_addresses(sort->operand_count()); for (int64 i = 0; i < sort->operand_count(); ++i) { ShapeIndex shape_index = @@ -540,105 +577,47 @@ Status IrEmitter::HandleSort(HloInstruction* hlo) { lower_dimensions *= normalized_keys_shape.dimensions(i); } - PrimitiveType keys_type = keys_shape.element_type(); - const char* fn_name = nullptr; - llvm::Type* keys_native_type = nullptr; - switch (keys_type) { - case PRED: - fn_name = runtime::kKeyValueSortPREDSymbolName; - keys_native_type = b_.getInt8PtrTy(); - break; - case S8: - fn_name = runtime::kKeyValueSortS8SymbolName; - keys_native_type = b_.getInt8PtrTy(); - break; - case U8: - fn_name = runtime::kKeyValueSortU8SymbolName; - keys_native_type = b_.getInt8PtrTy(); - break; - case S16: - fn_name = runtime::kKeyValueSortS16SymbolName; - keys_native_type = b_.getInt16Ty()->getPointerTo(); - break; - case U16: - fn_name = runtime::kKeyValueSortU16SymbolName; - keys_native_type = b_.getInt16Ty()->getPointerTo(); - break; - case F16: - fn_name = runtime::kKeyValueSortF16SymbolName; - keys_native_type = b_.getHalfTy()->getPointerTo(); - break; - case S32: - fn_name = runtime::kKeyValueSortS32SymbolName; - keys_native_type = b_.getInt32Ty()->getPointerTo(); - break; - case U32: - fn_name = runtime::kKeyValueSortU32SymbolName; - keys_native_type = b_.getInt32Ty()->getPointerTo(); - break; - case F32: - fn_name = runtime::kKeyValueSortF32SymbolName; - keys_native_type = b_.getFloatTy()->getPointerTo(); - break; - case S64: - fn_name = runtime::kKeyValueSortS64SymbolName; - keys_native_type = b_.getInt64Ty()->getPointerTo(); - break; - case U64: - fn_name = runtime::kKeyValueSortU64SymbolName; - keys_native_type = b_.getInt64Ty()->getPointerTo(); - break; - case F64: - fn_name = runtime::kKeyValueSortF64SymbolName; - keys_native_type = b_.getDoubleTy()->getPointerTo(); - break; - default: - return Unimplemented( - "Element type %s not supported in the Sort op on CPU.", - PrimitiveType_Name(keys_type)); - } - + auto less_than_function = FindOrDie(emitted_functions_, sort->to_apply()); + CHECK(absl::c_binary_search(thread_local_computations_, sort->to_apply())); llvm::FunctionType* key_value_sort_type = llvm::FunctionType::get( b_.getVoidTy(), - {keys_native_type, b_.getInt64Ty(), b_.getInt64Ty(), b_.getInt64Ty(), + {b_.getInt64Ty(), b_.getInt64Ty(), b_.getInt64Ty(), b_.getInt8PtrTy()->getPointerTo(), b_.getInt32Ty(), - b_.getInt32Ty()->getPointerTo()}, + b_.getInt32Ty()->getPointerTo(), b_.getInt8PtrTy(), + b_.getInt64Ty()->getPointerTo(), less_than_function->getType()}, /*isVarArg=*/false); - auto* key_value_sort_func = llvm::cast( - module_->getOrInsertFunction(fn_name, key_value_sort_type)); + auto* key_value_sort_func = llvm::dyn_cast( + module_ + ->getOrInsertFunction(runtime::kKeyValueSortSymbolName, + key_value_sort_type) + .getCallee()); key_value_sort_func->setCallingConv(llvm::CallingConv::C); key_value_sort_func->setDoesNotThrow(); - llvm::Value* values; - llvm::Value* sizes; - if (sort->values_count() == 0) { - values = llvm::Constant::getNullValue(b_.getInt8PtrTy()->getPointerTo()); - sizes = llvm::Constant::getNullValue(b_.getInt32Ty()->getPointerTo()); - } else { - values = llvm_ir::EmitAllocaAtFunctionEntryWithCount( - b_.getInt8PtrTy(), b_.getInt32(sort->values_count()), - "cc_values_alloca", &b_); - sizes = llvm_ir::EmitAllocaAtFunctionEntryWithCount( - b_.getInt32Ty(), b_.getInt32(sort->values_count()), "cc_sizes_alloca", - &b_); - for (int64 i = 0; i < sort->values_count(); ++i) { - llvm::Value* value_as_i8ptr = - PointerCast(destination_addresses[i + 1], b_.getInt8PtrTy()); - llvm::Value* slot_in_values_alloca = - ConstInBoundsGEP1_32(b_.getInt8PtrTy(), values, i); - Store(value_as_i8ptr, slot_in_values_alloca); - llvm::Value* slot_in_sizes_alloca = - ConstInBoundsGEP1_32(b_.getInt32Ty(), sizes, i); - llvm::Value* size = b_.getInt32(ShapeUtil::ByteSizeOfPrimitiveType( - sort->operand(i + 1)->shape().element_type())); - Store(size, slot_in_sizes_alloca); - } + llvm::Value* values = llvm_ir::EmitAllocaAtFunctionEntryWithCount( + b_.getInt8PtrTy(), b_.getInt32(sort->operand_count()), "cc_values_alloca", + &b_); + llvm::Value* sizes = llvm_ir::EmitAllocaAtFunctionEntryWithCount( + b_.getInt32Ty(), b_.getInt32(sort->operand_count()), "cc_sizes_alloca", + &b_); + for (int64 i = 0; i < sort->operand_count(); ++i) { + llvm::Value* value_as_i8ptr = + PointerCast(destination_addresses[i], b_.getInt8PtrTy()); + llvm::Value* slot_in_values_alloca = + ConstInBoundsGEP1_32(b_.getInt8PtrTy(), values, i); + Store(value_as_i8ptr, slot_in_values_alloca); + llvm::Value* slot_in_sizes_alloca = + ConstInBoundsGEP1_32(b_.getInt32Ty(), sizes, i); + llvm::Value* size = b_.getInt32(ShapeUtil::ByteSizeOfPrimitiveType( + sort->operand(i)->shape().element_type())); + Store(size, slot_in_sizes_alloca); } Call(key_value_sort_func, - {PointerCast(destination_addresses[0], keys_native_type), - b_.getInt64(higher_dimensions), b_.getInt64(sort_dimension_elements), + {b_.getInt64(higher_dimensions), b_.getInt64(sort_dimension_elements), b_.getInt64(lower_dimensions), values, - b_.getInt32(sort->values_count()), sizes}); + b_.getInt32(sort->operand_count()), sizes, + GetExecutableRunOptionsArgument(), GetProfileCountersArgument(), + less_than_function}); if (sort->values_count() > 0) { llvm_ir::EmitTuple(GetIrArrayFor(sort), destination_addresses, &b_, @@ -747,11 +726,6 @@ StatusOr IrEmitter::EmitTargetElementLoopBodyForReduceWindow( } Status IrEmitter::HandleReduceWindow(HloInstruction* reduce_window) { - TF_RETURN_IF_ERROR(ElementTypesSameAndSupported( - /*instruction=*/*reduce_window, - /*operands=*/{reduce_window->operand(0)}, - /*supported_types=*/{F32, BF16, S32, F16})); - // Pseudo code for reduce window: // // for (coordinates O in the output) @@ -942,12 +916,8 @@ Status IrEmitter::HandleDot(HloInstruction* dot) { auto rhs = dot->operand(1); TF_RETURN_IF_ERROR(ElementTypesSameAndSupported( /*instruction=*/*dot, /*operands=*/{lhs, rhs}, - /*supported_types=*/{F16, F32, F64, C64})); + /*supported_types=*/{F16, F32, F64, C64, C128})); const DotDimensionNumbers& dnums = dot->dot_dimension_numbers(); - if (dnums.lhs_batch_dimensions_size() > 0 || - dnums.rhs_batch_dimensions_size() > 0) { - return Unimplemented("Dot with batch dimensions not implemented."); - } if (dnums.lhs_contracting_dimensions_size() != 1) { // This is disallowed by ShapeInference today. @@ -970,10 +940,10 @@ Status IrEmitter::HandleDot(HloInstruction* dot) { << llvm_ir::DumpToString(*target_array.GetBasePointer()); // Dot operation is complicated so we delegate to a helper class. - return DotOpEmitter::EmitDotOperation( - *dot, target_array, lhs_array, rhs_array, /*addend_array=*/nullptr, - GetExecutableRunOptionsArgument(), &b_, hlo_module_config_, - target_machine_features_); + return EmitDotOperation(*dot, target_array, lhs_array, rhs_array, + /*addend_array=*/nullptr, + GetExecutableRunOptionsArgument(), &b_, + hlo_module_config_, target_machine_features_); } StatusOr IrEmitter::EmitTargetElementLoopBodyForConvolution( @@ -1118,7 +1088,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})); + /*supported_types=*/{F16, F32, C64, C128})); // TODO(tonywy): Add PotentiallyImplementedAsMKLCovolution to support // different data layouts. @@ -1231,8 +1201,8 @@ Status IrEmitter::HandleConvolution(HloInstruction* convolution) { LOG(WARNING) << "Using Eigen instead of MKL-DNN for single-threaded " "conv2d function."; } - llvm::Function* conv_func = llvm::cast( - module_->getOrInsertFunction(fn_name, conv_type)); + llvm::Function* conv_func = llvm::dyn_cast( + module_->getOrInsertFunction(fn_name, conv_type).getCallee()); conv_func->setCallingConv(llvm::CallingConv::C); conv_func->setDoesNotThrow(); conv_func->setOnlyAccessesArgMemory(); @@ -1315,8 +1285,8 @@ Status IrEmitter::HandleFft(HloInstruction* fft) { ? runtime::kEigenFftSymbolName : runtime::kEigenSingleThreadedFftSymbolName; - llvm::Function* fft_func = llvm::cast( - module_->getOrInsertFunction(fn_name, fft_type)); + llvm::Function* fft_func = llvm::dyn_cast( + module_->getOrInsertFunction(fn_name, fft_type).getCallee()); fft_func->setCallingConv(llvm::CallingConv::C); fft_func->setDoesNotThrow(); fft_func->setOnlyAccessesInaccessibleMemOrArgMem(); @@ -1399,7 +1369,7 @@ static bool ReductionPreservesLayout(const HloInstruction& reduce) { int64 delta = 0; for (int64 i = 0; i < operand_shape.dimensions_size(); i++) { - if (reduced_dims.count(i)) { + if (reduced_dims.contains(i)) { delta++; } else { InsertOrDie(&unreduced_dim_map, i, i - delta); @@ -1412,7 +1382,7 @@ static bool ReductionPreservesLayout(const HloInstruction& reduce) { for (int64 operand_dim_idx = 0; operand_dim_idx < operand_shape.dimensions_size(); operand_dim_idx++) { int64 operand_dim = operand_shape.layout().minor_to_major(operand_dim_idx); - if (!reduced_dims.count(operand_dim)) { + if (!reduced_dims.contains(operand_dim)) { if (FindOrDie(unreduced_dim_map, operand_dim) != result_shape.layout().minor_to_major(result_dim_idx++)) { return false; @@ -1709,10 +1679,8 @@ StatusOr IrEmitter::EmitVectorizedReduce( vectorization_factor_in_bytes / ShapeUtil::ByteSizeOfPrimitiveType(reduce->shape().element_type()); - bool is_reduction_over_minor_dimension = - std::find(dimensions.begin(), dimensions.end(), - LayoutUtil::Minor(arg->shape().layout(), 0)) != - dimensions.end(); + bool is_reduction_over_minor_dimension = absl::c_linear_search( + dimensions, LayoutUtil::Minor(arg->shape().layout(), 0)); unsigned element_alignment = tensorflow::MathUtil::GCD( ShapeUtil::ByteSizeOfPrimitiveType(reduce->shape().element_type()), @@ -1890,7 +1858,7 @@ StatusOr IrEmitter::EmitTargetElementLoopBodyForReduce( } Status IrEmitter::HandleReduce(HloInstruction* reduce) { - // TODO(b/112040122): Support variadic reduce. + // TODO(b/118333695): Support variadic reduce. if (!reduce->shape().IsArray()) { return Unimplemented("Variadic reduce is not supported on CPU"); } @@ -1990,7 +1958,7 @@ Status IrEmitter::HandleSlice(HloInstruction* slice) { // The memcpy will copy elements that are logically this shape (allowed to be // scalar). const Shape logical_element_shape = ShapeUtil::FilterDimensions( - [&inner_dims](int64 dim) -> bool { return inner_dims.count(dim); }, + [&inner_dims](int64 dim) { return inner_dims.contains(dim); }, operand->shape()); const int64 primitive_elements_per_logical_element = @@ -2205,10 +2173,10 @@ Status IrEmitter::HandleFusion(HloInstruction* fusion) { llvm_ir::IrArray addend_array( GetIrArrayFor(fusion->operand(addend_param_number))); - TF_RETURN_IF_ERROR(DotOpEmitter::EmitDotOperation( - *dot, target_array, lhs_array, rhs_array, &addend_array, - GetExecutableRunOptionsArgument(), &b_, hlo_module_config_, - target_machine_features_)); + TF_RETURN_IF_ERROR( + EmitDotOperation(*dot, target_array, lhs_array, rhs_array, + &addend_array, GetExecutableRunOptionsArgument(), &b_, + hlo_module_config_, target_machine_features_)); return Status::OK(); } else { return Unimplemented("Fusion kind not implemented on CPU"); @@ -2257,13 +2225,15 @@ Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) { InBoundsGEP(operands_alloca, {b_.getInt64(i)}); Store(operand_as_i8ptr, slot_in_operands_alloca); } - auto* custom_call_ir_function = - llvm::cast(module_->getOrInsertFunction( - AsStringRef(custom_call_target), - llvm::FunctionType::get( - /*Result=*/b_.getVoidTy(), - /*Params=*/{i8_ptr_type, operands_alloca->getType()}, - /*isVarArg=*/false))); + auto* custom_call_ir_function = llvm::dyn_cast( + module_ + ->getOrInsertFunction( + AsStringRef(custom_call_target), + llvm::FunctionType::get( + /*Result=*/b_.getVoidTy(), + /*Params=*/{i8_ptr_type, operands_alloca->getType()}, + /*isVarArg=*/false)) + .getCallee()); TF_RETURN_IF_ERROR(EmitTargetAddressForOp(custom_call)); // Write the tuple table if the output is a tuple. @@ -2401,8 +2371,7 @@ StatusOr IrEmitter::EmitFastConcatenate( int64 concat_dim = concatenate->dimensions(0); const Layout& output_layout = output_shape.layout(); auto output_min2maj = LayoutUtil::MinorToMajor(output_layout); - auto concat_dim_layout_itr = - std::find(output_min2maj.begin(), output_min2maj.end(), concat_dim); + auto concat_dim_layout_itr = absl::c_find(output_min2maj, concat_dim); std::vector inner_dims(output_min2maj.begin(), concat_dim_layout_itr); std::vector outer_dims(std::next(concat_dim_layout_itr), @@ -2956,8 +2925,7 @@ Status IrEmitter::ElementTypesSameAndSupported( TF_RET_CHECK(!operands.empty()); PrimitiveType primitive_type = operands[0]->shape().element_type(); - if (std::find(supported_types.begin(), supported_types.end(), - primitive_type) == supported_types.end()) { + if (!absl::c_linear_search(supported_types, primitive_type)) { return Unimplemented("unsupported operand type %s in op %s", PrimitiveType_Name(primitive_type), HloOpcodeString(instruction.opcode())); diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.h b/tensorflow/compiler/xla/service/cpu/ir_emitter.h index db76de4bb2b8ed568bf2557a30fa216d0cbe518d..974dd7cd3f2254bfbc86fffae02c06c481af8902 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.h +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.h @@ -250,14 +250,6 @@ class IrEmitter : public DfsHloVisitorWithDefault, llvm::Value* EmitBufferPointer(const BufferAllocation::Slice& slice, const Shape& target_shape); - // Emits a function into the current module. This can be used for - // computations embedded inside other computations, such as the - // function that a map operation applies. - StatusOr EmitFunction( - HloComputation* function, // The function to emit. - absl::string_view - function_name_suffix); // Used for LLVM IR register names. - // Emits a call to a thread local function (e.g. to the computation nested // within a reduce or a map). Thread local callees (by definition) only write // to and read from thread local allocations. @@ -448,7 +440,7 @@ class IrEmitter : public DfsHloVisitorWithDefault, computation_to_profile_idx_; // Maps HLOs to Values emitted for them. - std::unordered_map emitted_value_; + absl::flat_hash_map emitted_value_; llvm_ir::AliasAnalysis alias_analysis_; diff --git a/tensorflow/compiler/xla/service/cpu/ir_function.cc b/tensorflow/compiler/xla/service/cpu/ir_function.cc index adfb8392bf6fa356f0a5cdab3ff74036eca8918e..84a5b058cfb11c899eb6ae03478ed550b84dc819 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_function.cc +++ b/tensorflow/compiler/xla/service/cpu/ir_function.cc @@ -266,9 +266,11 @@ Status EmitCallToParallelForkJoin( /*Params=*/compute_function_params, /*isVarArg=*/false); - llvm::Function* fork_join_func = - llvm::cast(module->getOrInsertFunction( - runtime::kParallelForkJoinSymbolName, fork_join_type)); + llvm::Function* fork_join_func = llvm::dyn_cast( + module + ->getOrInsertFunction(runtime::kParallelForkJoinSymbolName, + fork_join_type) + .getCallee()); fork_join_func->setCallingConv(llvm::CallingConv::C); fork_join_func->setDoesNotThrow(); diff --git a/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc index f9722ffadac801521ddcbb568dd4435fd02e951b..93ef51754d21ad3ff4e24298c89649ef4c2742fb 100644 --- a/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc +++ b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc @@ -36,57 +36,88 @@ const char* const kLogV4F32SymbolName = "__xla_cpu_runtime_LogV4F32AVX"; const char* const kLogV8F32SymbolName = "__xla_cpu_runtime_LogV8F32AVX"; namespace { -llvm::Function* EmitVectorF32TanhIfNeeded(llvm::Module* module, - llvm::StringRef function_name, - int vector_width, - bool enable_fast_math) { - llvm::Function* vector_tanh_function = module->getFunction(function_name); - if (vector_tanh_function == nullptr) { + +// Replaces calls to the function `fn_name` with the code generated by +// fn_body_generator. +// +// We assume that fn_name accepts either a scalar f32 or a vector of +// vector_width f32s, and that fn_body_generator generates a function body with +// the same inputs/outputs as fn_name. +void RewriteCalls( + llvm::Module* module, const char* fn_name, + std::function* b, llvm::Value* input, + int32 vector_width)> + fn_body_generator, + int32 vector_width, bool enable_fast_math) { + llvm::Function* fn = module->getFunction(fn_name); + if (fn == nullptr) { // If the function declaration is not present in the module, there can't be // any calls to resolve. Don't emit the function in this case. - return nullptr; + return; } - llvm::LLVMContext* context = &module->getContext(); + // Our task is to generate a function body for `fn`, but we can't generate a + // function body for an LLVM intrinsic. So if fn is an intrinsic, replace it + // with a new function. + if (fn->isIntrinsic()) { + llvm::Function* new_fn = llvm::Function::Create( + fn->getFunctionType(), llvm::GlobalValue::InternalLinkage, + llvm::Twine("xla_impl.") + fn_name, module); + fn->replaceAllUsesWith(new_fn); + fn->eraseFromParent(); + fn = new_fn; + } - llvm::BasicBlock* vector_tanh_body = - llvm::BasicBlock::Create(*context, "body", vector_tanh_function); + llvm::LLVMContext* context = &module->getContext(); - llvm::IRBuilder<> b(vector_tanh_body); + llvm::BasicBlock* fn_body = llvm::BasicBlock::Create(*context, "body", fn); + llvm::IRBuilder<> b(fn_body); llvm::FastMathFlags fast_math_flags; fast_math_flags.setFast(enable_fast_math); b.setFastMathFlags(fast_math_flags); - llvm::Value* input = &*vector_tanh_function->arg_begin(); - CHECK_EQ(vector_width, input->getType()->getVectorNumElements()); - b.CreateRet(llvm_ir::EmitFastTanh(&b, input)); - - DCHECK(!llvm::verifyFunction(*vector_tanh_function)); - return vector_tanh_function; -} + llvm::Value* input = &*fn->arg_begin(); -llvm::Function* EmitVectorF32ExpIfNeeded(llvm::Module* module, - llvm::StringRef function_name, - int vector_width, - bool enable_fast_math) { - llvm::Function* vector_exp_function = module->getFunction(function_name); - if (vector_exp_function == nullptr) { - // If the function declaration is not present in the module, there can't be - // any calls to resolve. Don't emit the function in this case. - return nullptr; + // Upcast to vector type if input is a scalar. + if (vector_width == 1) { + llvm::Type* v1_type = llvm::VectorType::get(input->getType(), 1); + input = b.CreateInsertElement(llvm::UndefValue::get(v1_type), input, + uint64_t{0}); } - llvm::LLVMContext* context = &module->getContext(); + // Generate the vectorized code. + CHECK_EQ(vector_width, input->getType()->getVectorNumElements()); + llvm::Value* result = fn_body_generator(&b, input, vector_width); + + // Downcast result to scalar type if necessary. + if (vector_width == 1) { + result = b.CreateExtractElement(result, uint64_t{0}); + } + b.CreateRet(result); + DCHECK(!llvm::verifyFunction(*fn)); - llvm::BasicBlock* vector_exp_body = - llvm::BasicBlock::Create(*context, "body", vector_exp_function); + // Force-inline `fn` into all of its callers and then delete `fn`. + // + // TODO(b/73081976): Should we avoid inlining these in some cases? + std::vector calls_to_inline; + for (auto* user : fn->users()) { + calls_to_inline.push_back(llvm::cast(user)); + } + for (auto* call_to_inline : calls_to_inline) { + llvm::InlineFunctionInfo inline_function_info; + CHECK(llvm::InlineFunction(call_to_inline, inline_function_info)); + } + fn->eraseFromParent(); +} - llvm::IRBuilder<> b(vector_exp_body); - llvm::FastMathFlags fast_math_flags; - fast_math_flags.setFast(); - b.setFastMathFlags(fast_math_flags); +llvm::Value* GenerateVF32Tanh(llvm::IRBuilder<>* b, llvm::Value* input, + int32 /*vector_width*/) { + return llvm_ir::EmitFastTanh(b, input); +} - VectorSupportLibrary vsl(F32, vector_width, &b, "exp_f32"); +llvm::Value* GenerateVF32Exp(llvm::IRBuilder<>* b, llvm::Value* input, + int32 vector_width) { + VectorSupportLibrary vsl(F32, vector_width, b, "exp_f32"); // This implements the same polynomial approximation as implemented in Eigen3. @@ -107,7 +138,6 @@ llvm::Function* EmitVectorF32ExpIfNeeded(llvm::Module* module, const llvm::APFloat cephes_exp_p4 = GetIeeeF32(1.6666665459E-1); const llvm::APFloat cephes_exp_p5 = GetIeeeF32(5.0000001201E-1); - llvm::Value* input = &*vector_exp_function->arg_begin(); llvm::Value* input_clamped = vsl.Clamp(input, /*low=*/exp_lo, /*high=*/exp_hi); llvm::Value* fx = vsl.Floor(vsl.MulAdd(input_clamped, cephes_LOG2EF, half)); @@ -128,49 +158,24 @@ llvm::Function* EmitVectorF32ExpIfNeeded(llvm::Module* module, // VectorSupportLibrary (intentionally) can't juggle more than one type at a // time so drop down to IRBuilder for this bit. llvm::Value* vector_constant_0x7f = - b.CreateVectorSplat(vector_width, b.getInt32(0x7f)); + b->CreateVectorSplat(vector_width, b->getInt32(0x7f)); llvm::Value* vector_constant_23 = - b.CreateVectorSplat(vector_width, b.getInt32(23)); + b->CreateVectorSplat(vector_width, b->getInt32(23)); llvm::Type* i32_vector_type = - llvm::VectorType::get(b.getInt32Ty(), vector_width); + llvm::VectorType::get(b->getInt32Ty(), vector_width); // fx is clamped so we don't have to worry about it being out of range for // i32. - llvm::Value* emm0 = b.CreateFPToSI(fx, i32_vector_type); - emm0 = b.CreateAdd(emm0, vector_constant_0x7f); - emm0 = b.CreateShl(emm0, vector_constant_23); - llvm::Value* emm0_f32 = b.CreateBitCast(emm0, vsl.vector_type()); - - llvm::Value* result = vsl.Max(vsl.Mul(y, emm0_f32), input); + llvm::Value* emm0 = b->CreateFPToSI(fx, i32_vector_type); + emm0 = b->CreateAdd(emm0, vector_constant_0x7f); + emm0 = b->CreateShl(emm0, vector_constant_23); + llvm::Value* emm0_f32 = b->CreateBitCast(emm0, vsl.vector_type()); - b.CreateRet(result); - - DCHECK(!llvm::verifyFunction(*vector_exp_function)); - return vector_exp_function; + return vsl.Max(vsl.Mul(y, emm0_f32), input); } -llvm::Function* EmitVectorF32LogIfNeeded(llvm::Module* module, - llvm::StringRef function_name, - int vector_width, - bool enable_fast_math) { - llvm::Function* vector_log_function = module->getFunction(function_name); - if (vector_log_function == nullptr) { - // If the function declaration is not present in the module, there can't be - // any calls to resolve. Don't emit the function in this case. - return nullptr; - } - - llvm::LLVMContext* context = &module->getContext(); - - llvm::BasicBlock* vector_log_body = - llvm::BasicBlock::Create(*context, "body", vector_log_function); - - llvm::IRBuilder<> b(vector_log_body); - llvm::FastMathFlags fast_math_flags; - fast_math_flags.setFast(); - b.setFastMathFlags(fast_math_flags); - - llvm::Value* input = &*vector_log_function->arg_begin(); - VectorSupportLibrary vsl(F32, vector_width, &b, "log_f32"); +llvm::Value* GenerateVF32Log(llvm::IRBuilder<>* b, llvm::Value* input, + int32 vector_width) { + VectorSupportLibrary vsl(F32, vector_width, b, "log_f32"); const llvm::APFloat half = GetIeeeF32(0.5); const llvm::APFloat one = GetIeeeF32(1.0); @@ -193,129 +198,107 @@ llvm::Function* EmitVectorF32LogIfNeeded(llvm::Module* module, // The smallest non denormalized float number. const llvm::APFloat min_norm_pos = GetIeeeF32FromBitwiseRep(0x00800000); const llvm::APFloat minus_inf = GetIeeeF32FromBitwiseRep(0xff800000); + const llvm::APFloat pos_inf = GetIeeeF32FromBitwiseRep(0x7f800000); const llvm::APFloat inv_mant_mask = GetIeeeF32FromBitwiseRep(~0x7f800000); // invalid_mask is set if x is negative or NaN (and therefore output // must be NaN). llvm::Value* invalid_mask = vsl.FCmpULEMask(input, vsl.GetZeroVector()); - llvm::Value* iszero_mask = vsl.FCmpEQMask(input, vsl.GetZeroVector()); + llvm::Value* is_zero_mask = vsl.FCmpEQMask(input, vsl.GetZeroVector()); + llvm::Value* is_pos_inf_mask = vsl.FCmpEQMask(input, pos_inf); // Cut off denormalized stuff. - input = vsl.Max(min_norm_pos, input); + llvm::Value* tmp0 = vsl.Max(min_norm_pos, input); // VectorSupportLibrary (intentionally) can't juggle more than one type at a // time so drop down to IRBuilder for this bit. llvm::Value* vector_constant_0x7f = - b.CreateVectorSplat(vector_width, b.getInt32(0x7f)); + b->CreateVectorSplat(vector_width, b->getInt32(0x7f)); llvm::Value* vector_constant_23 = - b.CreateVectorSplat(vector_width, b.getInt32(23)); + b->CreateVectorSplat(vector_width, b->getInt32(23)); llvm::Type* i32_vector_type = - llvm::VectorType::get(b.getInt32Ty(), vector_width); + llvm::VectorType::get(b->getInt32Ty(), vector_width); - llvm::Value* emm0 = - b.CreateLShr(b.CreateBitCast(input, i32_vector_type), vector_constant_23); + llvm::Value* emm0 = b->CreateLShr(b->CreateBitCast(tmp0, i32_vector_type), + vector_constant_23); // Keep only the fractional part. - input = vsl.FloatAnd(input, inv_mant_mask); - input = vsl.FloatOr(input, half); + tmp0 = vsl.FloatAnd(tmp0, inv_mant_mask); + tmp0 = vsl.FloatOr(tmp0, half); - emm0 = b.CreateSub(emm0, vector_constant_0x7f); - llvm::Value* e = vsl.Add(one, b.CreateSIToFP(emm0, vsl.vector_type())); + emm0 = b->CreateSub(emm0, vector_constant_0x7f); + llvm::Value* e = vsl.Add(one, b->CreateSIToFP(emm0, vsl.vector_type())); // part2: // if( x < SQRTHF ) { // e -= 1; // x = x + x - 1.0; // } else { x = x - 1.0; } - llvm::Value* mask = vsl.FCmpOLTMask(input, cephes_SQRTHF); - llvm::Value* tmp = vsl.FloatAnd(input, mask); - input = vsl.Sub(input, one); + llvm::Value* mask = vsl.FCmpOLTMask(tmp0, cephes_SQRTHF); + llvm::Value* tmp1 = vsl.FloatAnd(tmp0, mask); + tmp0 = vsl.Sub(tmp0, one); e = vsl.Sub(e, vsl.FloatAnd(mask, one)); - input = vsl.Add(input, tmp); + tmp0 = vsl.Add(tmp0, tmp1); - llvm::Value* x2 = vsl.Mul(input, input); - llvm::Value* x3 = vsl.Mul(x2, input); + llvm::Value* x2 = vsl.Mul(tmp0, tmp0); + llvm::Value* x3 = vsl.Mul(x2, tmp0); llvm::Value *y, *y1, *y2; - y = vsl.MulAdd(input, cephes_log_p0, cephes_log_p1); - y1 = vsl.MulAdd(input, cephes_log_p3, cephes_log_p4); - y2 = vsl.MulAdd(input, cephes_log_p6, cephes_log_p7); - y = vsl.MulAdd(y, input, cephes_log_p2); - y1 = vsl.MulAdd(y1, input, cephes_log_p5); - y2 = vsl.MulAdd(y2, input, cephes_log_p8); + y = vsl.MulAdd(tmp0, cephes_log_p0, cephes_log_p1); + y1 = vsl.MulAdd(tmp0, cephes_log_p3, cephes_log_p4); + y2 = vsl.MulAdd(tmp0, cephes_log_p6, cephes_log_p7); + y = vsl.MulAdd(y, tmp0, cephes_log_p2); + y1 = vsl.MulAdd(y1, tmp0, cephes_log_p5); + y2 = vsl.MulAdd(y2, tmp0, cephes_log_p8); y = vsl.MulAdd(y, x3, y1); y = vsl.MulAdd(y, x3, y2); y = vsl.Mul(y, x3); y1 = vsl.Mul(cephes_log_q1, e); - tmp = vsl.Mul(half, x2); + llvm::Value* tmp2 = vsl.Mul(half, x2); y = vsl.Add(y, y1); - input = vsl.Sub(input, tmp); + tmp0 = vsl.Sub(tmp0, tmp2); y2 = vsl.Mul(cephes_log_q2, e); - input = vsl.Add(input, y); - input = vsl.Add(input, y2); + tmp0 = vsl.Add(tmp0, y); + tmp0 = vsl.Add(tmp0, y2); - // Negative arg will be NAN, 0 will be -INF. - llvm::Value* or_lhs = - vsl.FloatAndNot(iszero_mask, vsl.FloatOr(input, invalid_mask)); - llvm::Value* or_rhs = vsl.FloatAnd(iszero_mask, minus_inf); - llvm::Value* result = vsl.FloatOr(or_lhs, or_rhs); + // Contains +/-inf where +/-inf is the correct answer, otherwise 0. + llvm::Value* result_inf = vsl.FloatOr(vsl.FloatAnd(is_zero_mask, minus_inf), + vsl.FloatAnd(is_pos_inf_mask, pos_inf)); - b.CreateRet(result); + // Contains a finite result or nan. This is the correct answer only if both + // result_minus_inf and result_pos_inf are both 0. + // + // (This implementation works because 0xffffffff is a nan.) + llvm::Value* result_finite_or_nan = vsl.FloatOr(tmp0, invalid_mask); - DCHECK(!llvm::verifyFunction(*vector_log_function)); - return vector_log_function; + // Combine the above into a final result. + return vsl.FloatOr(result_inf, + vsl.FloatAndNot(vsl.FloatOr(is_zero_mask, is_pos_inf_mask), + result_finite_or_nan)); } } // namespace void RewriteIRRuntimeFunctions(llvm::Module* module, bool enable_fast_math) { - auto* tanh_v4f32 = - EmitVectorF32TanhIfNeeded(module, kTanhV4F32SymbolName, - /*vector_width=*/4, enable_fast_math); - auto* tanh_v8f32 = - EmitVectorF32TanhIfNeeded(module, kTanhV8F32SymbolName, - /*vector_width=*/8, enable_fast_math); - - auto* exp_v4f32 = - EmitVectorF32ExpIfNeeded(module, kExpV4F32SymbolName, - /*vector_width=*/4, enable_fast_math); - auto* exp_v8f32 = - EmitVectorF32ExpIfNeeded(module, kExpV8F32SymbolName, - /*vector_width=*/8, enable_fast_math); - - auto* log_v4f32 = - EmitVectorF32LogIfNeeded(module, kLogV4F32SymbolName, - /*vector_width=*/4, enable_fast_math); - auto* log_v8f32 = - EmitVectorF32LogIfNeeded(module, kLogV8F32SymbolName, - /*vector_width=*/8, enable_fast_math); - - // Gather all the call sites, force inline them and then delete the vector - // function bodies. - // - // TODO(b/73081976): Should we avoid inlining these intrinsics in some cases? - - std::vector calls_to_inline; - for (auto* function : - {tanh_v4f32, tanh_v8f32, exp_v4f32, exp_v8f32, log_v4f32, log_v8f32}) { - if (function != nullptr) { - for (auto* user : function->users()) { - calls_to_inline.push_back(llvm::cast(user)); - } - } - } - - for (auto* call_to_inline : calls_to_inline) { - llvm::InlineFunctionInfo inline_function_info; - CHECK(llvm::InlineFunction(call_to_inline, inline_function_info)); - } - - for (auto* function : - {tanh_v4f32, tanh_v8f32, exp_v4f32, exp_v8f32, log_v4f32, log_v8f32}) { - if (function != nullptr) { - function->eraseFromParent(); - } - } + // Curry some params to RewriteCalls. + auto rewrite_calls = + std::bind(RewriteCalls, module, std::placeholders::_1, + std::placeholders::_2, std::placeholders::_3, enable_fast_math); + + rewrite_calls("tanhf", GenerateVF32Tanh, /*vector_width=*/1); + rewrite_calls("llvm.tanh.f32", GenerateVF32Tanh, /*vector_width=*/1); + rewrite_calls(kTanhV4F32SymbolName, GenerateVF32Tanh, /*vector_width=*/4); + rewrite_calls(kTanhV8F32SymbolName, GenerateVF32Tanh, /*vector_width=*/8); + + rewrite_calls("expf", GenerateVF32Exp, /*vector_width=*/1); + rewrite_calls("llvm.exp.f32", GenerateVF32Exp, /*vector_width=*/1); + rewrite_calls(kExpV4F32SymbolName, GenerateVF32Exp, /*vector_width=*/4); + rewrite_calls(kExpV8F32SymbolName, GenerateVF32Exp, /*vector_width=*/8); + + rewrite_calls("logf", GenerateVF32Log, /*vector_width=*/1); + rewrite_calls("llvm.log.f32", GenerateVF32Log, /*vector_width=*/1); + rewrite_calls(kLogV4F32SymbolName, GenerateVF32Log, /*vector_width=*/4); + rewrite_calls(kLogV8F32SymbolName, GenerateVF32Log, /*vector_width=*/8); } } // namespace runtime diff --git a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc index 3b423f639137149f7efdee91a02bbcf01d1cc6e1..6121d1ca9a5c785cedd947200d3e7e320aa06bc2 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc +++ b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc @@ -146,8 +146,6 @@ int64 ParallelTaskAssignment::GetTargetParallelTaskCount( (opcode == HloOpcode::kConvolution && PotentiallyImplementedAsEigenConvolution(*instruction, target_machine_features_)) || - PotentiallyImplementedAsEigenDot(*instruction, - target_machine_features_) || (opcode == HloOpcode::kFusion && instruction->fusion_kind() != HloInstruction::FusionKind::kLoop) || instruction->shape().IsTuple()) { diff --git a/tensorflow/compiler/xla/service/cpu/runtime_fork_join.cc b/tensorflow/compiler/xla/service/cpu/runtime_fork_join.cc index 2d9492eacfea34bec3b0f1115e171a5328b7cdc3..6f72ddadf94d4c5b9add2ee66e0f4ac9a8ae9099 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_fork_join.cc +++ b/tensorflow/compiler/xla/service/cpu/runtime_fork_join.cc @@ -69,8 +69,13 @@ TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_ParallelForkJoin( CHECK_EQ(params, nullptr); CHECK_GT(num_partitions, 1); CHECK_GT(num_partitioned_dims, 0); + CHECK_NE(function_ptr, nullptr); + CHECK_NE(partitions, nullptr); const xla::ExecutableRunOptions* run_options = static_cast(run_options_ptr); + CHECK_NE(run_options, nullptr); + CHECK_NE(run_options->intra_op_thread_pool(), nullptr); + ComputeFunctionType function = reinterpret_cast(function_ptr); // Compute partition stride in 'partitions' array. diff --git a/tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.cc b/tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.cc index 722aa3120ef4d8c957873ac58c361f19632dde1f..cb46674138acf8e1daea24102f988a9a355ec5c8 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.cc +++ b/tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.cc @@ -15,12 +15,10 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.h" #include -#include #include -#include #include +#include #include -#include #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/core/platform/dynamic_annotations.h" @@ -28,80 +26,15 @@ limitations under the License. #include "tensorflow/core/platform/types.h" namespace { -using tensorflow::int16; using tensorflow::int32; using tensorflow::int64; -using tensorflow::int8; -using tensorflow::uint16; -using tensorflow::uint32; -using tensorflow::uint64; -using tensorflow::uint8; - -template -void KeyValueSort(std::pair* row_to_sort, int64 num_elements) { - std::sort(row_to_sort, row_to_sort + num_elements); -} - -// We would like a total order of floating point numbers so that the -// sort has a predictable behavior in the presence of NaNs. Rather -// than using floating point comparison, we use the following trick: -// If f is a float, and -// x = bit_cast(f); -// y = x < 0 ? 0x7FFFFFFF - x : x; -// then y is ordered as an int32 such that finite values have the -// obvious order, -0 is ordered before 0, and -NaN and NaN appear at -// the beginning and end of the ordering. -template -CastType Convert(KeyType value) { - CastType casted_value; - memcpy(&casted_value, &value, sizeof(CastType)); - if (casted_value < 0) { - return static_cast(std::numeric_limits::max()) - - casted_value; - } - return casted_value; -} - -template -bool LessThan(KeyType lhs, KeyType rhs) { - return Convert(lhs) < - Convert(rhs); -} - -template <> -void KeyValueSort(std::pair* row_to_sort, int64 num_elements) { - std::stable_sort(row_to_sort, row_to_sort + num_elements, - [](const std::pair& lhs, - const std::pair& rhs) -> bool { - return LessThan(lhs.first, rhs.first); - }); -} - -template <> -void KeyValueSort(std::pair* row_to_sort, int64 num_elements) { - std::stable_sort(row_to_sort, row_to_sort + num_elements, - [](const std::pair& lhs, - const std::pair& rhs) -> bool { - return LessThan(lhs.first, rhs.first); - }); -} - -template <> -void KeyValueSort(std::pair* row_to_sort, - int64 num_elements) { - std::stable_sort(row_to_sort, row_to_sort + num_elements, - [](const std::pair& lhs, - const std::pair& rhs) -> bool { - return LessThan( - Eigen::half_impl::half_to_float(lhs.first), - Eigen::half_impl::half_to_float(rhs.first)); - }); -} +} // namespace -template -void KeyValueSortImpl(KeyType* keys, int64 a, int64 b, int64 c, char** values, - int32 values_count, - int32* values_primitive_type_size_in_bytes) { +TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSort( + int64 a, int64 b, int64 c, char** values, int32 values_count, + int32* values_primitive_type_size_in_bytes, char* run_options, + int64* prof_counters, + void (*less_than)(char*, char*, char**, char**, tensorflow::int64*)) { // 'values' and 'values_primitive_type_size_in_bytes' are managed by the JIT // code, so msan can't tell they are initialized. TF_ANNOTATE_MEMORY_IS_INITIALIZED(values, values_count * sizeof(char*)); @@ -121,8 +54,9 @@ void KeyValueSortImpl(KeyType* keys, int64 a, int64 b, int64 c, char** values, int64 num_iteration_elements = a * c; int64 sort_dimension_offset = c; - std::unique_ptr[]> row_to_sort( - new std::pair[sort_dimension_elements]); + std::unique_ptr indices(new int64[sort_dimension_elements]); + std::unique_ptr comparison_values(new char*[2 * values_count]); + std::iota(indices.get(), indices.get() + sort_dimension_elements, 0); std::unique_ptr reordered_values( new std::string[sort_dimension_elements]); for (int64 index = 0; index < num_iteration_elements; ++index) { @@ -135,24 +69,28 @@ void KeyValueSortImpl(KeyType* keys, int64 a, int64 b, int64 c, char** values, int64 base_offset = index % sort_dimension_offset + (index - index % sort_dimension_offset) * sort_dimension_elements; - // TODO(b/26783907): We could define a custom iterator class that references - // all arrays. Then we could avoid the intermediate copy. However this - // would become more complicated, and it is not clear if the benefit is high - // enough. - for (int64 i = 0; i < sort_dimension_elements; ++i) { - row_to_sort[i] = - std::make_pair(keys[base_offset + i * sort_dimension_offset], i); - } - KeyValueSort(row_to_sort.get(), sort_dimension_elements); - for (int64 i = 0; i < sort_dimension_elements; ++i) { - keys[base_offset + i * sort_dimension_offset] = row_to_sort[i].first; - } - - // Reorder the values according to the order defined by the keys. + std::stable_sort( + indices.get(), indices.get() + sort_dimension_elements, + [&](int64 a, int64 b) -> bool { + int64 memory_index_lhs = (base_offset + a * sort_dimension_offset) * + values_primitive_type_size_in_bytes[0]; + int64 memory_index_rhs = (base_offset + b * sort_dimension_offset) * + values_primitive_type_size_in_bytes[0]; + for (int32 i = 0; i < values_count; ++i) { + comparison_values[i * 2] = values[i] + memory_index_lhs; + comparison_values[i * 2 + 1] = values[i] + memory_index_rhs; + } + char result = 0; // Overwritten by less_than. + less_than(&result, run_options, comparison_values.get(), nullptr, + prof_counters); + return result != 0u; + }); + + // Reorder the values according to the order defined by 'indices'. for (int32 idx = 0; idx < values_count; ++idx) { for (int64 i = 0; i < sort_dimension_elements; ++i) { int64 memory_index = - (base_offset + row_to_sort[i].second * sort_dimension_offset) * + (base_offset + indices[i] * sort_dimension_offset) * values_primitive_type_size_in_bytes[idx]; reordered_values[i] = @@ -168,88 +106,3 @@ void KeyValueSortImpl(KeyType* keys, int64 a, int64 b, int64 c, char** values, } } } -} // namespace - -TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSortPRED( - bool* keys, int64 a, int64 b, int64 c, char** values, int32 values_count, - int32* values_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, values_count, - values_primitive_type_size_in_bytes); -} - -TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSortS8( - int8* keys, int64 a, int64 b, int64 c, char** values, int32 values_count, - int32* values_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, values_count, - values_primitive_type_size_in_bytes); -} - -TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSortU8( - uint8* keys, int64 a, int64 b, int64 c, char** values, int32 values_count, - int32* values_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, values_count, - values_primitive_type_size_in_bytes); -} - -TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSortS16( - int16* keys, int64 a, int64 b, int64 c, char** values, int32 values_count, - int32* values_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, values_count, - values_primitive_type_size_in_bytes); -} - -TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSortU16( - uint16* keys, int64 a, int64 b, int64 c, char** values, int32 values_count, - int32* values_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, values_count, - values_primitive_type_size_in_bytes); -} - -TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSortF16( - Eigen::half* keys, int64 a, int64 b, int64 c, char** values, - int32 values_count, int32* values_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, values_count, - values_primitive_type_size_in_bytes); -} - -TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSortS32( - int32* keys, int64 a, int64 b, int64 c, char** values, int32 values_count, - int32* values_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, values_count, - values_primitive_type_size_in_bytes); -} - -TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSortU32( - uint32* keys, int64 a, int64 b, int64 c, char** values, int32 values_count, - int32* values_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, values_count, - values_primitive_type_size_in_bytes); -} - -TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSortF32( - float* keys, int64 a, int64 b, int64 c, char** values, int32 values_count, - int32* values_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, values_count, - values_primitive_type_size_in_bytes); -} - -TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSortS64( - int64* keys, int64 a, int64 b, int64 c, char** values, int32 values_count, - int32* values_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, values_count, - values_primitive_type_size_in_bytes); -} - -TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSortU64( - uint64* keys, int64 a, int64 b, int64 c, char** values, int32 values_count, - int32* values_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, values_count, - values_primitive_type_size_in_bytes); -} - -TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSortF64( - double* keys, int64 a, int64 b, int64 c, char** values, int32 values_count, - int32* values_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, values_count, - values_primitive_type_size_in_bytes); -} diff --git a/tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.h b/tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.h index 7821099386969e855ea1737cf53ef49c15c6e93b..4813de9ee67282110959f943cd70fcfe2ca94d9d 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.h +++ b/tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.h @@ -21,76 +21,27 @@ limitations under the License. extern "C" { -// 'keys' represents a 3-dimensional shape with dimensions [a, b, c]. The 'b' -// dimension of 'keys' is sorted into ascending order. If 'values_count' is <= -// 0, 'values' and 'values_primitive_type_size_in_bytes' can be nullptr. -// If 'values_count' > 0, they contain exactly 'values_count' many elements. -// Each element of 'values' also represents a 3-dimensional shape with -// dimensions [a, b, c], and the size of the primitive type of the i-th shape -// has exactly 'values_primitive_type_size_in_bytes[i]' bytes. The elements in -// each 'values' shape are reordered in such a way that if the element at index -// 'i' in 'keys' was moved to index 'j', the element at index 'i' in a 'values' -// shape is also moved to index 'j' (which means that the same elements -// correspond to each other as before). -extern void __xla_cpu_runtime_KeyValueSortPRED( - bool* keys, tensorflow::int64 a, tensorflow::int64 b, tensorflow::int64 c, +// Each entry in 'values' represents a 3-dimensional shape with dimensions +// [a, b, c]. The 'b' dimension of the first shape is sorted into ascending +// order according to the results of comparisons using the provided 'less_than' +// function. 'values_count' must be > 0 and specifies the number of entries in +// 'values' and 'values_primitive_type_size_in_bytes'. The size of the primitive +// type of the i-th shape has exactly 'values_primitive_type_size_in_bytes[i]' +// bytes. The elements in each 'values' shape are reordered in the same way +// according to the comparisons using the first shape. 'run_options' and +// 'prof_counters' are passed through to the less-than function, which expects +// the following arguments: +// - pointer to the return value buffer (char*) +// - xla::ExecutableRunOptions = 'run_options' (char*) +// - pointers to the parameter buffers (char**) +// - pointers to the buffer tables = nullptr for thread local functions (char**) +// - profile counters = 'prof_counters' (int64*) +extern void __xla_cpu_runtime_KeyValueSort( + tensorflow::int64 a, tensorflow::int64 b, tensorflow::int64 c, char** values, tensorflow::int32 values_count, - tensorflow::int32* values_primitive_type_size_in_bytes); - -extern void __xla_cpu_runtime_KeyValueSortS8( - tensorflow::int8* keys, tensorflow::int64 a, tensorflow::int64 b, - tensorflow::int64 c, char** values, tensorflow::int32 values_count, - tensorflow::int32* values_primitive_type_size_in_bytes); - -extern void __xla_cpu_runtime_KeyValueSortU8( - tensorflow::uint8* keys, tensorflow::int64 a, tensorflow::int64 b, - tensorflow::int64 c, char** values, tensorflow::int32 values_count, - tensorflow::int32* values_primitive_type_size_in_bytes); - -extern void __xla_cpu_runtime_KeyValueSortS16( - tensorflow::int16* keys, tensorflow::int64 a, tensorflow::int64 b, - tensorflow::int64 c, char** values, tensorflow::int32 values_count, - tensorflow::int32* values_primitive_type_size_in_bytes); - -extern void __xla_cpu_runtime_KeyValueSortU16( - tensorflow::uint16* keys, tensorflow::int64 a, tensorflow::int64 b, - tensorflow::int64 c, char** values, tensorflow::int32 values_count, - tensorflow::int32* values_primitive_type_size_in_bytes); - -extern void __xla_cpu_runtime_KeyValueSortF16( - Eigen::half* keys, tensorflow::int64 a, tensorflow::int64 b, - tensorflow::int64 c, char** values, tensorflow::int32 values_count, - tensorflow::int32* values_primitive_type_size_in_bytes); - -extern void __xla_cpu_runtime_KeyValueSortS32( - tensorflow::int32* keys, tensorflow::int64 a, tensorflow::int64 b, - tensorflow::int64 c, char** values, tensorflow::int32 values_count, - tensorflow::int32* values_primitive_type_size_in_bytes); - -extern void __xla_cpu_runtime_KeyValueSortU32( - tensorflow::uint32* keys, tensorflow::int64 a, tensorflow::int64 b, - tensorflow::int64 c, char** values, tensorflow::int32 values_count, - tensorflow::int32* values_primitive_type_size_in_bytes); - -extern void __xla_cpu_runtime_KeyValueSortF32( - float* keys, tensorflow::int64 a, tensorflow::int64 b, tensorflow::int64 c, - char** values, tensorflow::int32 values_count, - tensorflow::int32* values_primitive_type_size_in_bytes); - -extern void __xla_cpu_runtime_KeyValueSortS64( - tensorflow::int64* keys, tensorflow::int64 a, tensorflow::int64 b, - tensorflow::int64 c, char** values, tensorflow::int32 values_count, - tensorflow::int32* values_primitive_type_size_in_bytes); - -extern void __xla_cpu_runtime_KeyValueSortU64( - tensorflow::uint64* keys, tensorflow::int64 a, tensorflow::int64 b, - tensorflow::int64 c, char** values, tensorflow::int32 values_count, - tensorflow::int32* values_primitive_type_size_in_bytes); - -extern void __xla_cpu_runtime_KeyValueSortF64( - double* keys, tensorflow::int64 a, tensorflow::int64 b, tensorflow::int64 c, - char** values, tensorflow::int32 values_count, - tensorflow::int32* values_primitive_type_size_in_bytes); + tensorflow::int32* values_primitive_type_size_in_bytes, char* run_options, + tensorflow::int64* prof_counters, + void (*less_than)(char*, char*, char**, char**, tensorflow::int64*)); } #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_KEY_VALUE_SORT_H_ diff --git a/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.cc b/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.cc index 1ed743afc30af7c7ff38c7d2a738f2e376270952..1f7204e67a413efabd34cd7d88ced4c82ee7a5df 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.cc +++ b/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.cc @@ -20,6 +20,10 @@ limitations under the License. #include "tensorflow/core/platform/dynamic_annotations.h" #include "tensorflow/core/platform/types.h" +#if defined(TENSORFLOW_USE_CUSTOM_CONTRACTION_KERNEL) +#include "tensorflow/core/kernels/eigen_contraction_kernel.h" +#endif + using tensorflow::int32; using tensorflow::int64; diff --git a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc index 296f39a4853f2d3f7030209a921001e92c39d609..9c2685674fbc133de1220caef81ac3b60a1c0f7c 100644 --- a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc +++ b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc @@ -116,13 +116,26 @@ SimpleOrcJIT::SimpleOrcJIT(const llvm::TargetOptions& target_options, orc_jit_memory_mapper::GetInstance()); result.Resolver = symbol_resolver_; return result; + }, + /*NotifyLoaded=*/ + llvm::orc::LegacyRTDyldObjectLinkingLayer::NotifyLoadedFtor(), + /*NotifyFinalized=*/ + [this](VModuleKeyT, const llvm::object::ObjectFile& object, + const llvm::RuntimeDyld::LoadedObjectInfo& object_info) { + this->NotifyObjectFinalized(object, object_info); + }, + /*NotifyFreed=*/ + [this](VModuleKeyT, const llvm::object::ObjectFile& object) { + this->NotifyObjectFreed(object); }), compile_layer_(object_layer_, CompilerFunctor(target_machine_.get(), &disassembler_, opt_level, optimize_for_size, enable_fast_math, disable_expensive_passes, std::move(pre_optimization_hook), - std::move(post_optimization_hook))) { + std::move(post_optimization_hook))), + gdb_jit_event_listener_( + llvm::JITEventListener::createGDBRegistrationListener()) { VLOG(1) << "CPU target: " << target_machine_->getTargetCPU().str() << " features: " << target_machine_->getTargetFeatureString().str(); } @@ -147,6 +160,20 @@ llvm::JITSymbol SimpleOrcJIT::ResolveRuntimeSymbol(const std::string& name) { return symbol_info; } +void SimpleOrcJIT::NotifyObjectFinalized( + const llvm::object::ObjectFile& object, + const llvm::RuntimeDyld::LoadedObjectInfo& object_info) { + uint64_t key = static_cast( + reinterpret_cast(object.getData().data())); + gdb_jit_event_listener_->notifyObjectLoaded(key, object, object_info); +} + +void SimpleOrcJIT::NotifyObjectFreed(const llvm::object::ObjectFile& object) { + uint64_t key = static_cast( + reinterpret_cast(object.getData().data())); + gdb_jit_event_listener_->notifyFreeingObject(key); +} + SimpleOrcJIT::VModuleKeyT SimpleOrcJIT::AddModule( std::unique_ptr module) { auto key = execution_session_.allocateVModule(); @@ -213,18 +240,7 @@ bool RegisterKnownJITSymbols() { REGISTER_CPU_RUNTIME_SYMBOL(ParallelForkJoin); REGISTER_CPU_RUNTIME_SYMBOL(ReleaseInfeedBufferAfterDequeue); REGISTER_CPU_RUNTIME_SYMBOL(ReleaseOutfeedBufferAfterPopulation); - REGISTER_CPU_RUNTIME_SYMBOL(KeyValueSortPRED); - REGISTER_CPU_RUNTIME_SYMBOL(KeyValueSortS8); - REGISTER_CPU_RUNTIME_SYMBOL(KeyValueSortU8); - REGISTER_CPU_RUNTIME_SYMBOL(KeyValueSortS16); - REGISTER_CPU_RUNTIME_SYMBOL(KeyValueSortU16); - REGISTER_CPU_RUNTIME_SYMBOL(KeyValueSortF16); - REGISTER_CPU_RUNTIME_SYMBOL(KeyValueSortS32); - REGISTER_CPU_RUNTIME_SYMBOL(KeyValueSortU32); - REGISTER_CPU_RUNTIME_SYMBOL(KeyValueSortF32); - REGISTER_CPU_RUNTIME_SYMBOL(KeyValueSortS64); - REGISTER_CPU_RUNTIME_SYMBOL(KeyValueSortU64); - REGISTER_CPU_RUNTIME_SYMBOL(KeyValueSortF64); + REGISTER_CPU_RUNTIME_SYMBOL(KeyValueSort); registry->Register("__gnu_f2h_ieee", reinterpret_cast(__gnu_f2h_ieee)); registry->Register("__gnu_h2f_ieee", reinterpret_cast(__gnu_h2f_ieee)); diff --git a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.h b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.h index 78406ba143570183aea09d79db3f9b708c21bf70..3307c2f93d796bbdcd49af7f68e9f6c388e402ca 100644 --- a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.h +++ b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.h @@ -21,6 +21,7 @@ limitations under the License. #include #include "llvm/ADT/Triple.h" +#include "llvm/ExecutionEngine/JITEventListener.h" #include "llvm/ExecutionEngine/Orc/Core.h" #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" @@ -99,6 +100,11 @@ class SimpleOrcJIT { private: llvm::JITSymbol ResolveRuntimeSymbol(const std::string& name); + void NotifyObjectFinalized( + const llvm::object::ObjectFile& object, + const llvm::RuntimeDyld::LoadedObjectInfo& object_info); + void NotifyObjectFreed(const llvm::object::ObjectFile& object); + std::vector module_keys_; std::unique_ptr target_machine_; const Disassembler disassembler_; @@ -107,6 +113,15 @@ class SimpleOrcJIT { std::shared_ptr symbol_resolver_; ObjLayerT object_layer_; CompileLayerT compile_layer_; + + // Non owning pointer to a JIT event listener that registers the JIT events + // with an attached GDB. + // + // Note: we get a pointer to this event listener using + // `createGDBRegistrationListener` which makes it look like we're supposed to + // free this, but the function is poorly named and really just returns a + // pointer to a static object. + llvm::JITEventListener* gdb_jit_event_listener_; }; } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_eigen_dot_operation_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_eigen_dot_operation_test.cc index f8f5f392da8ab3348e63185aecf7b639daacaa42..8b7f843582b697058fe328fe69990122d868ada4 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_eigen_dot_operation_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_eigen_dot_operation_test.cc @@ -16,7 +16,6 @@ limitations under the License. // Tests that we call into Eigen for dot operations as needed. #include -#include #include #include "absl/strings/str_cat.h" @@ -102,10 +101,10 @@ std::vector GetDotTestCases() { return result; } -INSTANTIATE_TEST_CASE_P(CpuEigenDotOperationTestInstantiation, - CpuEigenDotOperationTest, - ::testing::ValuesIn(GetDotTestCases()), - DotTestSpecToString); +INSTANTIATE_TEST_SUITE_P(CpuEigenDotOperationTestInstantiation, + CpuEigenDotOperationTest, + ::testing::ValuesIn(GetDotTestCases()), + DotTestSpecToString); } // namespace } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_intrinsic_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_intrinsic_test.cc index 9b10c49f4f547edfb2164f98c49cceb031148bdc..9078b8fd1ff6cb0ddac89d5fcd13a9ccfae07763 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_intrinsic_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_intrinsic_test.cc @@ -14,9 +14,9 @@ limitations under the License. ==============================================================================*/ #include -#include #include +#include "absl/strings/ascii.h" #include "absl/strings/str_cat.h" #include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h" #include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h" @@ -59,8 +59,9 @@ class CpuUnaryIntrinsicTest string features{spec.features.data(), spec.features.size()}; if (!features.empty()) { - std::replace_if(features.begin(), features.end(), - [](char c) { return c != '_' && !isalnum(c); }, '_'); + std::replace_if( + features.begin(), features.end(), + [](char c) { return c != '_' && !absl::ascii_isalnum(c); }, '_'); } else { features = ""; } @@ -140,10 +141,10 @@ IntrinsicTestSpec CpuUnaryIntrinsicTestCases[] = { HloOpcode::kLog, kTriple_android_arm, "", R"(CHECK: fadd fast <4 x float> )"}}; -INSTANTIATE_TEST_CASE_P(CpuUnaryIntrinsicTestInstantiation, - CpuUnaryIntrinsicTest, - ::testing::ValuesIn(CpuUnaryIntrinsicTestCases), - CpuUnaryIntrinsicTest::Name); +INSTANTIATE_TEST_SUITE_P(CpuUnaryIntrinsicTestInstantiation, + CpuUnaryIntrinsicTest, + ::testing::ValuesIn(CpuUnaryIntrinsicTestCases), + CpuUnaryIntrinsicTest::Name); } // namespace } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_key_value_sort_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_key_value_sort_test.cc index 3934c03a04c978009282b3cd0d39bacf9b12a356..762ee67db9a1b2a753c6ec5538dee1d13282942e 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_key_value_sort_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_key_value_sort_test.cc @@ -26,10 +26,16 @@ TEST_F(CpuKeyValueSortTest, SortR1) { const string hlo_text = R"( HloModule KeyValueSort +compare { + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + ROOT lt = pred[] less-than(p.0.lhs, p.0.rhs) +} + ENTRY main { a = f32[10] parameter(0) - ROOT result = f32[10] sort(f32[10] a), dimensions={0} + ROOT result = f32[10] sort(f32[10] a), dimensions={0}, to_apply=compare } )"; diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_noalias_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_noalias_test.cc index a7702c2aeeaff8a46a2c4f2785ccb873ea2c08e5..030bd41c2fc73eac41fe43c1acdf862d5dc97f98 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_noalias_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_noalias_test.cc @@ -75,8 +75,9 @@ TEST_F(CpuNoAliasTest, Concat) { // the buffers in the HLO module. We'll inspect these loads to ensure that // they have the expected alias information. llvm::Module ir_module("test", context); - llvm::Function* func = llvm::cast( - ir_module.getOrInsertFunction("test_fn", llvm::Type::getVoidTy(context))); + llvm::Function* func = llvm::dyn_cast( + ir_module.getOrInsertFunction("test_fn", llvm::Type::getVoidTy(context)) + .getCallee()); llvm::BasicBlock* bb = llvm::BasicBlock::Create(context, "body", func); llvm::IRBuilder<> b(bb); auto* zero = llvm::ConstantInt::get(llvm::Type::getInt32Ty(context), 0); diff --git a/tensorflow/compiler/xla/service/cpu/tiled_dot_emitter.cc b/tensorflow/compiler/xla/service/cpu/tiled_dot_emitter.cc new file mode 100644 index 0000000000000000000000000000000000000000..eb6c44b70ab34d0a294880b5de4fe0b3ba5e19e5 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/tiled_dot_emitter.cc @@ -0,0 +1,1014 @@ +/* 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/cpu/tiled_dot_emitter.h" + +#include "tensorflow/compiler/xla/service/cpu/vector_support_library.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h" +#include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" + +namespace xla { +namespace cpu { +namespace { + +using tensorflow::int64; + +// Provides tiled access to an in-memory rank 2 array. +class MemoryTile { + public: + // Constructs a MemoryTile that can operate on tiles consisting of + // `tile_size_along_major_dim` vectors from the matrix `matrix`, starting at + // `major_dim_offset` in the major dimension. The tile size along the minor + // dimension is the vector size, and that is implicitly determined by `vsl`. + MemoryTile(VectorSupportLibrary* vsl, llvm::IRBuilder<>* b, + llvm::Value* matrix, int64 matrix_size_along_minor_dim, + llvm::Value* major_dim_offset, int64 tile_size_along_major_dim) + : vsl_(vsl), b_(b) { + pointers_.reserve(tile_size_along_major_dim); + for (int64 i = 0; i < tile_size_along_major_dim; i++) { + llvm::Value* total_offset = + b->CreateMul(b->getInt64(matrix_size_along_minor_dim), + b->CreateAdd(b->getInt64(i), major_dim_offset)); + pointers_.push_back(vsl_->ComputeOffsetPointer(matrix, total_offset)); + } + } + + // Load a tile consisting of `tile_size_along_major_dim` vectors from position + // {major: `major_dim_offset`, minor: `minor_dim_offset`}. + // + // Note: `major_dim_offset` is a parameter to the constructor. + std::vector LoadTile(llvm::Value* minor_dim_offset) const { + std::vector result; + result.reserve(pointers_.size()); + for (const auto& pointer : pointers_) { + result.push_back(vsl_->LoadVector(pointer, minor_dim_offset)); + } + return result; + } + + // Stores `tile` to position {major: `major_dim_offset`, minor: + // `minor_dim_offset`}. + // + // Note: `major_dim_offset` is a parameter to the constructor. + void StoreTile(absl::Span tile, + llvm::Value* minor_dim_offset) const { + CHECK_EQ(tile.size(), pointers_.size()); + for (int64 i = 0; i < pointers_.size(); i++) { + vsl_->StoreVector(tile[i], pointers_[i], minor_dim_offset); + } + } + + // Loads a tile of size [`tile_size_along_major_dim`, + // `tile_size_along_middle_dim`] from position {major: `major_dim_offset`, + // minor: `minor_dim_offset`} and then broadcasts each element into a vector + // of size vsl_.vector_size(). The (i,j)'th element of the return value is + // the (i,j)'th element in the tile broadcasted into an LLVM vector. + // + // Note: `major_dim_offset` is a parameter to the constructor. + std::vector> LoadBroadcastTile( + llvm::Value* minor_dim_offset, int64 tile_size_along_middle_dim) const { + std::vector> result; + result.resize(pointers_.size()); + for (int64 i = 0; i < pointers_.size(); i++) { + for (int64 j = 0; j < tile_size_along_middle_dim; j++) { + result[i].push_back(vsl_->LoadBroadcast( + pointers_[i], b_->CreateAdd(minor_dim_offset, b_->getInt64(j)))); + } + } + return result; + } + + private: + VectorSupportLibrary* vsl_; + llvm::IRBuilder<>* b_; + std::vector pointers_; +}; + +// The base class for the classes representing the GEMV emitter configurations. +// +// The IR emitted (modulo the LLVM values representing the input and output +// buffers) by the row major and column major GEMV emitters should be a function +// of their configuration. This is important because their configuration is +// used as a key to cache the generated IR. +class GemvConfig { + public: + // Mixin for convenience. + template + struct User { + public: + PrimitiveType scalar_type() const { + return derived().config().scalar_type(); + } + int64 tile_rows() const { return derived().config().tile_rows(); } + int64 tile_cols() const { return derived().config().tile_cols(); } + int64 m() const { return derived().config().m(); } + int64 k() const { return derived().config().k(); } + int64 has_addend() const { return derived().config().has_addend(); } + + private: + const T& derived() const { return *static_cast(this); } + }; + + PrimitiveType scalar_type() const { return scalar_type_; } + int64 tile_rows() const { return tile_rows_; } + int64 tile_cols() const { return tile_cols_; } + int64 m() const { return m_; } + int64 k() const { return k_; } + bool has_addend() const { return has_addend_; } + + string GetCacheKey() const { + return absl::StrCat(name_, "_", PrimitiveType_Name(scalar_type()), "_", + tile_rows(), "_", tile_cols(), "_", m(), "_", k(), + has_addend() ? "_with_addend" : ""); + } + + protected: + explicit GemvConfig(string name, PrimitiveType scalar_type, int64 tile_rows, + int64 tile_cols, int64 m, int64 k, bool has_addend) + : name_(std::move(name)), + scalar_type_(scalar_type), + tile_rows_(tile_rows), + tile_cols_(tile_cols), + m_(m), + k_(k), + has_addend_(has_addend) {} + + private: + string name_; + PrimitiveType scalar_type_; + int64 tile_rows_; + int64 tile_cols_; + int64 m_; + int64 k_; + bool has_addend_; +}; + +// Computes a dot product between "[M,K]{0,1} lhs" with a [K,1] vector (the +// layout of the vector does not matter). This implementation uses a tiling +// scheme to improve performance. +// +// We logically separate the LHS matrix into four segments: +// +// +----------------------+---+ +// | | | +// | | | +// | A | B | +// | | | +// | | | +// | | | +// +----------------------+---+ +// | C | D | +// +----------------------+---+ +// +// where A is the largest submatrix of the LHS that can be evenly dividied into +// tiles. For each tile in A, assuming tile_rows_ == tile_cols_ == 4, we have: +// +// +---+---+---+---+ +--+--+--+--+ +// |M00|M10|M20|M30| |V0|V1|V2|V3| +// +---+---+---+---+ +--+--+--+--+ +// |M01|M11|M21|M31| and |V0|V1|V2|V3| +// +---+---+---+---+ +--+--+--+--+ +// |M02|M12|M22|M32| |V0|V1|V2|V3| +// +---+---+---+---+ +--+--+--+--+ +// |M03|M13|M23|M33| |V0|V1|V2|V3| +// +---+---+---+---+ +--+--+--+--+ +// +// (Legend: rows are horizontal and columns are vertical; and each column is one +// llvm::Value of a vector type) +// +// where: +// +// a. The left tile is from the column major left matrix. +// b. The right tile is an elementwise broadcast of a [V0, V1, V2, V3] +// vector loaded from the RHS vector. +// +// As we iterate through the column dimension, we compute the change to the +// result vector by an elementwise multiplication between the two tiles above +// followed by a reduction along the major dimension: +// +// +-----------------------------------+ +// | M00*V0 + M10*V1 + M20*V2 + M30*V3 | +// +-----------------------------------+ +// | M01*V0 + M11*V1 + M21*V2 + M31*V3 | +// Result[R:R+4] += +-----------------------------------+ +// | M02*V0 + M12*V1 + M22*V2 + M32*V3 | +// +-----------------------------------+ +// | M03*V0 + M13*V1 + M23*V2 + M33*V3 | +// +-----------------------------------+ +// +// Where R is the starting row for the tile. +// +// We have an inner epilogue loop to deal with the "C" submatrix and an outer +// epilogue loop to deal with the B,D submarix. +// +// TODO(sanjoy): We should investigate if using gather loads and scatter stores +// can be used here have the same inner loop for both column-major and row-major +// matrix-vector products. +class ColumnMajorMatrixVectorProductEmitter + : public GemvConfig::User { + public: + class Config : public GemvConfig { + public: + explicit Config(PrimitiveType scalar_type, int64 tile_rows, int64 tile_cols, + int64 m, int64 k, bool has_addend) + : GemvConfig(/*name=*/"col_major_gemv", scalar_type, + /*tile_rows=*/tile_rows, /*tile_cols=*/tile_cols, /*m=*/m, + /*k=*/k, /*has_addend=*/has_addend) {} + }; + + ColumnMajorMatrixVectorProductEmitter(const Config& config, llvm::Value* lhs, + llvm::Value* rhs, llvm::Value* addend, + llvm::Value* result, + llvm::IRBuilder<>* b) + : config_(config), + lhs_(lhs), + rhs_(rhs), + addend_(addend), + result_(result), + b_(b), + ksl_(b_), + vsl_(config.scalar_type(), /*vector_size=*/config.tile_rows(), b_, "") { + CHECK(tile_rows() > 0 && IsPowerOfTwo(static_cast(tile_rows()))); + CHECK(!has_addend() || addend != nullptr); + } + + void Emit(); + + const Config& config() const { return config_; } + + private: + void EmitOuterLoopBody(llvm::Value* column, int64 column_count, + bool is_first_column); + + MemoryTile GetLhsMemoryTile(llvm::Value* column_start, int64 column_count) { + return MemoryTile(&vsl_, b_, /*matrix=*/lhs_, + /*matrix_size_along_minor_dim=*/m(), + /*major_dim_offset=*/column_start, + /*tile_size_along_major_dim=*/column_count); + } + + // Load a tile of values from the RHS. For the RHS a "tile" is a contiguous + // sequence of `count` values, each one broadcasted to the vector width. + std::vector LoadRhsTile(llvm::Value* offset, int64 count) { + llvm::Value* base_pointer = vsl_.ComputeOffsetPointer(rhs_, offset); + std::vector result; + result.reserve(count); + for (int64 i = 0; i < count; i++) { + result.push_back(vsl_.LoadBroadcast(base_pointer, i)); + } + return result; + } + + void EmitInnerLoopTiled(MemoryTile* lhs_memory_tile, + const std::vector& rhs_tile, + int64 columns, bool is_first_column); + + void EmitInnerLoopEpilogue(llvm::Value* current_tile_col, int64 columns, + bool is_first_tiled_column); + + Config config_; + llvm::Value* lhs_; + llvm::Value* rhs_; + llvm::Value* addend_; + llvm::Value* result_; + llvm::IRBuilder<>* b_; + KernelSupportLibrary ksl_; + VectorSupportLibrary vsl_; +}; + +void ColumnMajorMatrixVectorProductEmitter::EmitOuterLoopBody( + llvm::Value* column, int64 column_count, bool is_first_column) { + MemoryTile lhs_memory_tile = GetLhsMemoryTile(/*column_start=*/column, + /*column_count=*/column_count); + + std::vector rhs_tile = + LoadRhsTile(column, /*count=*/column_count); + EmitInnerLoopTiled(&lhs_memory_tile, rhs_tile, + /*columns=*/column_count, is_first_column); + EmitInnerLoopEpilogue(column, /*columns=*/column_count, is_first_column); +} + +void ColumnMajorMatrixVectorProductEmitter::Emit() { + // See the comment on the class declaration for the algorithm used here. + int64 column_remainder = k() % tile_cols(); + int64 column_limit = k() - column_remainder; + + ksl_.For("dot.outer.tiled", + /*start=*/0, /*end=*/column_limit, /*step=*/tile_cols(), + [&](llvm::Value* column, bool is_first_column) { + EmitOuterLoopBody(column, tile_cols(), is_first_column); + }); + + if (column_remainder != 0) { + EmitOuterLoopBody(b_->getInt64(column_limit), column_remainder, + column_limit == 0); + } +} + +void ColumnMajorMatrixVectorProductEmitter::EmitInnerLoopTiled( + MemoryTile* lhs_memory_tile, const std::vector& rhs_tile, + int64 columns, bool is_first_column) { + int64 row_limit = m() - (m() % tile_rows()); + + ksl_.For("dot.inner.tiled", /*start=*/0, /*end=*/row_limit, + /*step=*/tile_rows(), [&](llvm::Value* row) { + std::vector lhs_tile = + lhs_memory_tile->LoadTile(/*minor_dim_offset=*/row); + llvm::Value* accumulator = + is_first_column ? (addend_ ? vsl_.LoadVector(addend_, row) + : vsl_.GetZeroVector()) + : vsl_.LoadVector(result_, row); + for (int i = 0; i < columns; i++) { + accumulator = vsl_.MulAdd(lhs_tile[i], rhs_tile[i], accumulator); + } + vsl_.StoreVector(accumulator, result_, row); + }); +} + +void ColumnMajorMatrixVectorProductEmitter::EmitInnerLoopEpilogue( + llvm::Value* current_tile_col, int64 columns, bool is_first_tiled_column) { + int64 row_start = m() - (m() % tile_rows()); + if (row_start == m()) { + return; + } + + llvm::Value* columns_llvm = b_->getInt64(columns); + + // for (col = current_tile_col; col < (columns + current_tile_col); col++) + // for (row = row_start, row < m_; row++) { + // result[row] += lhs[row, col] * rhs[col] + // // Also take into account that if col is 0 then result[row] is not + // // initialized. + // } + + ksl_.For( + "dot.inner.epilg.outer", /*start=*/current_tile_col, + /*end=*/b_->CreateAdd(columns_llvm, current_tile_col), + /*step=*/1, /*peel_first_iteration=*/false, + [&](llvm::Value* col, llvm::Value* is_first_scalar_col) { + llvm::Value* rhs_element = vsl_.LoadScalar(rhs_, col); + llvm::Value* total_offset = b_->CreateMul(col, b_->getInt64(m())); + llvm::Value* lhs_base_pointer = + vsl_.ComputeOffsetPointer(lhs_, total_offset); + ksl_.For( + "dot.inner.epilg.inner", /*start=*/row_start, /*end=*/m(), + /*step=*/1, [&](llvm::Value* scalar_row) { + llvm::Value* product = vsl_.Mul( + vsl_.LoadScalar(lhs_base_pointer, scalar_row), rhs_element); + llvm::Value* setting_result_first_time = b_->CreateAnd( + is_first_scalar_col, b_->getInt1(is_first_tiled_column)); + ksl_.If( + setting_result_first_time, + /*true_block_generator=*/ + [&]() { + if (addend_) { + vsl_.StoreScalar( + vsl_.Add(vsl_.LoadScalar(addend_, scalar_row), + product), + result_, scalar_row); + } else { + vsl_.StoreScalar(product, result_, scalar_row); + } + }, + /*false_block_generator=*/ + [&]() { + vsl_.StoreScalar( + vsl_.Add(vsl_.LoadScalar(result_, scalar_row), product), + result_, scalar_row); + }); + }); + }); +} + +// Computes a dot product between "[M,K]{1,0} lhs" with a [K,1] vector (the +// layout of the vector does not matter). This implementation uses a tiling +// scheme to improve performance. +// +// We logically separate the LHS matrix into four segments: +// +// +----------------------+---+ +// | | | +// | | | +// | A | B | +// | | | +// | | | +// | | | +// +----------------------+---+ +// | C | D | +// +----------------------+---+ +// +// where A is the largest submatrix of the LHS that can be evenly dividied into +// tiles. For each tile in A, assuming tile_rows_ == tile_cols_ == 4, we have: +// +// +---+---+---+---+ +// |M00|M10|M20|M30| +// +---+---+---+---+ +--+--+--+--+ +// |M01|M11|M21|M31| and |V0|V1|V2|V3| +// +---+---+---+---+ +--+--+--+--+ +// |M02|M12|M22|M32| +// +---+---+---+---+ +// |M03|M13|M23|M33| +// +---+---+---+---+ +// +// (Legend: rows are horizontal and columns are vertical; and each row is one +// llvm::Value of a vector type) +// +// where: +// +// a. The left tile is loaded from the row major left matrix. +// b. The right vector is loaded from the RHS vector. +// +// We keep 4 vector accumulators accumulating the following four vector +// expressions as we iterate over the row dimension: +// +// +------+------+------+------+ +// |M0I*V0|M1I*V1|M2I*V2|M3I*V3| for I in [0,4) +// +------+------+------+------+ +// +// In the end we do a horizontal reduction over these 4 vector accumulators to +// get 4 values in the result vector. +// +// We have an inner epilogue loop to deal with the "B" sub-matrix and an outer +// epilogue loop to deal with the C,D submatrix. +class RowMajorMatrixVectorProductEmitter + : public GemvConfig::User { + public: + class Config : public GemvConfig { + public: + explicit Config(PrimitiveType scalar_type, int64 tile_rows, int64 tile_cols, + int64 m, int64 k, bool has_addend) + : GemvConfig(/*name=*/"row_major_gemv", scalar_type, + /*tile_rows=*/tile_rows, /*tile_cols=*/tile_cols, /*m=*/m, + /*k=*/k, /*has_addend=*/has_addend) {} + }; + + RowMajorMatrixVectorProductEmitter(const Config& config, llvm::Value* lhs, + llvm::Value* rhs, llvm::Value* addend, + llvm::Value* result, llvm::IRBuilder<>* b) + : config_(config), + lhs_(lhs), + rhs_(rhs), + addend_(addend), + result_(result), + b_(b), + ksl_(b_), + vsl_(scalar_type(), /*vector_size=*/tile_cols(), b_, "") { + CHECK(tile_cols() > 0 && IsPowerOfTwo(static_cast(tile_cols()))); + CHECK(!has_addend() || addend != nullptr); + } + + void Emit(); + + const Config& config() const { return config_; } + + private: + MemoryTile GetLhsMemoryTile(llvm::Value* row_start, int64 row_count) { + return MemoryTile(&vsl_, b_, /*matrix=*/lhs_, + /*matrix_size_along_minor_dim=*/k(), + /*major_dim_offset=*/row_start, + /*tile_size_along_major_dim=*/row_count); + } + + void EmitOuterLoopBody(llvm::Value* row, int64 row_count); + + void EmitInnerLoopTiled(MemoryTile* lhs_memory_tile, int64 rows, + std::vector* vector_accumulators); + + void EmitInnerLoopEpilogue(llvm::Value* current_tile_row, int64 rows, + std::vector* scalar_accumulators); + + Config config_; + llvm::Value* lhs_; + llvm::Value* rhs_; + llvm::Value* addend_; + llvm::Value* result_; + llvm::IRBuilder<>* b_; + KernelSupportLibrary ksl_; + VectorSupportLibrary vsl_; +}; + +void RowMajorMatrixVectorProductEmitter::EmitOuterLoopBody(llvm::Value* row, + int64 row_count) { + MemoryTile lhs_memory_tile = GetLhsMemoryTile(/*row_start=*/row, + /*row_count=*/row_count); + std::vector vector_accumulators; + std::vector scalar_accumulators; + for (int i = 0; i < row_count; i++) { + vector_accumulators.emplace_back(&vsl_, vsl_.GetZeroVector()); + scalar_accumulators.emplace_back(&vsl_, vsl_.GetZeroScalar()); + } + EmitInnerLoopTiled(&lhs_memory_tile, /*rows=*/row_count, + &vector_accumulators); + EmitInnerLoopEpilogue(/*current_tile_row=*/row, /*rows=*/row_count, + &scalar_accumulators); + + std::vector accumulator_values; + std::transform( + vector_accumulators.begin(), vector_accumulators.end(), + std::back_inserter(accumulator_values), + [](const VectorVariable& vector_var) { return vector_var.Get(); }); + + std::vector horizontal_sums; + if (row_count == vsl_.vector_size()) { + if (addend_) { + horizontal_sums = vsl_.ComputeHorizontalSums( + std::move(accumulator_values), vsl_.LoadVector(addend_, row)); + } else { + horizontal_sums = + vsl_.ComputeHorizontalSums(std::move(accumulator_values)); + } + } else { + horizontal_sums = vsl_.ComputeHorizontalSums(std::move(accumulator_values)); + } + + for (int i = 0; i < row_count; i++) { + llvm::Value* result_value = + vsl_.Add(horizontal_sums[i], scalar_accumulators[i].Get()); + llvm::Value* offset = b_->CreateAdd(b_->getInt64(i), row); + if (addend_ && row_count != vsl_.vector_size()) { + result_value = vsl_.Add(vsl_.LoadScalar(addend_, offset), result_value); + } + vsl_.StoreScalar(result_value, result_, offset); + } +} + +void RowMajorMatrixVectorProductEmitter::Emit() { + // See the comment on the class declaration for the algorithm used here. + int64 row_remainder = m() % tile_rows(); + int64 row_limit = m() - row_remainder; + + ksl_.For("dot.outer.tiled", + /*start=*/0, /*end=*/row_limit, /*step=*/tile_rows(), + [&](llvm::Value* row) { EmitOuterLoopBody(row, tile_rows()); }); + + if (row_remainder != 0) { + EmitOuterLoopBody(b_->getInt64(row_limit), row_remainder); + } +} + +void RowMajorMatrixVectorProductEmitter::EmitInnerLoopTiled( + MemoryTile* lhs_memory_tile, int64 rows, + std::vector* vector_accumulators) { + int64 column_limit = k() - (k() % tile_cols()); + + ksl_.For("dot.inner.tiled", /*start=*/0, /*end=*/column_limit, + /*step=*/tile_cols(), [&](llvm::Value* col) { + std::vector lhs_tile = + lhs_memory_tile->LoadTile(/*minor_dim_offset=*/col); + llvm::Value* rhs_value = vsl_.LoadVector(rhs_, col); + for (int i = 0; i < rows; i++) { + llvm::Value* old_sum = (*vector_accumulators)[i].Get(); + (*vector_accumulators)[i].Set( + vsl_.Add(old_sum, vsl_.Mul(rhs_value, lhs_tile[i]))); + } + }); +} + +void RowMajorMatrixVectorProductEmitter::EmitInnerLoopEpilogue( + llvm::Value* current_tile_row, int64 rows, + std::vector* scalar_accumulators) { + int64 column_start = k() - (k() % tile_cols()); + if (column_start == k()) { + return; + } + + for (int r = 0; r < rows; r++) { + llvm::Value* total_offset = b_->CreateMul( + b_->CreateAdd(b_->getInt64(r), current_tile_row), b_->getInt64(k())); + llvm::Value* lhs_base_pointer = + vsl_.ComputeOffsetPointer(lhs_, total_offset); + ksl_.For("dot.inner.epilg.inner", /*start=*/column_start, /*end=*/k(), + /*step=*/1, [&](llvm::Value* scalar_col) { + llvm::Value* product = + vsl_.Mul(vsl_.LoadScalar(lhs_base_pointer, scalar_col), + vsl_.LoadScalar(rhs_, scalar_col)); + llvm::Value* old_value = (*scalar_accumulators)[r].Get(); + (*scalar_accumulators)[r].Set(vsl_.Add(old_value, product)); + }); + } +} + +// This class implements a tiled matrix multiplication algorithm, intended for +// multiplying small matrices that don't need cache tiling. +// +// In the future this can be used as the innermost GEBP loop in a GEMM kernel as +// described in "Goto, Kazushige, and Robert A. Geijn. "Anatomy of +// high-performance matrix multiplication." ACM Transactions on Mathematical +// Software (TOMS) 34.3 (2008): 12.". +// +// This only supports canonical dot operations (i.e. where the lhs contraction +// dimension is 1 and the rhs contraction dimension is 0) over row major +// matrices. +class TiledSmallGemmEmitter { + public: + // Describe the dimensions of the kernel. + class Dimensions { + public: + explicit Dimensions(int64 m, int64 k, int64 n) : m_(m), k_(k), n_(n) {} + + int64 m() const { return m_; } + int64 k() const { return k_; } + int64 n() const { return n_; } + + string ToString() const { return absl::StrCat(m(), "x", k(), "x", n()); } + + private: + const int64 m_; + const int64 k_; + const int64 n_; + }; + + // Represents the configuration of the emitter. The LLVM IR emitted by the + // emitter, modulo the LLVM values holding the input and output buffers, must + // be a function of the instance of `Config` passed to it. + // + // `dims` holds the matrix multiplication dimensions. + // + // `max_vectorization_width` is the maximum vector width (i.e. the width of + // the largest vector register we will use). This can be larger than the + // largest vector register supported by the machine -- LLVM will legalize + // these large vector widths into legally sized vectors. + // + // `max_vector_count` is the maximum number of vectors of size + // `max_vectorization_width` that we will attempt to process at once. + // + // `min_vectorization_width` is the smallest vector width the emitter will use + // -- below that it will devolve to using a scalar loop. + // + // The innermost reduction loop executes the matrix multiply in tiles of size + // [`tile_size_m`, `tile_size_k`] from the LHS and [`tile_size_k`, + // ] in the RHS. + class Config { + public: + explicit Config(PrimitiveType scalar_type, Dimensions dims, + int64 max_vectorization_width, int64 max_vector_count, + int64 min_vectorization_width, int64 tile_size_m, + int64 tile_size_k) + : scalar_type_(scalar_type), + dims_(dims), + max_vectorization_width_(max_vectorization_width), + max_vector_count_(max_vector_count), + min_vectorization_width_(min_vectorization_width), + tile_size_m_(tile_size_m), + tile_size_k_(tile_size_k) {} + + string GetCacheKey() const { + return absl::StrCat("gemm_", PrimitiveType_Name(scalar_type()), "_", + dims().ToString(), "_", max_vectorization_width(), + "_", min_vectorization_width(), "_", tile_size_m(), + "_", tile_size_k()); + } + + PrimitiveType scalar_type() const { return scalar_type_; } + Dimensions dims() const { return dims_; } + int64 max_vectorization_width() const { return max_vectorization_width_; } + int64 max_vector_count() const { return max_vector_count_; } + int64 min_vectorization_width() const { return min_vectorization_width_; } + + int64 tile_size_m() const { return tile_size_m_; } + int64 tile_size_k() const { return tile_size_k_; } + + private: + PrimitiveType scalar_type_; + Dimensions dims_; + int64 max_vectorization_width_; + int64 max_vector_count_; + int64 min_vectorization_width_; + int64 tile_size_m_; + int64 tile_size_k_; + }; + + // Creates an instance of TiledSmallGemmEmitter that matrix-multiplies + // `lhs` with `rhs` and stores the result in `result`. + explicit TiledSmallGemmEmitter(Config config, llvm::Value* lhs, + llvm::Value* rhs, llvm::Value* result, + llvm::IRBuilder<>* b) + : lhs_(lhs), + rhs_(rhs), + result_(result), + config_(config), + b_(b), + ksl_(b_) { + CHECK(max_vectorization_width() > 0 && + IsPowerOfTwo(static_cast(max_vectorization_width()))); + CHECK_GT(max_vector_count(), 0); + CHECK(min_vectorization_width() > 0 && + IsPowerOfTwo(static_cast(min_vectorization_width()))); + CHECK_GE(max_vectorization_width(), min_vectorization_width()); + CHECK_GT(tile_size_k(), 0); + } + + void Emit(); + + private: + // The HandleResiduesOnX helpers split the iteration space for dimension X + // into a multiple of the tile size on dimension X and an epilogue. These + // helpers ultimately call into `EmitTiledGemm` for emitting the + // tiled GEMM kernel. + + void HandleResiduesOnN(); + void HandleResiduesOnK(VectorSupportLibrary* vsl, llvm::Value* n_start, + llvm::Value* n_end); + void HandleResiduesOnM(VectorSupportLibrary* vsl, int64 tile_size_k, + llvm::Value* k_start, llvm::Value* k_end, + llvm::Value* n_start, llvm::Value* n_end); + + // This emits a tiled GEMM kernel. For a detailed description see the comment + // on the implementation. + void EmitTiledGemm(VectorSupportLibrary* vsl, int64 tile_size_k, + llvm::Value* k_start, llvm::Value* k_end, + llvm::Value* n_start, llvm::Value* n_end, + int64 tile_size_m, llvm::Value* m_start, + llvm::Value* m_end); + + llvm::Value* GetInt64(int64 value) { return b_->getInt64(value); } + + Config config() const { return config_; } + Dimensions dims() const { return config().dims(); } + + int64 max_vectorization_width() const { + return config().max_vectorization_width(); + } + int64 max_vector_count() const { return config().max_vector_count(); } + int64 min_vectorization_width() const { + return config().min_vectorization_width(); + } + int64 tile_size_m() const { return config().tile_size_m(); } + int64 tile_size_k() const { return config().tile_size_k(); } + PrimitiveType scalar_type() const { return config().scalar_type(); } + + llvm::Value* lhs_; + llvm::Value* rhs_; + llvm::Value* result_; + Config config_; + + llvm::IRBuilder<>* b_; + KernelSupportLibrary ksl_; +}; + +void TiledSmallGemmEmitter::Emit() { HandleResiduesOnN(); } + +void TiledSmallGemmEmitter::HandleResiduesOnN() { + // We can only iterate the `n` dimension for an extent that is divisible by + // the vectorization width. So we emit an outer loop that first processes the + // largest extent in `n` that is divisible by max_vectorization_width, then + // the largest remaining extent that is divisible by max_vectorization_width / + // 2 etc. + + int64 current_vectorization_width = + max_vector_count() * max_vectorization_width(); + int64 current_vector_count = max_vector_count(); + + int64 n_start = 0; + while (n_start != dims().n() && + current_vectorization_width >= min_vectorization_width()) { + int64 n_end = dims().n() - (dims().n() % current_vectorization_width); + if (n_start != n_end) { + VectorSupportLibrary vsl(scalar_type(), current_vectorization_width, b_, + "gemm"); + HandleResiduesOnK(&vsl, GetInt64(n_start), GetInt64(n_end)); + n_start = n_end; + } + if (current_vector_count == 1) { + current_vectorization_width /= 2; + } else { + current_vector_count--; + current_vectorization_width = + current_vector_count * max_vectorization_width(); + } + } + + if (n_start != dims().n()) { + VectorSupportLibrary vsl(scalar_type(), 1, b_, "gemm"); + ksl_.For("epi.n", n_start, dims().n(), 1, [&](llvm::Value* n_i) { + llvm::Value* n_i_next = b_->CreateAdd(n_i, b_->getInt64(1)); + HandleResiduesOnK(&vsl, n_i, n_i_next); + }); + } +} + +void TiledSmallGemmEmitter::HandleResiduesOnK(VectorSupportLibrary* vsl, + llvm::Value* n_start, + llvm::Value* n_end) { + int64 k_start = 0; + int64 k_end = dims().k() - (dims().k() % tile_size_k()); + if (k_end != k_start) { + HandleResiduesOnM(vsl, tile_size_k(), GetInt64(k_start), GetInt64(k_end), + n_start, n_end); + k_start = k_end; + } + + if (k_start != dims().k()) { + HandleResiduesOnM(vsl, dims().k() - k_start, GetInt64(k_start), + GetInt64(dims().k()), n_start, n_end); + } +} + +void TiledSmallGemmEmitter::HandleResiduesOnM( + VectorSupportLibrary* vsl, int64 tile_size_k, llvm::Value* k_start, + llvm::Value* k_end, llvm::Value* n_start, llvm::Value* n_end) { + const int64 m_end = dims().m() - dims().m() % tile_size_m(); + EmitTiledGemm(vsl, tile_size_k, k_start, k_end, n_start, n_end, tile_size_m(), + GetInt64(0), GetInt64(m_end)); + + if (m_end != dims().m()) { + EmitTiledGemm(vsl, tile_size_k, k_start, k_end, n_start, n_end, + dims().m() - m_end, GetInt64(m_end), GetInt64(dims().m())); + } +} + +// The loop structure is: +// +// Iterate over dimension M as m: +// Iterate over dimension N as n: +// Iterate over dimension K as k: +// OutputTile[m,n] += Dot(LhsTile[m,k], RhsTile[k,n]) +// +// I.e. a just a tiled version of a "naive" GEMM. +// +// The tiling scheme is as follows: +// +// Let the LHS be: +// +// +----+----+----+ +// | a0 | b0 | c0 | . +// +----+----+----+ . +// | a1 | b1 | c1 | . +// +----+----+----+ +// .. .. +// +// and the RHS be: +// +// +----+----+----+----+ +// | p0 | p1 | p2 | p3 | . +// +----+----+----+----+ . +// | q0 | q1 | q2 | q3 | . +// +----+----+----+----+ +// | r0 | r1 | r2 | r3 | . +// +----+----+----+----+ . +// ...... ...... +// +// and let tile_size_m=2, tile_size_k=3 and the vector width (implicitly denoted +// by `vsl`) be 4. Then we want to matrix multiply this tile to get a [2,4] +// matrix that we can increment the result matrix by. +// +// First broadcast the rows row in LHS to 3 vectors of width 4, giving us a rank +// 3 array, L, of dimension [2,3,4]: +// +// L[0,_,_] * L[1,_,_] +// * +// +----+----+----+----+ * +----+----+----+----+ +// | a0 | a0 | a0 | a0 | * | a1 | a1 | a1 | a1 | +// +----+----+----+----+ * +----+----+----+----+ +// | b0 | b0 | b0 | b0 | * | b1 | b1 | b1 | b1 | +// +----+----+----+----+ * +----+----+----+----+ +// | c0 | c0 | c0 | c0 | * | c1 | c1 | c1 | c1 | +// +----+----+----+----+ * +----+----+----+----+ +// +// +// Then we FMA L[0,_,_] with the RHS to get the first row of the result and +// L[1,_,_] with the RHS to get the second row of the result. For example, +// L[0,_,_] is computed as: +// +// +----+----+----+----+ +----+----+----+----+ +// | a0 | a0 | a0 | a0 | * | p0 | p1 | p2 | p3 | + +// +----+----+----+----+ +----+----+----+----+ +// +// +----+----+----+----+ +----+----+----+----+ +// | b0 | b0 | b0 | b0 | * | q0 | q1 | q2 | q3 | + +// +----+----+----+----+ +----+----+----+----+ +// +// +----+----+----+----+ +----+----+----+----+ +// | c0 | c0 | c0 | c0 | * | r0 | r1 | r2 | r3 | +// +----+----+----+----+ +----+----+----+----+ +// +// to get: +// +// +-------------------+-------------------+-------------------+--------- +// | a0*p0+b0*q0+c0*r0 | a0*p1+b0*q1+c0*r1 | a0*p2+b0*q2+c0*r2 | ... +// +-------------------+-------------------+-------------------+--------- +void TiledSmallGemmEmitter::EmitTiledGemm( + VectorSupportLibrary* vsl, int64 tile_size_k, llvm::Value* k_start, + llvm::Value* k_end, llvm::Value* n_start, llvm::Value* n_end, + int64 tile_size_m, llvm::Value* m_start, llvm::Value* m_end) { + ksl_.For("dot.m", m_start, m_end, tile_size_m, [&](llvm::Value* m_i) { + MemoryTile result_memory_tile(vsl, b_, /*matrix=*/result_, + /*matrix_size_along_minor_dim=*/dims().n(), + /*major_dim_offset=*/m_i, + /*tile_size_along_major_dim=*/tile_size_m); + MemoryTile lhs_memory_tile(vsl, b_, /*matrix=*/lhs_, + /*matrix_size_along_minor_dim=*/dims().k(), + /*major_dim_offset=*/m_i, + /*tile_size_along_major_dim=*/tile_size_m); + ksl_.For( + "dot.n", n_start, n_end, vsl->vector_size(), [&](llvm::Value* n_i) { + TileVariable result_tile_var(vsl, result_memory_tile.LoadTile(n_i)); + ksl_.For("dot.k", k_start, k_end, tile_size_k, [&](llvm::Value* k_i) { + MemoryTile rhs_memory_tile(vsl, b_, rhs_, dims().n(), k_i, + tile_size_k); + std::vector> lhs_tile = + lhs_memory_tile.LoadBroadcastTile(k_i, tile_size_k); + std::vector rhs_tile = rhs_memory_tile.LoadTile(n_i); + std::vector result_tile = result_tile_var.Get(); + for (int64 r_m_i = 0; r_m_i < tile_size_m; r_m_i++) { + for (int64 r_k_i = 0; r_k_i < tile_size_k; r_k_i++) { + result_tile[r_m_i] = + vsl->MulAdd(lhs_tile[r_m_i][r_k_i], rhs_tile[r_k_i], + result_tile[r_m_i]); + } + } + result_tile_var.Set(result_tile); + }); + + result_memory_tile.StoreTile(result_tile_var.Get(), n_i); + }); + }); +} + +} // namespace + +void EmitRowMajorGemv(PrimitiveType scalar_type, int64 tile_rows, + int64 tile_cols, int64 m, int64 k, llvm::Value* lhs, + llvm::Value* rhs, llvm::Value* addend, + llvm::Value* result, llvm::IRBuilder<>* b, + bool enable_fast_math, bool optimize_for_size) { + RowMajorMatrixVectorProductEmitter::Config config( + /*scalar_type=*/scalar_type, + /*tile_rows=*/tile_rows, /*tile_cols=*/tile_cols, + /*m=*/m, /*k=*/k, /*has_addend=*/addend != nullptr); + + KernelSupportLibrary::EmitAndCallOutlinedKernel( + /*enable_fast_math=*/enable_fast_math, + /*optimize_for_size=*/optimize_for_size, b, config.GetCacheKey(), lhs, + rhs, addend, result, + [&](llvm::Value* lhs, llvm::Value* rhs, llvm::Value* addend, + llvm::Value* result) { + RowMajorMatrixVectorProductEmitter emitter(config, lhs, rhs, addend, + result, b); + emitter.Emit(); + }); +} + +void EmitColumnMajorGemv(PrimitiveType scalar_type, int64 tile_rows, + int64 tile_cols, int64 m, int64 k, llvm::Value* lhs, + llvm::Value* rhs, llvm::Value* addend, + llvm::Value* result, llvm::IRBuilder<>* b, + bool enable_fast_math, bool optimize_for_size) { + ColumnMajorMatrixVectorProductEmitter::Config config( + /*scalar_type=*/scalar_type, + /*tile_rows=*/tile_rows, /*tile_cols=*/tile_cols, + /*m=*/m, /*k=*/k, /*has_addend=*/addend != nullptr); + + KernelSupportLibrary::EmitAndCallOutlinedKernel( + /*enable_fast_math=*/enable_fast_math, + /*optimize_for_size=*/optimize_for_size, b, config.GetCacheKey(), lhs, + rhs, addend, result, + [&](llvm::Value* lhs, llvm::Value* rhs, llvm::Value* addend, + llvm::Value* result) { + ColumnMajorMatrixVectorProductEmitter emitter(config, lhs, rhs, addend, + result, b); + emitter.Emit(); + }); +} + +void EmitSmallGemm(PrimitiveType scalar_type, int64 m, int64 k, int64 n, + int64 max_vectorization_width, int64 max_vector_count, + int64 min_vectorization_width, int64 tile_size_m, + int64 tile_size_k, llvm::Value* lhs, llvm::Value* rhs, + llvm::Value* result, llvm::IRBuilder<>* b, + bool enable_fast_math, bool optimize_for_size) { + TiledSmallGemmEmitter::Config config( + /*scalar_type=*/scalar_type, + TiledSmallGemmEmitter::Dimensions{/*m=*/m, /*k=*/k, /*n=*/n}, + /*max_vectorization_width=*/max_vectorization_width, + /*max_vector_count=*/max_vector_count, + /*min_vectorization_width=*/min_vectorization_width, + /*tile_size_m=*/tile_size_m, /*tile_size_k=*/tile_size_k); + + KernelSupportLibrary::EmitAndCallOutlinedKernel( + /*enable_fast_math=*/enable_fast_math, + /*optimize_for_size=*/optimize_for_size, b, config.GetCacheKey(), lhs, + rhs, result, + [&](llvm::Value* lhs, llvm::Value* rhs, llvm::Value* result) { + TiledSmallGemmEmitter small_gemm_emitter(config, /*lhs=*/lhs, + /*rhs=*/rhs, + /*result=*/result, b); + small_gemm_emitter.Emit(); + }); +} + +} // namespace cpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/tiled_dot_emitter.h b/tensorflow/compiler/xla/service/cpu/tiled_dot_emitter.h new file mode 100644 index 0000000000000000000000000000000000000000..0a82326cc3704bce8c122261383249c60eda1f3a --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/tiled_dot_emitter.h @@ -0,0 +1,55 @@ +/* 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_CPU_TILED_DOT_EMITTER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_TILED_DOT_EMITTER_H_ + +#include "llvm/IR/IRBuilder.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { +namespace cpu { + +// These routines emit LLVM IR implementing tiled GEMM and GEMV routines. + +void EmitRowMajorGemv(PrimitiveType scalar_type, tensorflow::int64 tile_rows, + tensorflow::int64 tile_cols, tensorflow::int64 m, + tensorflow::int64 k, llvm::Value* lhs, llvm::Value* rhs, + llvm::Value* addend, llvm::Value* result, + llvm::IRBuilder<>* b, bool enable_fast_math, + bool optimize_for_size); + +void EmitColumnMajorGemv(PrimitiveType scalar_type, tensorflow::int64 tile_rows, + tensorflow::int64 tile_cols, tensorflow::int64 m, + tensorflow::int64 k, llvm::Value* lhs, + llvm::Value* rhs, llvm::Value* addend, + llvm::Value* result, llvm::IRBuilder<>* b, + bool enable_fast_math, bool optimize_for_size); + +void EmitSmallGemm(PrimitiveType scalar_type, tensorflow::int64 m, + tensorflow::int64 k, tensorflow::int64 n, + tensorflow::int64 max_vectorization_width, + tensorflow::int64 max_vector_count, + tensorflow::int64 min_vectorization_width, + tensorflow::int64 tile_size_m, tensorflow::int64 tile_size_k, + llvm::Value* lhs, llvm::Value* rhs, llvm::Value* result, + llvm::IRBuilder<>* b, bool enable_fast_math, + bool optimize_for_size); + +} // namespace cpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_TILED_DOT_EMITTER_H_ diff --git a/tensorflow/compiler/xla/service/cpu/vector_support_library.h b/tensorflow/compiler/xla/service/cpu/vector_support_library.h index 5690d2be2fe3e21c96b51a5226e0b29148217fd1..c444fd7d4aa88fa21b1aa2b2f058bd689b234b15 100644 --- a/tensorflow/compiler/xla/service/cpu/vector_support_library.h +++ b/tensorflow/compiler/xla/service/cpu/vector_support_library.h @@ -114,6 +114,9 @@ class VectorSupportLibrary { // raison d'etre) less cluttered. llvm::Value* FCmpEQMask(llvm::Value* lhs, llvm::Value* rhs); + llvm::Value* FCmpEQMask(llvm::Value* lhs, const llvm::APFloat& rhs) { + return FCmpEQMask(lhs, GetConstantFloat(lhs->getType(), rhs)); + } llvm::Value* FCmpULEMask(llvm::Value* lhs, llvm::Value* rhs); llvm::Value* FCmpOLTMask(llvm::Value* lhs, llvm::Value* rhs); llvm::Value* FCmpOLTMask(llvm::Value* lhs, const llvm::APFloat& rhs) { diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h index 2132468b9067ad4d5644d6cf3908a488a20ced05..af039776b7f157a407f8f4cf3d7cabdbee45b014 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h @@ -105,9 +105,11 @@ class DfsHloVisitorBase { } virtual Status HandleConvolution(HloInstructionPtr hlo) = 0; virtual Status HandleFft(HloInstructionPtr fft) = 0; + virtual Status HandleTriangularSolve(HloInstructionPtr hlo) = 0; virtual Status HandleAllReduce(HloInstructionPtr hlo) = 0; virtual Status HandleAllToAll(HloInstructionPtr hlo) = 0; virtual Status HandleCollectivePermute(HloInstructionPtr hlo) = 0; + virtual Status HandleReplicaId(HloInstructionPtr hlo) = 0; virtual Status HandleGetDimensionSize(HloInstructionPtr hlo) = 0; virtual Status HandleCompare(HloInstructionPtr hlo) { return HandleElementwiseBinary(hlo); 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 680dd256bb15bd3a9eaff7241174c1d2833002c6..341bb37b8355e9987a0331d0a66bb8fe87f019cf 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h @@ -91,6 +91,9 @@ class DfsHloVisitorWithDefaultBase Status HandleFft(HloInstructionPtr fft) override { return DefaultAction(fft); } + Status HandleTriangularSolve(HloInstructionPtr hlo) override { + return DefaultAction(hlo); + } Status HandleAllReduce(HloInstructionPtr crs) override { return DefaultAction(crs); } @@ -100,6 +103,9 @@ class DfsHloVisitorWithDefaultBase Status HandleCollectivePermute(HloInstructionPtr hlo) override { return DefaultAction(hlo); } + Status HandleReplicaId(HloInstructionPtr hlo) override { + return DefaultAction(hlo); + } Status HandleRng(HloInstructionPtr random) override { return DefaultAction(random); } diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default_test.cc b/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default_test.cc index 825e1436f0ec6d49b555e5e3e9c2c7a19fb7b062..70173d43d79e931b75f131ad380ad98359cc78b8 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default_test.cc +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default_test.cc @@ -73,15 +73,14 @@ ENTRY TestComputation { abs = f32[] abs(arg) add = f32[] add(arg, gte) broadcast = f32[42] broadcast(add), dimensions={} - slice = f32[0] slice(broadcast), slice={[1:2]} + slice = f32[1] slice(broadcast), slice={[1:2]} copy = f32[] copy(arg) eq = pred[] equal-to(arg, gte) neg = f32[] negate(arg) ROOT convert = f64[] convert(f32[] arg) })"; std::unique_ptr module = - HloRunner::CreateModuleFromString(hlo_string, GetDebugOptionsForTest()) - .ConsumeValueOrDie(); + ParseAndReturnVerifiedModule(hlo_string).ConsumeValueOrDie(); ElementwiseTestVisitor visitor; TF_EXPECT_OK(module->entry_computation()->Accept(&visitor)); } diff --git a/tensorflow/compiler/xla/service/dot_decomposer.cc b/tensorflow/compiler/xla/service/dot_decomposer.cc index b2ba2617902104bfea06713332fa1c2aedea536d..559b9c1f2c9f341293ca89adc61e3312fd9f313c 100644 --- a/tensorflow/compiler/xla/service/dot_decomposer.cc +++ b/tensorflow/compiler/xla/service/dot_decomposer.cc @@ -15,6 +15,8 @@ limitations under the License. #include "tensorflow/compiler/xla/service/dot_decomposer.h" +#include "absl/algorithm/container.h" +#include "absl/strings/str_join.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" @@ -156,29 +158,192 @@ Status DecomposeBatchDot(HloInstruction* dot) { return computation->ReplaceInstruction(dot, new_dot); } +// Convert a dot into a canonical form where non-contracting and contracting +// dimensions are reshaped together and batch dimensions are the most major +// dimensions. The requires transposing and reshapes the lhs and rhs and +// reshaping the output batch to the original shape. +Status CanonicalizeDot(HloInstruction* original_dot) { + auto computation = original_dot->parent(); + const auto& original_dnums = original_dot->dot_dimension_numbers(); + const int64 num_batch_dims = original_dnums.lhs_batch_dimensions_size(); + const int64 num_contracting_dims = + original_dnums.lhs_contracting_dimensions_size(); + + const auto& lhs_shape = original_dot->operand(0)->shape(); + const int64 lhs_rank = lhs_shape.rank(); + const int64 num_lhs_non_contracting_dims = + lhs_rank - num_batch_dims - num_contracting_dims; + + std::vector lhs_non_contracting_dims; + lhs_non_contracting_dims.reserve(num_lhs_non_contracting_dims); + int64 lhs_contracting_size = 1; + int64 lhs_non_contracting_size = 1; + std::vector batch_dim_sizes; + batch_dim_sizes.reserve(num_batch_dims); + for (int64 i = 0; i < lhs_rank; ++i) { + if (absl::c_linear_search(original_dnums.lhs_contracting_dimensions(), i)) { + lhs_contracting_size *= lhs_shape.dimensions(i); + } else if (absl::c_linear_search(original_dnums.lhs_batch_dimensions(), + i)) { + batch_dim_sizes.push_back(lhs_shape.dimensions(i)); + } else { + lhs_non_contracting_dims.push_back(i); + lhs_non_contracting_size *= lhs_shape.dimensions(i); + } + } + // The canonical form of the lhs is + // [BatchDims, NonContractingDims, ContractingsDims] + std::vector lhs_transpose; + lhs_transpose.reserve(lhs_rank); + lhs_transpose.insert(lhs_transpose.end(), + original_dnums.lhs_batch_dimensions().begin(), + original_dnums.lhs_batch_dimensions().end()); + lhs_transpose.insert(lhs_transpose.end(), lhs_non_contracting_dims.begin(), + lhs_non_contracting_dims.end()); + lhs_transpose.insert(lhs_transpose.end(), + original_dnums.lhs_contracting_dimensions().begin(), + original_dnums.lhs_contracting_dimensions().end()); + HloInstruction* transposed_lhs = + computation->AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::PermuteDimensions(InversePermutation(lhs_transpose), + lhs_shape), + original_dot->mutable_operand(0), lhs_transpose)); + std::vector lhs_reshape_dims = batch_dim_sizes; + lhs_reshape_dims.push_back(lhs_non_contracting_size); + lhs_reshape_dims.push_back(lhs_contracting_size); + // Reshape the contracting and non-contracting dimensions together. + HloInstruction* reshaped_lhs = + computation->AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(lhs_shape.element_type(), lhs_reshape_dims), + transposed_lhs)); + + const auto& rhs_shape = original_dot->operand(1)->shape(); + const int64 rhs_rank = rhs_shape.rank(); + const int64 num_rhs_non_contracting_dims = + rhs_rank - num_batch_dims - num_contracting_dims; + std::vector rhs_non_contracting_dims; + rhs_non_contracting_dims.reserve(num_rhs_non_contracting_dims); + int64 rhs_non_contracting_size = 1; + int64 rhs_contracting_size = 1; + for (int64 i = 0; i < rhs_rank; ++i) { + if (absl::c_linear_search(original_dnums.rhs_contracting_dimensions(), i)) { + rhs_contracting_size *= rhs_shape.dimensions(i); + } else if (!absl::c_linear_search(original_dnums.rhs_batch_dimensions(), + i)) { + rhs_non_contracting_dims.push_back(i); + rhs_non_contracting_size *= rhs_shape.dimensions(i); + } + } + + // The canonical form of the rhs is + // [BatchDims, ContractingsDims, NonContractingDims] + std::vector rhs_transpose; + rhs_transpose.reserve(rhs_rank); + rhs_transpose.insert(rhs_transpose.end(), + original_dnums.rhs_batch_dimensions().begin(), + original_dnums.rhs_batch_dimensions().end()); + rhs_transpose.insert(rhs_transpose.end(), + original_dnums.rhs_contracting_dimensions().begin(), + original_dnums.rhs_contracting_dimensions().end()); + rhs_transpose.insert(rhs_transpose.end(), rhs_non_contracting_dims.begin(), + rhs_non_contracting_dims.end()); + HloInstruction* transposed_rhs = + computation->AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::PermuteDimensions(InversePermutation(rhs_transpose), + rhs_shape), + original_dot->mutable_operand(1), rhs_transpose)); + + std::vector rhs_reshape_dims = batch_dim_sizes; + rhs_reshape_dims.push_back(rhs_contracting_size); + rhs_reshape_dims.push_back(rhs_non_contracting_size); + // Reshape the contracting and non-contracting dimensions together. + HloInstruction* reshaped_rhs = + computation->AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(rhs_shape.element_type(), rhs_reshape_dims), + transposed_rhs)); + + std::vector dot_dims = batch_dim_sizes; + dot_dims.push_back(lhs_non_contracting_size); + dot_dims.push_back(rhs_non_contracting_size); + + DotDimensionNumbers dot_dnums; + for (int64 i = 0; i < num_batch_dims; ++i) { + dot_dnums.add_lhs_batch_dimensions(i); + dot_dnums.add_rhs_batch_dimensions(i); + } + dot_dnums.add_lhs_contracting_dimensions(num_batch_dims + 1); + dot_dnums.add_rhs_contracting_dimensions(num_batch_dims); + + HloInstruction* dot = computation->AddInstruction(HloInstruction::CreateDot( + ShapeUtil::MakeShape(original_dot->shape().element_type(), dot_dims), + reshaped_lhs, reshaped_rhs, dot_dnums, original_dot->precision_config())); + + return computation->ReplaceInstruction( + original_dot, computation->AddInstruction(HloInstruction::CreateReshape( + original_dot->shape(), dot))); +} + } // namespace StatusOr DotDecomposer::Run(HloModule* module) { XLA_VLOG_LINES(2, "DotDecomposer ENTRY\n" + module->ToString()); - // Gather all batch Dot operations. - std::vector batch_dots; + // Gather all Non-canonical Dot operations. + std::vector non_canonical_dots; for (auto* computation : module->MakeNonfusionComputations()) { for (auto* instruction : computation->instructions()) { if (instruction->opcode() != HloOpcode::kDot) { continue; } const DotDimensionNumbers& dnums = instruction->dot_dimension_numbers(); - if (dnums.lhs_batch_dimensions_size() > 0 && decompose_batch_dot_) { - batch_dots.push_back(instruction); + // A dot it not canonical if there are more than one contracting + // dimension. + if (dnums.lhs_contracting_dimensions_size() != 1) { + non_canonical_dots.push_back(instruction); + continue; + } + if (dnums.lhs_batch_dimensions().empty() && + dnums.lhs_contracting_dimensions().empty()) { + non_canonical_dots.push_back(instruction); + continue; + } + if (dnums.lhs_batch_dimensions().empty()) { + continue; + } + std::vector canonical_batch_dims( + dnums.lhs_batch_dimensions_size()); + absl::c_iota(canonical_batch_dims, 0); + if (!absl::c_equal(dnums.lhs_batch_dimensions(), canonical_batch_dims) || + !absl::c_equal(dnums.rhs_batch_dimensions(), canonical_batch_dims)) { + non_canonical_dots.push_back(instruction); } } } - // Decompose each batch Dot in 'batch_dots'. bool changed = false; - for (auto* dot : batch_dots) { - TF_RETURN_IF_ERROR(DecomposeBatchDot(dot)); + for (auto* dot : non_canonical_dots) { + TF_RETURN_IF_ERROR(CanonicalizeDot(dot)); changed = true; } + + if (decompose_batch_dot_) { + std::vector batch_dots; + for (auto* computation : module->MakeNonfusionComputations()) { + for (auto* instruction : computation->instructions()) { + if (instruction->opcode() != HloOpcode::kDot) { + continue; + } + const DotDimensionNumbers& dnums = instruction->dot_dimension_numbers(); + if (!dnums.lhs_batch_dimensions().empty()) { + batch_dots.push_back(instruction); + } + } + } + // Decompose each batch Dot in 'batch_dots'. + + for (auto* dot : batch_dots) { + TF_RETURN_IF_ERROR(DecomposeBatchDot(dot)); + changed = true; + } + } XLA_VLOG_LINES(2, "DotDecompose EXIT\n" + module->ToString()); return changed; } diff --git a/tensorflow/compiler/xla/service/dynamic_dimension_inference.cc b/tensorflow/compiler/xla/service/dynamic_dimension_inference.cc index f2b7e9a186135d6246d2dc6c5fb3fa0b70145d10..de3b508064bfadd88396f050142e682de2294434 100644 --- a/tensorflow/compiler/xla/service/dynamic_dimension_inference.cc +++ b/tensorflow/compiler/xla/service/dynamic_dimension_inference.cc @@ -18,6 +18,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/while_util.h" namespace xla { @@ -53,6 +54,8 @@ class DynamicDimensionInferenceVisitor : public DfsHloVisitorWithDefault { Status HandleDot(HloInstruction* hlo) override; + Status HandleTuple(HloInstruction* hlo) override; + Status HandleTranspose(HloInstruction* hlo) override; Status HandleReshape(HloInstruction* hlo) override; @@ -77,6 +80,8 @@ class DynamicDimensionInferenceVisitor : public DfsHloVisitorWithDefault { Status HandleElementwiseBinary(HloInstruction* hlo) override; + Status HandleWhile(HloInstruction* hlo) override; + private: using OperandDynamicDimensionFn = std::functionSetDynamicSize(hlo, index, dimension, dynamic_size); + return Status::OK(); + }); +} + Status DynamicDimensionInferenceVisitor::HandleBroadcast(HloInstruction* hlo) { return ForEachOperandDynamicDimension( hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, @@ -383,6 +398,120 @@ Status DynamicDimensionInferenceVisitor::HandleSelectAndScatter( }); } +Status DynamicDimensionInferenceVisitor::HandleWhile(HloInstruction* hlo) { + // While loop is handled by passing dynamic size hlos as parameters into the + // hlo while loop. This is done by replacing the original while with a new + // one. + // + // Before: + // + // op1 = ... + // op2 = ... + // op1_x = ... // dynamic dimension size of op1 + // while = while(op1, op2) + // + // + // After: + // + // op1 = ... + // op2 = ... + // op1_x = ... // dynamic dimension size of op1 + // while = while(op1, op2, op1_x) + // + // In the above graph, op_x is the bound of the dynamic dimension size of op1 + // and is wired into the while loop as new parameter. + // + // TODO(b/119843103): Once we implement dynamic bounds in XLA backend, dynamic + // bound can be propagated through native xla values instead of relying on + // additional parameter. + + // dynamic_size_to_operand_id_index_map keeps track of dynamic size operations + // to their operand ids in the new while loop. + absl::flat_hash_map + dynamic_size_to_operand_id_index_map; + + // operands_to_add collects dynamic sizes that need to be added to the while + // loop as parameters. Note that a dynamic size is ignored if it is already + // part of the parameter. i.e.: + // + // We don't do: + // + // op1 = ... + // op2 = ... + // op_x = ... // dynamic dimension size of both op1 and op2 + // while = while(op1, op2, op_x, op_x) // 4 parameters + // + // But we do: + // + // op1 = ... + // op2 = ... + // op_x = ... // dynamic dimension size of both op1 and op2 + // while = while(op1, op2, op_x) + // + // An alternative is to do this in a while loop CSE pass. + // + std::vector operands_to_add; + int64 operand_count = hlo->shape().tuple_shapes_size(); + TF_RETURN_IF_ERROR(ForEachOperandDynamicDimension( + hlo, [&](HloInstruction*, ShapeIndex, int64, int64, + HloInstruction* dynamic_size) { + const HloInstruction* tuple_operand = hlo->operand(0); + for (int64 i = 0; i < tuple_operand->operand_count(); ++i) { + if (dynamic_size == tuple_operand->operand(i)) { + dynamic_size_to_operand_id_index_map[dynamic_size] = i; + return Status::OK(); + } + } + auto iter = dynamic_size_to_operand_id_index_map.find(dynamic_size); + if (iter == dynamic_size_to_operand_id_index_map.end()) { + operands_to_add.push_back(dynamic_size); + dynamic_size_to_operand_id_index_map[dynamic_size] = operand_count++; + } + return Status::OK(); + })); + + if (!operands_to_add.empty()) { + // Only replace the while loop if there are new parameters to add. + HloInstruction* old_tuple_operand = hlo->mutable_operand(0); + TF_ASSIGN_OR_RETURN( + WhileUtil::MakeInstructionsLiveInResult result, + WhileUtil::MakeInstructionsLiveIn(hlo, operands_to_add)); + // WhileUtil creates a new while hlo and tuple. Update the dynamic size + // mapping for the newly created tuple. + HloInstruction* new_tuple_operand = + result.new_while_instr->mutable_operand(0); + parent_->CopyMapping(/*from=*/old_tuple_operand, /*to=*/new_tuple_operand); + hlo = result.new_while_instr; + } + + // We have replaced the while loop, now set the dynamic dimensions for the + // newly created while loop so that the hlos that consumes the while loop can + // see the dynamic dimensions. Also sets the dynamic parameter binding for + // running inference in the while loop. + DynamicParameterBinding binding_for_while; + TF_RETURN_IF_ERROR(ForEachOperandDynamicDimension( + hlo, [&](HloInstruction*, ShapeIndex index, int64 dimension, + int64 operand_index, HloInstruction* dynamic_size) { + DynamicParameterBinding::DynamicParameter dynamic_parameter{ + operand_index, + {dynamic_size_to_operand_id_index_map[dynamic_size]}}; + DynamicParameterBinding::DynamicDimension dynamic_dimension{ + operand_index, index, dimension}; + TF_RETURN_IF_ERROR( + binding_for_while.Bind(dynamic_parameter, dynamic_dimension)); + parent_->SetDynamicSize(hlo, index, dimension, dynamic_size); + return Status::OK(); + })); + + // Run inference in while body and condition. + TF_RETURN_IF_ERROR(DynamicDimensionInferenceVisitor::Run( + hlo->while_body(), binding_for_while, parent_)); + TF_RETURN_IF_ERROR(DynamicDimensionInferenceVisitor::Run( + hlo->while_condition(), binding_for_while, parent_)); + + return Status::OK(); +} + Status DynamicDimensionInferenceVisitor::HandleParameter(HloInstruction* hlo) { return param_bindings_.ForEachBinding( [&](const DynamicParameterBinding::DynamicParameter& dynamic_parameter, @@ -430,15 +559,43 @@ Status DynamicDimensionInferenceVisitor::ForEachOperandDynamicDimension( return Status::OK(); } +void DynamicDimensionInference::CopyMapping(HloInstruction* from, + HloInstruction* to) { + auto iter = per_hlo_dynamic_dimensions_.find(from); + if (iter != per_hlo_dynamic_dimensions_.end()) { + for (auto& dynamic_dimension : iter->second) { + HloInstruction* dynamic_size = + GetDynamicSize(dynamic_dimension.inst, dynamic_dimension.index, + dynamic_dimension.dim); + SetDynamicSize(to, dynamic_dimension.index, dynamic_dimension.dim, + dynamic_size); + } + } +} + /* static */ StatusOr DynamicDimensionInference::Run( HloModule* module) { - VLOG(0) << "Param Config " << module->dynamic_parameter_binding().ToString(); + VLOG(2) << "Param Config " << module->dynamic_parameter_binding().ToString(); DynamicDimensionInference inference(module); TF_RETURN_IF_ERROR(inference.AnalyzeDynamicDimensions()); return inference; } +string DynamicDimensionInference::ToString() const { + std::vector pieces; + pieces.push_back("DynamicDimensionInference: "); + for (const auto& mapping : dynamic_mapping_) { + const DynamicDimension& dynamic_dimension = mapping.first; + pieces.push_back(absl::StrFormat( + " -- instruction %s at %s has dim %lld as dynamic" + " dimension, which is represented by instruction %s", + dynamic_dimension.inst->ToString(), dynamic_dimension.index.ToString(), + dynamic_dimension.dim, mapping.second->ToString())); + } + return absl::StrJoin(pieces, "\n"); +} + DynamicDimensionInference::DynamicDimensionInference(HloModule* module) : module_(module) {} diff --git a/tensorflow/compiler/xla/service/dynamic_dimension_inference.h b/tensorflow/compiler/xla/service/dynamic_dimension_inference.h index 164d15bf111a92e3da957f609b54ee0662ef18b1..d0f2998328f3028ccbd5b33690a514371a03b5a1 100644 --- a/tensorflow/compiler/xla/service/dynamic_dimension_inference.h +++ b/tensorflow/compiler/xla/service/dynamic_dimension_inference.h @@ -88,6 +88,11 @@ class DynamicDimensionInference { iter.first->second.emplace(DynamicDimension{inst, index, dim}); } + // Copies the internal mapping from instruction `from` to instruction `to`. + // This is useful when an instruction is replaced by the other during the + // inferencing process. + void CopyMapping(HloInstruction* from, HloInstruction* to); + // AnalyzeDynamicDimensions starts the analysis of the dynamic dimensions in // module_. Status AnalyzeDynamicDimensions(); @@ -101,6 +106,8 @@ class DynamicDimensionInference { using DynamicMapping = absl::flat_hash_map; DynamicMapping dynamic_mapping_; + // A convenient mapping from an hlo to the set of dynamic dimensions that it + // holds. using PerHloDynamicDimensions = absl::flat_hash_map>; diff --git a/tensorflow/compiler/xla/service/dynamic_dimension_inference_test.cc b/tensorflow/compiler/xla/service/dynamic_dimension_inference_test.cc index 1dd196821c05cc820e2a3bf53a04d96b15484cd4..597cdf27c3318b3cf8bd5bb5f9b3239cf23a4c73 100644 --- a/tensorflow/compiler/xla/service/dynamic_dimension_inference_test.cc +++ b/tensorflow/compiler/xla/service/dynamic_dimension_inference_test.cc @@ -62,6 +62,17 @@ class DynamicDimensionInferenceTest : public HloTestBase { return module_->AddEmbeddedComputation(embedded_builder.Build()); } + HloComputation* GetGe() { + auto embedded_builder = HloComputation::Builder("ge"); + auto lhs = embedded_builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {}), "lhs")); + auto rhs = embedded_builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {}), "rhs")); + embedded_builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(PRED, {}), HloOpcode::kGe, lhs, rhs)); + return module_->AddEmbeddedComputation(embedded_builder.Build()); + } + std::unique_ptr module_; std::unique_ptr inference_; const Shape scalar_shape_ = ShapeUtil::MakeShape(S32, {}); @@ -434,6 +445,96 @@ TEST_F(DynamicDimensionInferenceTest, BroadcastTest) { EXPECT_EQ(inference_->GetDynamicSize(broadcast, {}, 2), nullptr); } +TEST_F(DynamicDimensionInferenceTest, WhileTest) { + // Test the ability to trace into while loops. + auto builder = HloComputation::Builder(TestName()); + auto input_shape = ShapeUtil::MakeShape(F32, {2, 4, 4}); + auto output_shape = ShapeUtil::MakeShape(F32, {2, 2, 2}); + auto tuple_shape = ShapeUtil::MakeTupleShape({input_shape, input_shape}); + + // Body: + // + // Param + // | | + // GTE1 GTE2 + // | | + // ADD + auto body_builder = HloComputation::Builder("body"); + auto body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto gte_0 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(input_shape, body_param, 0)); + auto gte_1 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(input_shape, body_param, 1)); + auto add = body_builder.AddInstruction( + HloInstruction::CreateBinary(input_shape, HloOpcode::kAdd, gte_0, gte_1)); + body_builder.AddInstruction(HloInstruction::CreateTuple({add, add})); + + HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); + + auto cond_builder = HloComputation::Builder("condition"); + cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + cond_builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(false))); + HloComputation* condition = + module_->AddEmbeddedComputation(cond_builder.Build()); + + // Entry: + // + // Param + // | + // While + auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/0, tuple_shape, "A")); + auto* size_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/1, scalar_shape_, "size_param")); + builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition, body, a_param)); + + module_->AddEntryComputation(builder.Build()); + + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{1, {}}, + DynamicParameterBinding::DynamicDimension{0, {0}, 0})); + + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{1, {}}, + DynamicParameterBinding::DynamicDimension{0, {1}, 0})); + + // Test that dynamic dimension inference does the right thing. A lambda is + // used here since we want to test twice by running inference again + // (idempotency). + auto test_dynamic_dimension = [&]() { + HloInstruction* while_hlo = nullptr; + // The while hlo has been replaced, find the new one. + for (HloInstruction* inst : module_->entry_computation()->instructions()) { + if (inst->opcode() == HloOpcode::kWhile) { + while_hlo = inst; + } + } + ASSERT_NE(while_hlo, nullptr); + // The original while shape has 2 parameters. With dynamic size passed in + // as an extra parameter, the tuple should have 3 elements. + EXPECT_EQ(while_hlo->shape().tuple_shapes_size(), 3); + HloInstruction* add = nullptr; + for (HloInstruction* inst : while_hlo->while_body()->instructions()) { + if (inst->opcode() == HloOpcode::kAdd) { + add = inst; + } + } + EXPECT_NE(add, nullptr); + EXPECT_NE(inference_->GetDynamicSize(add, {}, 0), nullptr); + EXPECT_EQ(inference_->GetDynamicSize(while_hlo, {0}, 0), size_param); + EXPECT_EQ(inference_->GetDynamicSize(while_hlo, {1}, 0), size_param); + }; + + TF_ASSERT_OK(RunInference()); + test_dynamic_dimension(); + TF_ASSERT_OK(RunInference()); + test_dynamic_dimension(); +} + TEST_F(DynamicDimensionInferenceTest, ReduceWindowBatchTest) { // Test the ability to trace reduce window batch dimensions. auto builder = HloComputation::Builder(TestName()); @@ -487,7 +588,7 @@ TEST_F(DynamicDimensionInferenceTest, SelectAndScatterTest) { // Test the ability to trace select and scatter batch dimensions. auto builder = HloComputation::Builder(TestName()); auto input_shape = ShapeUtil::MakeShape(F32, {2, 4, 4}); - auto output_shape = ShapeUtil::MakeShape(F32, {2, 2, 2}); + auto source_shape = ShapeUtil::MakeShape(F32, {2, 2, 2}); Window window; // First dimension is unchanged. @@ -514,22 +615,26 @@ TEST_F(DynamicDimensionInferenceTest, SelectAndScatterTest) { /*parameter_number=*/0, input_shape, "A")); auto* size_param = builder.AddInstruction(HloInstruction::CreateParameter( /*parameter_number=*/1, scalar_shape_, "size_param")); + auto* source = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/2, source_shape, "B")); auto init = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0))); - auto* reduce_window = - builder.AddInstruction(HloInstruction::CreateReduceWindow( - output_shape, a_param, init, window, GetAdd())); + auto* sns = builder.AddInstruction(HloInstruction::CreateSelectAndScatter( + input_shape, a_param, GetGe(), window, source, init, GetAdd())); module_->AddEntryComputation(builder.Build()); TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( DynamicParameterBinding::DynamicParameter{1, {}}, DynamicParameterBinding::DynamicDimension{0, {}, 0})); + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{1, {}}, + DynamicParameterBinding::DynamicDimension{2, {}, 0})); TF_ASSERT_OK(RunInference()); - EXPECT_EQ(inference_->GetDynamicSize(reduce_window, {}, 0), size_param); + EXPECT_EQ(inference_->GetDynamicSize(sns, {}, 0), size_param); } } // namespace diff --git a/tensorflow/compiler/xla/service/dynamic_index_splitter.cc b/tensorflow/compiler/xla/service/dynamic_index_splitter.cc index ede4a5afdb7d02b31b9d0839fc8716f9b814e544..e34adfd2d2bbb7214cfa2da28291b133538845e5 100644 --- a/tensorflow/compiler/xla/service/dynamic_index_splitter.cc +++ b/tensorflow/compiler/xla/service/dynamic_index_splitter.cc @@ -45,17 +45,10 @@ StatusOr DynamicIndexSplitter::Run(HloModule* module) { } auto parent = dynamic_op->parent(); bool is_update = dynamic_op->opcode() == HloOpcode::kDynamicUpdateSlice; - int64 index_operand_number = Cast(dynamic_op) - ->first_index_operand_number(); - auto index_operand = dynamic_op->mutable_operand(index_operand_number); - if (ShapeUtil::IsScalar(index_operand->shape())) { - // This DS/DUS already uses scalar indices. - continue; - } - TF_RET_CHECK(index_operand->shape().rank() == 1); - int64 num_indices = index_operand->shape().dimensions(0); + int64 num_indices = dynamic_op->operand(0)->shape().rank(); + if (num_indices == 0) { - // If the operand dimension is 0, directly replace R0 DS/DUS with the + // If the operand rank is 0, directly replace R0 DS/DUS with the // operand (for DS) or update (for DUS). if (is_update) { TF_CHECK_OK(parent->ReplaceInstruction( @@ -67,6 +60,15 @@ StatusOr DynamicIndexSplitter::Run(HloModule* module) { changed = true; continue; } + + int64 index_operand_number = Cast(dynamic_op) + ->first_index_operand_number(); + auto index_operand = dynamic_op->mutable_operand(index_operand_number); + if (ShapeUtil::IsScalar(index_operand->shape())) { + // This DS/DUS already uses scalar indices. + continue; + } + TF_RET_CHECK(index_operand->shape().rank() == 1); auto index_element_type = index_operand->shape().element_type(); std::vector index_array; for (int64 dim = 0; dim < num_indices; ++dim) { diff --git a/tensorflow/compiler/xla/service/dynamic_padder.cc b/tensorflow/compiler/xla/service/dynamic_padder.cc new file mode 100644 index 0000000000000000000000000000000000000000..4db280f817141bd52e3a5b9564600a618f81aeac --- /dev/null +++ b/tensorflow/compiler/xla/service/dynamic_padder.cc @@ -0,0 +1,161 @@ +/* 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/dynamic_padder.h" + +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" + +#include "absl/container/flat_hash_set.h" +#include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/dynamic_dimension_inference.h" +#include "tensorflow/compiler/xla/service/hlo_dce.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/util.h" + +#include "tensorflow/core/lib/core/errors.h" + +namespace xla { + +namespace { + +// ChooseIdentityValue looks at the instruction and returns a identity value +// which, when padded, doesn't change the result of the instruction. +// +// nullopt is returned if padding doesn't need to be reset. +StatusOr ChooseIdentityValue(HloInstruction* inst) { + HloComputation* comp = inst->parent(); + // Padding on elementwise operation doesn't affect the result of the effective + // data. + if (inst->IsElementwise()) { + return nullptr; + } + + switch (inst->opcode()) { + case HloOpcode::kReduce: + case HloOpcode::kReduceWindow: { + // Because of the way we do reduce, we already require the `init` operand + // of hlo reduce instruction to be identity value. Here we reuse the + // operand. + return inst->mutable_operand(1); + } + + case HloOpcode::kConvolution: + case HloOpcode::kDot: { + // Use 0 as padding value for convolution and dot. + PrimitiveType ptype = inst->shape().element_type(); + return comp->AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::Zero(ptype))); + } + + case HloOpcode::kPad: { + return inst->mutable_operand(1); + } + case HloOpcode::kParameter: + case HloOpcode::kGetDimensionSize: + case HloOpcode::kReshape: + case HloOpcode::kTuple: + case HloOpcode::kAllReduce: + case HloOpcode::kBroadcast: + return nullptr; + default: + return UnimplementedStrCat("Unimplimented padding for instruction: ", + inst->ToString()); + } +} + +} // namespace + +StatusOr DynamicPadder::Run(HloModule* module) { + bool changed = false; + VLOG(2) << "Pre DynamicPadder HLO:"; + XLA_VLOG_LINES(2, module->ToString()); + TF_ASSIGN_OR_RETURN(DynamicDimensionInference dynamic_dimension_inference, + DynamicDimensionInference::Run(module)); + + for (HloComputation* computation : module->computations()) { + for (HloInstruction* inst : computation->instructions()) { + for (int64 operand_num = 0; operand_num < inst->operand_count(); + ++operand_num) { + HloInstruction* operand = inst->mutable_operand(operand_num); + if (!operand->shape().IsArray()) { + continue; + } + for (int64 dim = 0; dim < operand->shape().rank(); ++dim) { + HloInstruction* dynamic_size = + dynamic_dimension_inference.GetDynamicSize(operand, {}, dim); + if (dynamic_size == nullptr) { + continue; + } + VLOG(1) << "Has dynamic dimension of operand" << operand_num << " @" + << dim; + TF_ASSIGN_OR_RETURN(HloInstruction * identity_value, + ChooseIdentityValue(inst)); + if (identity_value == nullptr) { + continue; + } + + // For each dimension, first generates a mask representing the + // effective area of data and padded area of data using iota and + // dynamic_size. For example, given a dimension of 7 elements and 5 + // effective elements: + // + // iota = [0, 1, 2, 3, 4, 5, 6] + // broadcast_dynamic_size = [5, 5, 5, 5, 5, 5, 5] + // mask = lt(iota, broadcast_dynamic_size) = [t, t, t, t, t, f, f] + // + // Once the mask is generated, the input data is then padded using the + // mask and pad value. + // + const Shape mask_shape = + ShapeUtil::ChangeElementType(operand->shape(), xla::U32); + const Shape pred_shape = + ShapeUtil::ChangeElementType(operand->shape(), xla::PRED); + HloInstruction* iota = computation->AddInstruction( + HloInstruction::CreateIota(mask_shape, dim)); + + HloInstruction* broadcasted_effective_size = + computation->AddInstruction(HloInstruction::CreateBroadcast( + mask_shape, dynamic_size, {})); + HloInstruction* pred = computation->AddInstruction( + HloInstruction::CreateBinary(pred_shape, HloOpcode::kLt, iota, + broadcasted_effective_size)); + + HloInstruction* broadcasted_identity_value = + computation->AddInstruction(HloInstruction::CreateBroadcast( + operand->shape(), identity_value, {})); + HloInstruction* padded = + computation->AddInstruction(HloInstruction::CreateTernary( + operand->shape(), HloOpcode::kSelect, pred, operand, + broadcasted_identity_value)); + TF_RETURN_IF_ERROR(inst->ReplaceOperandWith(operand_num, padded)); + operand = inst->mutable_operand(operand_num); + changed = true; + } + } + } + } + HloDCE dce; + TF_ASSIGN_OR_RETURN(changed, dce.Run(module)); + VLOG(2) << "Post DynamicPadder HLO:"; + XLA_VLOG_LINES(2, module->ToString()); + return changed; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/dynamic_padder.h b/tensorflow/compiler/xla/service/dynamic_padder.h new file mode 100644 index 0000000000000000000000000000000000000000..509269f7f56746fa5516ad917a04221587c6dcca --- /dev/null +++ b/tensorflow/compiler/xla/service/dynamic_padder.h @@ -0,0 +1,44 @@ +/* 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_DYNAMIC_PADDER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_DYNAMIC_PADDER_H_ + +#include "tensorflow/compiler/xla/service/dynamic_dimension_inference.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" + +namespace xla { + +// With bounded shapes, only part of the shape contains effective data and the +// rest contains padded data, whose value can be anything depending on the +// source of the data. When a bounded shape is directly consumed by an +// instruction that collapses dimensions (reduce for example), the padding data +// would affect result of the instruction. +// +// DynamicPadder uses DynamicDimensionInference to detect bounded shapes in a +// hlo module, it then inserts certain instructions to reset the padding into an +// identity value so that in doesn't affect the result of subsequent +// instruction. For example, it'd reset the padding to 0 before a bounded shape +// is consumed by a reduce-sum. +class DynamicPadder : public HloModulePass { + public: + absl::string_view name() const override { return "dynamic_padder"; } + + StatusOr Run(HloModule* module) override; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_DYNAMIC_PADDER_H_ diff --git a/tensorflow/compiler/xla/service/dynamic_padder_test.cc b/tensorflow/compiler/xla/service/dynamic_padder_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..55a11286e4596d87c330315322cae704fc5cd707 --- /dev/null +++ b/tensorflow/compiler/xla/service/dynamic_padder_test.cc @@ -0,0 +1,152 @@ +/* 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/dynamic_padder.h" + +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#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_runner.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/xla_data.pb.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/platform/test_benchmark.h" + +namespace op = xla::testing::opcode_matchers; + +namespace xla { +namespace { + +class DynamicPadderTest : public HloTestBase { + protected: + DynamicPadderTest() : HloTestBase() { module_ = CreateNewVerifiedModule(); } + + StatusOr RunPadder() { + hlo_graph_dumper::MaybeDumpHloModule(*module_, "Before padder"); + + DynamicPadder padder; + + return padder.Run(module_.get()); + } + + void ExpectPadded(const HloInstruction* inst) { + EXPECT_THAT(inst, + op::Select(op::Lt(op::Iota(), op::Broadcast(op::Parameter())), + ::testing::_, op::Broadcast())); + } + + HloComputation* GetScalarAddComputation() { + auto embedded_builder = HloComputation::Builder("add"); + auto lhs = embedded_builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {}), "lhs")); + auto rhs = embedded_builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {}), "rhs")); + embedded_builder.AddInstruction( + HloInstruction::CreateBinary(lhs->shape(), HloOpcode::kAdd, lhs, rhs)); + return module_->AddEmbeddedComputation(embedded_builder.Build()); + } + + std::unique_ptr module_; + const Shape scalar_shape_ = ShapeUtil::MakeShape(U32, {}); +}; + +TEST_F(DynamicPadderTest, ReduceTest) { + auto builder = HloComputation::Builder(TestName()); + auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 2}); + auto reduce_shape = ShapeUtil::MakeShape(F32, {2}); + + auto data_param = builder.AddInstruction( + HloInstruction::CreateParameter(0, input_shape, "data_param")); + builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "size_param")); + + auto negate = builder.AddInstruction( + HloInstruction::CreateUnary(input_shape, HloOpcode::kNegate, data_param)); + + auto init = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0))); + + auto reduce = builder.AddInstruction(HloInstruction::CreateReduce( + reduce_shape, negate, init, {0, 2}, GetScalarAddComputation())); + + module_->AddEntryComputation(builder.Build()); + + // Set up dynamic parameter binding. + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{1, {}}, + DynamicParameterBinding::DynamicDimension{0, {}, 1})); + + TF_ASSERT_OK(RunPadder().status()); + + ExpectPadded(reduce->operand(0)); +} + +TEST_F(DynamicPadderTest, ConvolutionTest) { + 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_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{2, {}}, + DynamicParameterBinding::DynamicDimension{0, {}, 1})); + + TF_ASSERT_OK(RunPadder().status()); + + ExpectPadded(conv->operand(0)); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/dynamic_parameter_binding.cc b/tensorflow/compiler/xla/service/dynamic_parameter_binding.cc index 3d0eddb0d857ff1d9f3008f13df273fe0db6c7c2..7f0ae692f7414dbdcccda8b287c9059bcf920df1 100644 --- a/tensorflow/compiler/xla/service/dynamic_parameter_binding.cc +++ b/tensorflow/compiler/xla/service/dynamic_parameter_binding.cc @@ -29,7 +29,8 @@ Status DynamicParameterBinding::Bind( } absl::optional -DynamicParameterBinding::GetBinding(const DynamicDimension& dynamic_dimension) { +DynamicParameterBinding::GetBinding( + const DynamicDimension& dynamic_dimension) const { auto param_iter = bindings_.find(dynamic_dimension); if (param_iter == bindings_.end()) { return absl::nullopt; @@ -70,7 +71,7 @@ StatusOr DynamicParameterBinding::CreateFromProto( int64 target_param_num = binding.target_param_num(); ShapeIndex target_param_index(binding.target_param_index().begin(), binding.target_param_index().end()); - int64 target_dim_num = binding.target_param_num(); + int64 target_dim_num = binding.target_param_dim_num(); TF_RETURN_IF_ERROR( result.Bind(DynamicParameter{dynamic_param_num, dynamic_param_index}, @@ -111,7 +112,8 @@ Status DynamicParameterBinding::Verify(const HloModule& module) const { return ForEachBinding([&](const DynamicParameter& dynamic_parameter, const DynamicDimension& dynamic_dimension) -> Status { - TF_RET_CHECK(dynamic_parameter.parameter_num < entry->num_parameters()); + TF_RET_CHECK(dynamic_parameter.parameter_num >= 0 && + dynamic_parameter.parameter_num < entry->num_parameters()); TF_RET_CHECK(dynamic_dimension.parameter_num < entry->num_parameters()); TF_RET_CHECK(ShapeUtil::IndexIsValid( entry->parameter_instruction(dynamic_parameter.parameter_num)->shape(), diff --git a/tensorflow/compiler/xla/service/dynamic_parameter_binding.h b/tensorflow/compiler/xla/service/dynamic_parameter_binding.h index dd474d8eed1b2c30ddb8f624a864198c74eacaba..57af2c43d3c65f7340e6a9f04e5abbf052ebceea 100644 --- a/tensorflow/compiler/xla/service/dynamic_parameter_binding.h +++ b/tensorflow/compiler/xla/service/dynamic_parameter_binding.h @@ -89,7 +89,7 @@ class DynamicParameterBinding { // // Returns nullopt if the binding is not set. absl::optional GetBinding( - const DynamicDimension& dynamic_dimension); + const DynamicDimension& dynamic_dimension) const; using BindingFn = std::functionToProto(); + TF_ASSERT_OK_AND_ASSIGN(*binding, + DynamicParameterBinding::CreateFromProto(proto)); + } +}; TEST_F(DynamicParameterBindingTest, SimpleBinding) { // 'b' is a dynamic shape; 'a' represents the real size of b's first @@ -56,15 +64,20 @@ ENTRY main { binding.Bind(DynamicParameterBinding::DynamicParameter{0, {}}, DynamicParameterBinding::DynamicDimension{1, {}, 0})); - absl::optional param = - binding.GetBinding( - DynamicParameterBinding::DynamicDimension{/*parameter_num=*/1, - /*parameter_index=*/{}, - /*dimension=*/0}); - EXPECT_TRUE(param); - EXPECT_EQ(param->parameter_num, 0); - EXPECT_EQ(param->parameter_index, ShapeIndex({})); - TF_EXPECT_OK(binding.Verify(*module)); + auto test = [&](const DynamicParameterBinding& binding) { + absl::optional param = + binding.GetBinding( + DynamicParameterBinding::DynamicDimension{/*parameter_num=*/1, + /*parameter_index=*/{}, + /*dimension=*/0}); + EXPECT_TRUE(param); + EXPECT_EQ(param->parameter_num, 0); + EXPECT_EQ(param->parameter_index, ShapeIndex({})); + TF_EXPECT_OK(binding.Verify(*module)); + }; + test(binding); + SerializeAndDeserialize(&binding); + test(binding); } TEST_F(DynamicParameterBindingTest, TupleBinding) { @@ -89,16 +102,21 @@ ENTRY main { binding.Bind(DynamicParameterBinding::DynamicParameter{0, {0}}, DynamicParameterBinding::DynamicDimension{0, {1}, 0})); - absl::optional param = - binding.GetBinding( - DynamicParameterBinding::DynamicDimension{/*parameter_num=*/0, - /*parameter_index=*/{1}, - /*dimension=*/0}); - - EXPECT_TRUE(param); - EXPECT_EQ(param->parameter_num, 0); - EXPECT_EQ(param->parameter_index, ShapeIndex({0})); - TF_EXPECT_OK(binding.Verify(*module)); + auto test = [&](const DynamicParameterBinding& binding) { + absl::optional param = + binding.GetBinding( + DynamicParameterBinding::DynamicDimension{/*parameter_num=*/0, + /*parameter_index=*/{1}, + /*dimension=*/0}); + + EXPECT_TRUE(param); + EXPECT_EQ(param->parameter_num, 0); + EXPECT_EQ(param->parameter_index, ShapeIndex({0})); + TF_EXPECT_OK(binding.Verify(*module)); + }; + test(binding); + SerializeAndDeserialize(&binding); + test(binding); } TEST_F(DynamicParameterBindingTest, TupleBindingWithMultiDimension) { @@ -127,26 +145,35 @@ ENTRY main { binding.Bind(DynamicParameterBinding::DynamicParameter{0, {0}}, DynamicParameterBinding::DynamicDimension{0, {1}, 1})); - absl::optional param = - binding.GetBinding( - DynamicParameterBinding::DynamicDimension{/*parameter_num=*/0, - /*parameter_index=*/{1}, - /*dimension=*/0}); - - EXPECT_TRUE(param); - EXPECT_EQ(param->parameter_num, 0); - EXPECT_EQ(param->parameter_index, ShapeIndex({0})); - - absl::optional param2 = - binding.GetBinding( - DynamicParameterBinding::DynamicDimension{/*parameter_num=*/0, - /*parameter_index=*/{1}, - /*dimension=*/0}); - EXPECT_TRUE(param2); - EXPECT_EQ(param2->parameter_num, 0); - EXPECT_EQ(param2->parameter_index, ShapeIndex({0})); - - TF_EXPECT_OK(binding.Verify(*module)); + auto test = [&](const DynamicParameterBinding& binding) { + absl::optional param = + binding.GetBinding( + DynamicParameterBinding::DynamicDimension{/*parameter_num=*/0, + /*parameter_index=*/{1}, + /*dimension=*/0}); + + EXPECT_TRUE(param); + EXPECT_EQ(param->parameter_num, 0); + EXPECT_EQ(param->parameter_index, ShapeIndex({0})); + + absl::optional param2 = + + binding.GetBinding( + DynamicParameterBinding::DynamicDimension{/*parameter_num=*/0, + /*parameter_index=*/{1}, + /*dimension=*/0}); + EXPECT_TRUE(param2); + EXPECT_EQ(param2->parameter_num, 0); + EXPECT_EQ(param2->parameter_index, ShapeIndex({0})); + TF_EXPECT_OK(binding.Verify(*module)); + }; + + test(binding); + + SerializeAndDeserialize(&binding); + + // Test the binding again after deserialization. + test(binding); } } // namespace diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc index f84c115e0abec378f0401a15e6bd381983d0ff34..e868dc6d889c867001bf2a145bb9277c56950401 100644 --- a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc @@ -440,14 +440,16 @@ StatusOr ElementalIrEmitter::EmitFloatUnaryOp( {operand_value}, {operand_value->getType()}, b_); case HloOpcode::kSign: { - // TODO(b/32151903): Ensure consistent sign behavior for -0.0. auto type = operand_value->getType(); auto zero = llvm::ConstantFP::get(type, 0.0); - auto oeq = FCmpOEQ(operand_value, zero); - auto olt = FCmpOLT(operand_value, zero); - return Select(oeq, zero, - Select(olt, llvm::ConstantFP::get(type, -1.0), - llvm::ConstantFP::get(type, 1.0))); + auto ne0_i1 = FCmpONE(operand_value, zero); + auto ne0_float = UIToFP(ne0_i1, type); + llvm::Value* result = llvm_ir::EmitCallToIntrinsic( + llvm::Intrinsic::copysign, {ne0_float, operand_value}, + {operand_value->getType()}, b_); + auto is_nan = FCmpUNO(operand_value, operand_value); + result = Select(is_nan, operand_value, result); + return result; } case HloOpcode::kIsFinite: { // abs(x) o!= inf, this works because the comparison returns false if @@ -812,11 +814,14 @@ StatusOr ElementalIrEmitter::EmitComplexBinaryOp( 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); @@ -828,7 +833,13 @@ StatusOr ElementalIrEmitter::EmitComplexBinaryOp( 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)); - return EmitComposeComplex(op, FMul(coeff, cos_q), FMul(coeff, sin_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))); } default: return Unimplemented("binary complex op '%s'", @@ -846,6 +857,9 @@ llvm::Value* ElementalIrEmitter::EmitFloatMin(llvm::Value* lhs_value, return llvm_ir::EmitFloatMin(lhs_value, rhs_value, b_); } +// TODO(b/123355973): We have an implementation of erfinv in math.cc. We +// shouldn't have two implementations, especially since this one isn't testable +// (it's only observable via a normally-distributed RNG). StatusOr ElementalIrEmitter::EmitErfInv(PrimitiveType prim_type, llvm::Value* x) { if (prim_type != F16 && prim_type != F32 && prim_type != F64) { @@ -1758,18 +1772,10 @@ StatusOr ElementalIrEmitter::EmitElementalDynamicSlice( auto index_typed_const = [&](uint64 c) -> llvm::Constant* { return llvm::ConstantInt::get(index_type, c); }; - // TODO(b/118437727): Remove the R1 path. - llvm::Value* start_index_value; - if (hlo->operand(1)->shape().rank() == 1) { - llvm_ir::IrArray::Index dim_index(1, index_typed_const(i)); - TF_ASSIGN_OR_RETURN(start_index_value, - operand_to_generator.at(hlo->operand(1))(dim_index)); - } else { - llvm_ir::IrArray::Index zero_index(index_type); - TF_ASSIGN_OR_RETURN( - start_index_value, - operand_to_generator.at(hlo->operand(1 + i))(zero_index)); - } + llvm_ir::IrArray::Index zero_index(index_type); + TF_ASSIGN_OR_RETURN( + llvm::Value * start_index_value, + operand_to_generator.at(hlo->operand(1 + i))(zero_index)); // Clamp the start index so that the sliced portion fits in the operand: // start_index = clamp(start_index, 0, operand_dim_size - output_dim_size) @@ -1915,18 +1921,10 @@ StatusOr ElementalIrEmitter::EmitElementalDynamicUpdateSlice( return llvm::ConstantInt::get(index_type, c); }; - llvm::Value* start_index_value; - // TODO(b/118437727): Remove the R1 path. - if (hlo->operand(2)->shape().rank() == 1) { - llvm_ir::IrArray::Index dim_index(1, index_typed_const(i)); - TF_ASSIGN_OR_RETURN(start_index_value, - operand_to_generator.at(hlo->operand(2))(dim_index)); - } else { - llvm_ir::IrArray::Index zero_index(index_type); - TF_ASSIGN_OR_RETURN( - start_index_value, - operand_to_generator.at(hlo->operand(2 + i))(zero_index)); - } + llvm_ir::IrArray::Index zero_index(index_type); + TF_ASSIGN_OR_RETURN( + llvm::Value * start_index_value, + operand_to_generator.at(hlo->operand(2 + i))(zero_index)); // Clamp the start index so that the update region fits in the operand. // start_index = clamp(start_index, 0, input_dim_size - update_dim_size) diff --git a/tensorflow/compiler/xla/service/executable.cc b/tensorflow/compiler/xla/service/executable.cc index 10b8c01ff1383658fcfb2271c177ba54347f985a..1518d83083b3b0ce876da9344c483a23cd5b073c 100644 --- a/tensorflow/compiler/xla/service/executable.cc +++ b/tensorflow/compiler/xla/service/executable.cc @@ -26,7 +26,6 @@ limitations under the License. #include "tensorflow/core/lib/strings/proto_serialization.h" #include "tensorflow/core/platform/env.h" - namespace xla { StatusOr> Executable::ExecuteOnStreams( @@ -173,11 +172,13 @@ Status Executable::DumpHloSnapshot() { } filename = SanitizeFileName(std::move(filename)); string file_path = tensorflow::io::JoinPath(directory_path, filename); - string result; - TF_RET_CHECK( - tensorflow::SerializeToStringDeterministic(hlo_session, &result)); - return tensorflow::WriteStringToFile(tensorflow::Env::Default(), file_path, - result); + const size_t size = hlo_session.ByteSizeLong(); + auto serialized = absl::make_unique(size); + TF_RET_CHECK(tensorflow::SerializeToBufferDeterministic( + hlo_session, serialized.get(), size)); + return tensorflow::WriteStringToFile( + tensorflow::Env::Default(), file_path, + absl::string_view(serialized.get(), size)); } } // namespace xla diff --git a/tensorflow/compiler/xla/service/gather_expander.cc b/tensorflow/compiler/xla/service/gather_expander.cc index 01cef499665c050d4453382289168276028e1d26..a58ac39dffad56315308f784b08e6b6087b8e30a 100644 --- a/tensorflow/compiler/xla/service/gather_expander.cc +++ b/tensorflow/compiler/xla/service/gather_expander.cc @@ -153,10 +153,9 @@ static StatusOr> GatherLoopBody( dim_numbers.index_vector_dim() == gather.operand(1)->shape().dimensions_size()); - TF_ASSIGN_OR_RETURN( - HloInstruction * induction_var_as_vector, + HloInstruction* induction_var_as_vector = MakeBroadcastHlo(induction_var, /*broadcast_dimensions=*/{}, - /*result_shape_bounds=*/{1})); + /*result_shape_bounds=*/{1}); HloInstruction* index_vector; @@ -222,7 +221,7 @@ static StatusOr> GatherLoopBody( {operand, start_indices, updated_accumulator}}; } -static StatusOr CreateGatherLoopAccumulatorInitValue( +static HloInstruction* CreateGatherLoopAccumulatorInitValue( HloComputation* computation, PrimitiveType element_type, absl::Span slice_sizes, int64 gather_loop_trip_count, const GatherDimensionNumbers& dim_numbers) { @@ -297,7 +296,7 @@ static StatusOr PermuteBatchAndOffsetDims( // [3,1] out of operand into an accumulator of shape [4,3,1]. We then // reshape this result to [2,2,3] and finally transpose it to [2,3,2]. -StatusOr GatherExpander::ExpandGather( +StatusOr GatherExpander::ExpandInstruction( HloInstruction* gather_instr) { CHECK(!ShapeUtil::IsZeroElementArray(gather_instr->shape())); @@ -332,12 +331,10 @@ StatusOr GatherExpander::ExpandGather( CHECK_EQ(gather_loop_trip_count, canonical_start_indices->shape().dimensions(0)); - TF_ASSIGN_OR_RETURN( - HloInstruction * accumulator_init, - CreateGatherLoopAccumulatorInitValue( - computation, output_shape.element_type(), - gather_instr->gather_slice_sizes(), gather_loop_trip_count, - gather_instr->gather_dimension_numbers())); + HloInstruction* accumulator_init = CreateGatherLoopAccumulatorInitValue( + computation, output_shape.element_type(), + gather_instr->gather_slice_sizes(), gather_loop_trip_count, + gather_instr->gather_dimension_numbers()); StatusOr> gather_loop_result_or_error = WhileUtil::MakeCountedLoop( @@ -364,25 +361,11 @@ StatusOr GatherExpander::ExpandGather( output_rank); } -StatusOr GatherExpander::Run(HloModule* module) { - auto is_nontrivial_gather = [](HloInstruction* inst) { - return inst->opcode() == HloOpcode::kGather && - // Avoid expanding gather ops that produce zero sized tensors, - // instead punt these to ZeroSizedHloElimination. - !ShapeUtil::IsZeroElementArray(inst->shape()); - }; - - std::vector gather_instrs; - for (HloComputation* computation : module->MakeNonfusionComputations()) { - absl::c_copy_if(computation->instructions(), - std::back_inserter(gather_instrs), is_nontrivial_gather); - } - - for (HloInstruction* inst : gather_instrs) { - TF_ASSIGN_OR_RETURN(HloInstruction * expanded_root, ExpandGather(inst)); - TF_RETURN_IF_ERROR(inst->parent()->ReplaceInstruction(inst, expanded_root)); - } - - return !gather_instrs.empty(); +bool GatherExpander::InstructionMatchesPattern(HloInstruction* inst) { + return inst->opcode() == HloOpcode::kGather && + // Avoid expanding gather ops that produce zero sized tensors, + // instead punt these to ZeroSizedHloElimination. + !ShapeUtil::IsZeroElementArray(inst->shape()); } + } // namespace xla diff --git a/tensorflow/compiler/xla/service/gather_expander.h b/tensorflow/compiler/xla/service/gather_expander.h index 8af9c6b71fbc391bf7c0e9809e979b65135a6df3..5625a37cb46ca5b70f69d86bc424f6512bfb293f 100644 --- a/tensorflow/compiler/xla/service/gather_expander.h +++ b/tensorflow/compiler/xla/service/gather_expander.h @@ -16,20 +16,22 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GATHER_EXPANDER_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_GATHER_EXPANDER_H_ -#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" +#include "tensorflow/compiler/xla/service/op_expander_pass.h" namespace xla { // This pass rewrites gather operations into (roughly) while loops of dynamic // slices. This lets backends that don't support gather directly to // nevertheless have a minimum level of support. -class GatherExpander : public HloModulePass { +class GatherExpander : public OpExpanderPass { public: absl::string_view name() const override { return "gather_expander"; } - StatusOr Run(HloModule* module) override; protected: - StatusOr ExpandGather(HloInstruction* gather_instr); + bool InstructionMatchesPattern(HloInstruction* instruction) override; + + StatusOr ExpandInstruction( + HloInstruction* gather_inst) override; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/generic_transfer_manager.cc b/tensorflow/compiler/xla/service/generic_transfer_manager.cc index 7d450f4b53cdea209f2ef10ba785be6ec3b8bf8d..cb43c27be961262bf29d4a3958de62cfada19aed 100644 --- a/tensorflow/compiler/xla/service/generic_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/generic_transfer_manager.cc @@ -26,7 +26,6 @@ limitations under the License. #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD index 6c23f921f40cac0dc5df08494dc1b63e6d1d5e93..05980fe549c4d9235ff80916cfad77ab60e1c447 100644 --- a/tensorflow/compiler/xla/service/gpu/BUILD +++ b/tensorflow/compiler/xla/service/gpu/BUILD @@ -3,6 +3,11 @@ load("//tensorflow/compiler/xla/tests:build_defs.bzl", "xla_test") load("//tensorflow/compiler/xla:xla.bzl", "xla_proto_library") +load( + "//tensorflow/core:platform/default/build_config_root.bzl", + "tf_cuda_tests_tags", +) +load("//tensorflow:tensorflow.bzl", "tf_cc_test") licenses(["notice"]) # Apache 2.0 @@ -24,12 +29,6 @@ filegroup( ]), ) -load("//tensorflow:tensorflow.bzl", "tf_cc_test") -load( - "//tensorflow/core:platform/default/build_config_root.bzl", - "tf_cuda_tests_tags", -) - xla_proto_library( name = "backend_configs", srcs = ["backend_configs.proto"], @@ -94,8 +93,8 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_reachability", - "//tensorflow/core:lib", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/memory", ], ) @@ -135,6 +134,8 @@ cc_library( "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", "//tensorflow/compiler/xla/service/llvm_ir:tuple_ops", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", "@llvm//:core", @@ -263,7 +264,9 @@ cc_library( "//tensorflow/compiler/xla/service:buffer_assignment", "//tensorflow/compiler/xla/service:device_memory_allocator", "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", "//tensorflow/core:stream_executor_no_cuda", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/types:span", ], @@ -302,6 +305,7 @@ cc_library( "sequential_thunk.cc", "thunk.cc", "thunk_schedule.cc", + "triangular_solve_thunk.cc", "tuple_thunk.cc", "while_thunk.cc", ], @@ -321,6 +325,7 @@ cc_library( "sequential_thunk.h", "thunk.h", "thunk_schedule.h", + "triangular_solve_thunk.h", "tuple_thunk.h", "while_thunk.h", ], @@ -361,7 +366,10 @@ cc_library( "//tensorflow/core/platform/default/build_config:cufft_plugin", "//tensorflow/core/platform/default/build_config:stream_executor_cuda", # build_cleaner: keep "//tensorflow/stream_executor", + "//tensorflow/stream_executor:blas", + "//tensorflow/stream_executor:device_memory", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", @@ -392,18 +400,21 @@ cc_library( srcs = ["cudnn_conv_algorithm_picker.cc"], hdrs = ["cudnn_conv_algorithm_picker.h"], deps = [ + ":autotuning_proto", ":backend_configs", ":buffer_comparator", ":cudnn_conv_runner", ":gpu_executable", ":ir_emission_utils", "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:protobuf_util", "//tensorflow/compiler/xla/service:compiler", "//tensorflow/compiler/xla/service:device_memory_allocator", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_casting_utils", "//tensorflow/compiler/xla/service:hlo_pass", "//tensorflow/core:lib", + "//tensorflow/core:logger", "//tensorflow/core:stream_executor_no_cuda", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", @@ -551,6 +562,44 @@ cc_library( ], ) +cc_library( + name = "gpu_sanitize_constant_names", + srcs = ["gpu_sanitize_constant_names.cc"], + hdrs = ["gpu_sanitize_constant_names.h"], + deps = [ + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_pass", + "//tensorflow/compiler/xla/service/llvm_ir:buffer_assignment_util", + "//tensorflow/core:lib", + ], +) + +tf_cc_test( + name = "gpu_sanitize_constant_names_test", + srcs = ["gpu_sanitize_constant_names_test.cc"], + tags = tf_cuda_tests_tags(), + deps = [ + ":gpu_sanitize_constant_names", + ":ir_emission_utils", + "//tensorflow/compiler/xla:shape_layout", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service:computation_layout", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_matchers", + "//tensorflow/compiler/xla/service:hlo_module_config", + "//tensorflow/compiler/xla/service:hlo_parser", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:test_utils", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:test", + "@com_google_absl//absl/strings", + ], +) + cc_library( name = "fusion_merger", srcs = ["fusion_merger.cc"], @@ -675,6 +724,7 @@ cc_library( ":gpu_hlo_schedule", ":gpu_hlo_support_checker", ":gpu_layout_assignment", + ":gpu_sanitize_constant_names", ":instruction_fusion", ":ir_emission_utils", ":ir_emitter", @@ -695,6 +745,8 @@ cc_library( "//tensorflow/compiler/xla/service:call_inliner", "//tensorflow/compiler/xla/service:conditional_simplifier", "//tensorflow/compiler/xla/service:convolution_group_converter", + "//tensorflow/compiler/xla/service:dot_decomposer", + "//tensorflow/compiler/xla/service:dynamic_index_splitter", "//tensorflow/compiler/xla/service:executable", "//tensorflow/compiler/xla/service:flatten_call_graph", "//tensorflow/compiler/xla/service:hlo", @@ -712,6 +764,7 @@ cc_library( "//tensorflow/compiler/xla/service:llvm_compiler", "//tensorflow/compiler/xla/service:reduce_precision_insertion", "//tensorflow/compiler/xla/service:reshape_mover", + "//tensorflow/compiler/xla/service:sort_simplifier", "//tensorflow/compiler/xla/service:transpose_folding", "//tensorflow/compiler/xla/service:tuple_simplifier", "//tensorflow/compiler/xla/service:while_loop_constant_sinking", @@ -725,6 +778,7 @@ cc_library( "//tensorflow/core:lib_internal", "//tensorflow/core:regexp_internal", "//tensorflow/core:stream_executor_no_cuda", + "//tensorflow/stream_executor/cuda:cuda_diagnostics", "@com_google_absl//absl/container:node_hash_map", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", @@ -1005,14 +1059,10 @@ cc_library( srcs = ["variadic_op_splitter.cc"], hdrs = ["variadic_op_splitter.h"], deps = [ - ":ir_emission_utils", - "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla:window_util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service:hlo", - "//tensorflow/compiler/xla/service:hlo_casting_utils", "//tensorflow/compiler/xla/service:hlo_pass", "//tensorflow/core:lib", "@com_google_absl//absl/strings", @@ -1038,3 +1088,12 @@ tf_cc_test( "//tensorflow/compiler/xla/tests:hlo_test_base", ], ) + +xla_proto_library( + name = "autotuning_proto", + srcs = ["autotuning.proto"], + deps = [ + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service:hlo_proto", + ], +) diff --git a/tensorflow/compiler/xla/service/gpu/autotuning.proto b/tensorflow/compiler/xla/service/gpu/autotuning.proto new file mode 100644 index 0000000000000000000000000000000000000000..b4a08963b4f2ebc55c89ed57325093536f343bd1 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/autotuning.proto @@ -0,0 +1,81 @@ +// This file defines protos that store the results of autotuning XLA:GPU +// operations. +// +// They are in proto format because we want to log them structured. They offer +// tremendous statistical, testing, and debugging value. +syntax = "proto3"; + +package xla.gpu; + +import "google/protobuf/duration.proto"; +import "tensorflow/compiler/xla/xla_data.proto"; +import "tensorflow/compiler/xla/service/hlo.proto"; + +message CudnnVersion { + int32 major = 1; + int32 minor = 2; + int32 patch = 3; +} + +message ComputeCapability { + int32 major = 1; + int32 minor = 2; +} + +message AutotuneResult { + message SuccessResult { + int64 scratch_bytes = 1; + google.protobuf.Duration run_time = 2; + } + + message ConvKey { + int64 algorithm = 1; + bool tensor_ops_enabled = 2; + } + + // If the conv runs successfully, success will be populated with the + // autotuning result. Otherwise, the error message is propagated. + oneof result { + SuccessResult success = 3; + string error_string = 4; + } + + oneof key { + ConvKey conv = 5; + } + + // Sometimes we run a correctness checker during autotuning. It compares the + // result buffer content between two algorithms, say, "reference" and "test" + // algorithms. The "test" algorithm is the one associated with this + // AutotuneResult. + // + // This field records the reference algorithm used. Notice that naming it + // "reference" doesn't mean it's always correct. However, empirically it's + // more correct, as it's "algo 0", less fancy than the compared one. + // + // Notice that the checker_failure may exist even in the success case. + // This is because the error string in `result` comes from the underlying + // implementation like cuDNN, which isn't aware that it produced an incorrect + // result. And even if the checker detects an incorrect result, we can still + // retrieve scratch_bytes and runtime_ms. + oneof checker_failure { + ConvKey reference_conv = 6; + } +} + +message AutotuneLog { + message Instruction { + xla.HloInstructionProto instruction = 1; + repeated xla.ShapeProto operand_shapes = 2; + } + + oneof instr_oneof { + Instruction instr = 1; + } + + // Records all auto-tuning results per algorithm. + repeated AutotuneResult results = 3; + + CudnnVersion cudnn_version = 4; + ComputeCapability compute_capability = 5; +} diff --git a/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc b/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc index 528209abc75777440163c2e1512658b8ad36315b..eb59ee5a1d47b6b706ef3f53a76069b3538eb6b7 100644 --- a/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc +++ b/tensorflow/compiler/xla/service/gpu/buffer_allocations.cc @@ -24,6 +24,7 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/gtl/map_util.h" #include "tensorflow/core/lib/strings/numbers.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" @@ -57,16 +58,16 @@ StatusOr> BufferAllocations::Builder::Build( // If buffer #i's address is already registered (e.g. external arguments or // result buffers), use that registered buffer. - if (registered_buffers_.count(i)) { - se::DeviceMemoryBase address = FindOrDie(registered_buffers_, i); - if (reinterpret_cast(address.opaque()) % expected_alignment != + if (se::DeviceMemoryBase* address = + tensorflow::gtl::FindOrNull(registered_buffers_, i)) { + if (reinterpret_cast(address->opaque()) % expected_alignment != 0) { return InternalError( "Address of registered buffer %d must be a multiple of %x, but " "was %p", - i, kEntryParameterAlignBytes, address.opaque()); + i, kEntryParameterAlignBytes, address->opaque()); } - buffer_allocations->SetBuffer(i, FindOrDie(registered_buffers_, i)); + buffer_allocations->SetBuffer(i, *address); continue; } diff --git a/tensorflow/compiler/xla/service/gpu/buffer_allocations.h b/tensorflow/compiler/xla/service/gpu/buffer_allocations.h index 14186b8faa68ad8492ea4863fcd7bd746e2eae48..9413ac2cff7c8d3ec4be6662569c580060bf1173 100644 --- a/tensorflow/compiler/xla/service/gpu/buffer_allocations.h +++ b/tensorflow/compiler/xla/service/gpu/buffer_allocations.h @@ -20,6 +20,7 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_map.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/device_memory_allocator.h" @@ -52,7 +53,8 @@ class BufferAllocations { DeviceMemoryAllocator* memory_allocator); private: - std::map registered_buffers_; + absl::flat_hash_map + registered_buffers_; }; ~BufferAllocations(); 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 6d6780fa1c7b0c636eb771c40e74f074cd8c4c4b..603af5a654589e0b02c762b57d70a8b7628b1d0f 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.cc @@ -16,14 +16,17 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" +#include "absl/time/time.h" #include "absl/types/optional.h" #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/protobuf_util.h" #include "tensorflow/compiler/xla/service/gpu/backend_configs.pb.h" #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/hlo_casting_utils.h" #include "tensorflow/core/lib/strings/numbers.h" +#include "tensorflow/core/platform/logger.h" #include "tensorflow/core/platform/mutex.h" namespace xla { @@ -32,7 +35,6 @@ namespace { using absl::optional; using se::DeviceMemoryBase; -using se::dnn::AlgorithmConfig; using se::dnn::AlgorithmDesc; class ScratchAllocator : public se::ScratchAllocator { @@ -132,6 +134,31 @@ tensorflow::mutex_lock LockGpu(const se::StreamExecutor* stream_exec) { return tensorflow::mutex_lock{it->second}; } +xla::gpu::CudnnVersion GetCudnnVersion(se::StreamExecutor* stream_executor) { + xla::gpu::CudnnVersion cudnn_version; + if (auto* dnn = stream_executor->AsDnn()) { + StatusOr version_or = dnn->GetVersion(); + if (version_or.ok()) { + const auto& version = version_or.ValueOrDie(); + cudnn_version.set_major(version.major_version()); + cudnn_version.set_minor(version.minor_version()); + cudnn_version.set_patch(version.patch()); + } + } + return cudnn_version; +} + +xla::gpu::ComputeCapability GetComputeCapability( + se::StreamExecutor* stream_executor) { + xla::gpu::ComputeCapability cc; + int cc_major, cc_minor; + stream_executor->GetDeviceDescription().cuda_compute_capability(&cc_major, + &cc_minor); + cc.set_major(cc_major); + cc.set_minor(cc_minor); + return cc; +} + } // anonymous namespace // We could have caching here so that we don't redo this work for two identical @@ -145,8 +172,8 @@ tensorflow::mutex_lock LockGpu(const se::StreamExecutor* stream_exec) { // cache misses and doing extra work. Overall, caching doesn't seem worth the // trouble, but we may want to revisit this if we ever find a model where // caching would speed up compilation a lot. -StatusOr -CudnnConvAlgorithmPicker::PickBestAlgorithm(HloCustomCallInstruction* instr) { +StatusOr CudnnConvAlgorithmPicker::PickBestAlgorithm( + const HloCustomCallInstruction* instr) { // TODO(timshen): for now only check fp16. It can be expanded to other types, // with some work on the HLO routines. const bool cross_check_enabled = @@ -232,8 +259,6 @@ CudnnConvAlgorithmPicker::PickBestAlgorithm(HloCustomCallInstruction* instr) { &stream, ShapeUtil::ByteSizeOf(instr->shape().tuple_shapes(0)))); initialize_buffer(result_buffer); - se::dnn::ProfileResult best_result; - int64 best_result_bytes_used = 0; TF_ASSIGN_OR_RETURN(auto backend_config, instr->backend_config()); @@ -243,82 +268,119 @@ CudnnConvAlgorithmPicker::PickBestAlgorithm(HloCustomCallInstruction* instr) { // this algorithm considered correct, though. optional first_algorithm; TF_ASSIGN_OR_RETURN(CudnnConvKind kind, GetCudnnConvKind(instr)); + std::vector profile_results; for (const AlgorithmDesc& alg : GetAlgorithms(kind, stream_exec_)) { ScratchAllocator scratch_allocator(device_ordinal, allocator); se::dnn::ProfileResult profile_result; VLOG(3) << "Trying algorithm " << AlgorithmToString(alg) << " for " << instr->ToString(); - backend_config.set_algorithm(alg.algo_id()); - backend_config.set_tensor_ops_enabled(alg.tensor_ops_enabled()); - TF_RETURN_IF_ERROR(instr->set_backend_config(backend_config)); - bool launch_ok = + // Use assignment instead of brace-list to make GCC 4.9 happy. + RunConvOptions options; + options.profile_result = &profile_result; + options.algo_override = alg; + Status launch_status = RunCudnnConv(instr, absl::MakeSpan(operand_buffers), result_buffer, - &scratch_allocator, &stream, &profile_result) - .ok(); - - if (launch_ok && profile_result.is_valid()) { - const bool crash_on_checking_failure = - instr->GetModule() - ->config() - .debug_options() - .xla_gpu_crash_on_verification_failures(); - if (comparator.has_value()) { - StatusOr result = comparator->CompareEqual( - se::DeviceMemory(result_buffer)); - if (!result.ok()) { - LOG(ERROR) << "Unable to compare " - << AlgorithmToString(*first_algorithm) << " against " - << AlgorithmToString(alg) << " for " << instr->ToString() - << ": " << result.status(); - CHECK(!crash_on_checking_failure); - } else if (!result.ValueOrDie()) { - LOG(ERROR) << "Results mismatch between different convolution " - "algorithms. This is likely a bug in convolution, or " - "an excessive loss of precision in convolution. " - << instr->ToString() << " for " - << AlgorithmToString(*first_algorithm) << " vs " - << AlgorithmToString(alg); - CHECK(!crash_on_checking_failure); - } - } else if (cross_check_enabled) { - auto comp = F16BufferComparator::Create( - se::DeviceMemory(result_buffer), compiler_, allocator, - &stream); - if (comp.ok()) { - comparator.emplace(comp.ConsumeValueOrDie()); - first_algorithm.emplace(alg); - } else { - LOG(ERROR) << "Fail to initialize buffer comparator: " - << comp.status() << ", instruction: " << instr->ToString(); - CHECK(!crash_on_checking_failure); - } + &scratch_allocator, &stream, options); + + profile_results.emplace_back(); + AutotuneResult& result = profile_results.back(); + result.mutable_conv()->set_algorithm(alg.algo_id()); + result.mutable_conv()->set_tensor_ops_enabled(alg.tensor_ops_enabled()); + + if (!launch_status.ok()) { + result.set_error_string(launch_status.error_message()); + continue; + } + + if (!profile_result.is_valid()) { + result.set_error_string("Invalid profile result"); + continue; + } + + int64 scratch_bytes_used = scratch_allocator.TotalAllocatedBytes(); + result.mutable_success()->set_scratch_bytes(scratch_bytes_used); + *result.mutable_success()->mutable_run_time() = + protobuf_util::ToDurationProto( + absl::Milliseconds(profile_result.elapsed_time_in_ms())); + + const bool crash_on_checking_failure = + instr->GetModule() + ->config() + .debug_options() + .xla_gpu_crash_on_verification_failures(); + + if (comparator.has_value()) { + StatusOr compare_result = comparator->CompareEqual( + se::DeviceMemory(result_buffer)); + if (!compare_result.ok()) { + LOG(ERROR) << "Unable to compare " + << AlgorithmToString(*first_algorithm) << " against " + << AlgorithmToString(alg) << " for " << instr->ToString() + << ": " << compare_result.status(); + CHECK(!crash_on_checking_failure); + } else if (!compare_result.ValueOrDie()) { + LOG(ERROR) << "Results mismatch between different convolution " + "algorithms. This is likely a bug in convolution, or " + "an excessive loss of precision in convolution. " + << instr->ToString() << " for " + << AlgorithmToString(*first_algorithm) << " vs " + << AlgorithmToString(alg); + CHECK(!crash_on_checking_failure); + auto* failure = result.mutable_reference_conv(); + failure->set_algorithm(first_algorithm->algo_id()); + failure->set_tensor_ops_enabled(first_algorithm->tensor_ops_enabled()); } - int64 scratch_bytes_used = scratch_allocator.TotalAllocatedBytes(); - VLOG(3) << "Run of algorithm " << AlgorithmToString(alg) - << " succeeded, taking " << profile_result.elapsed_time_in_ms() - << "ms and using " << NumBytesToString(scratch_bytes_used) - << " of scratch (Best result: " - << best_result.elapsed_time_in_ms() << "ms, " - << NumBytesToString(best_result_bytes_used) << " of scratch)"; - if (profile_result.elapsed_time_in_ms() < - best_result.elapsed_time_in_ms()) { - best_result = profile_result; - best_result_bytes_used = scratch_bytes_used; + } else if (cross_check_enabled) { + auto comp = F16BufferComparator::Create( + se::DeviceMemory(result_buffer), compiler_, allocator, + &stream); + if (comp.ok()) { + comparator.emplace(comp.ConsumeValueOrDie()); + first_algorithm.emplace(alg); + } else { + LOG(ERROR) << "Fail to initialize buffer comparator: " << comp.status() + << ", instruction: " << instr->ToString(); + CHECK(!crash_on_checking_failure); } - } else { - VLOG(3) << "Run of algorithm " << AlgorithmToString(alg) << " failed."; } } - if (best_result.is_valid()) { - VLOG(2) << "Best algorithm for " << instr->ToString() << ": " - << AlgorithmToString(best_result.algorithm()) << ", takes " - << best_result.elapsed_time_in_ms() << "ms, and uses " - << best_result_bytes_used << "B of scratch memory."; - return AutotuneResult{best_result.algorithm().algo_id(), - best_result.algorithm().tensor_ops_enabled(), - best_result_bytes_used, - absl::Milliseconds(best_result.elapsed_time_in_ms())}; + + // Log the autotuning result. + { + AutotuneLog log; + *log.mutable_instr()->mutable_instruction() = instr->ToProto(); + for (const auto* op : instr->operands()) { + *log.mutable_instr()->add_operand_shapes() = op->shape().ToProto(); + } + for (const auto& profile : profile_results) { + *log.add_results() = profile; + } + *log.mutable_compute_capability() = GetComputeCapability(stream_exec_); + *log.mutable_cudnn_version() = GetCudnnVersion(stream_exec_); + VLOG(2) << "Autotuning result:\n" << log.DebugString(); + tensorflow::Logger::Singleton()->LogProto(log); + } + + auto* profile_results_end = profile_results.data() + profile_results.size(); + + const AutotuneResult* best_result = std::min_element( + profile_results.data(), profile_results_end, + [](const AutotuneResult& lhs, const AutotuneResult& rhs) { + // The successful one should have a smaller key, since we are doing + // min_element. If they are both unsuccessful, keep the earlier one in + // the vector by comparing pointers. + return std::make_tuple( + !lhs.has_success(), + protobuf_util::FromDurationProto(lhs.success().run_time()), + &lhs) < std::make_tuple(!rhs.has_success(), + protobuf_util::FromDurationProto( + rhs.success().run_time()), + &rhs); + }); + + if (best_result != profile_results_end && best_result->has_success()) { + return *best_result; } return InternalError( @@ -339,22 +401,23 @@ StatusOr CudnnConvAlgorithmPicker::RunOnInstruction( } auto best_algo = std::move(best_algo_or).ValueOrDie(); - VLOG(1) << "Setting cudnn conv to use algorithm " << best_algo.algorithm - << " and " << NumBytesToString(best_algo.scratch_bytes) + VLOG(1) << "Setting cudnn conv to use algorithm " + << best_algo.conv().algorithm() << " and " + << NumBytesToString(best_algo.success().scratch_bytes()) << " of scratch memory: " << instr->ToString() - << " tensor_ops_enabled: " << best_algo.tensor_ops_enabled; + << " tensor_ops_enabled: " << best_algo.conv().tensor_ops_enabled(); // Replace instr with a new CustomCall which has the correct algorithm, and // whose output shape has the appropriate amount of scratch memory. HloComputation* computation = instr->parent(); Shape new_call_shape = ShapeUtil::MakeTupleShape( {instr->shape().tuple_shapes(0), - ShapeUtil::MakeShape(U8, {best_algo.scratch_bytes})}); + ShapeUtil::MakeShape(U8, {best_algo.success().scratch_bytes()})}); TF_ASSIGN_OR_RETURN(CudnnConvBackendConfig backend_config, instr->backend_config()); - backend_config.set_algorithm(best_algo.algorithm); - backend_config.set_tensor_ops_enabled(best_algo.tensor_ops_enabled); + backend_config.set_algorithm(best_algo.conv().algorithm()); + backend_config.set_tensor_ops_enabled(best_algo.conv().tensor_ops_enabled()); HloInstruction* new_call = computation->AddInstruction( instr->CloneWithNewOperands(new_call_shape, instr->operands())); diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.h b/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.h index 642af787afc71586d722ecc7e529ed8b3fa64d33..2e34ba9672314a62290b8a557960a605a98996c7 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.h +++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.h @@ -20,6 +20,7 @@ limitations under the License. #include "absl/types/optional.h" #include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/service/device_memory_allocator.h" +#include "tensorflow/compiler/xla/service/gpu/autotuning.pb.h" #include "tensorflow/compiler/xla/service/gpu/cudnn_conv_runner.h" #include "tensorflow/compiler/xla/service/hlo_instructions.h" #include "tensorflow/compiler/xla/service/hlo_module.h" @@ -47,16 +48,10 @@ class CudnnConvAlgorithmPicker : public HloModulePass { StatusOr Run(HloModule* module) override; private: - struct AutotuneResult { - int64 algorithm; - bool tensor_ops_enabled; - int64 scratch_bytes; - absl::Duration runtime; - }; - StatusOr RunOnComputation(HloComputation* computation); StatusOr RunOnInstruction(HloInstruction* instr); - StatusOr PickBestAlgorithm(HloCustomCallInstruction* instr); + StatusOr PickBestAlgorithm( + const HloCustomCallInstruction* instr); se::StreamExecutor* stream_exec_; // never null DeviceMemoryAllocator* allocator_; // may be null diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_conv_runner.cc b/tensorflow/compiler/xla/service/gpu/cudnn_conv_runner.cc index 3425e1b4942aaf1011ba1bf1c50dd7e79c1f9807..b628f27f4b2ba8ccf17fd531d8a0c25cb99d9396 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_runner.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_runner.cc @@ -395,32 +395,36 @@ Status RunCudnnConv(const HloCustomCallInstruction* conv, absl::Span operand_buffers, se::DeviceMemoryBase result_buffer, se::DeviceMemoryBase scratch_buf, se::Stream* stream, - se::dnn::ProfileResult* profile_result) { + RunConvOptions options) { ScratchBufAllocator scratch_allocator(scratch_buf); return RunCudnnConv(conv, operand_buffers, result_buffer, &scratch_allocator, - stream, profile_result); + stream, options); } Status RunCudnnConv(const HloCustomCallInstruction* conv, absl::Span operand_buffers, se::DeviceMemoryBase result_buffer, se::ScratchAllocator* scratch_allocator, se::Stream* stream, - se::dnn::ProfileResult* profile_result) { + RunConvOptions options) { TF_ASSIGN_OR_RETURN(CudnnConvParams params, GetCudnnConvParams(conv, operand_buffers, result_buffer)); + if (options.algo_override) { + params.algorithm = AlgorithmConfig(*options.algo_override); + } + PrimitiveType output_primitive_type = conv->shape().tuple_shapes(0).element_type(); switch (output_primitive_type) { case F16: return RunCudnnConvImpl(params, scratch_allocator, stream, - profile_result); + options.profile_result); case F32: return RunCudnnConvImpl(params, scratch_allocator, stream, - profile_result); + options.profile_result); case F64: return RunCudnnConvImpl(params, scratch_allocator, stream, - profile_result); + options.profile_result); default: LOG(FATAL) << ShapeUtil::HumanString(*params.output_shape); } diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_conv_runner.h b/tensorflow/compiler/xla/service/gpu/cudnn_conv_runner.h index edbc75a94a1238540390b93f0fa5217852c7781f..25b2461ca61251c6cb7b89b1f91da0f1636a3647 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_runner.h +++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_runner.h @@ -28,6 +28,14 @@ limitations under the License. namespace xla { namespace gpu { +struct RunConvOptions { + // Nullable output-parameter pointer for profiling results. + se::dnn::ProfileResult* profile_result = nullptr; + + // Use this algorithm, instead of the one from the instruction. + absl::optional algo_override; +}; + // This file contains low-level routines for running cudnn convolutions. // Calls into cudnn to run the specified convolution. @@ -46,13 +54,13 @@ Status RunCudnnConv(const HloCustomCallInstruction* conv, absl::Span operand_buffers, se::DeviceMemoryBase result_buffer, se::DeviceMemoryBase scratch_buf, se::Stream* stream, - se::dnn::ProfileResult* profile_result = nullptr); + RunConvOptions = {}); Status RunCudnnConv(const HloCustomCallInstruction* conv, absl::Span operand_buffers, se::DeviceMemoryBase result_buffer, se::ScratchAllocator* scratch_allocator, se::Stream* stream, - se::dnn::ProfileResult* profile_result = nullptr); + RunConvOptions = {}); } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc index 2ab754a471070d5f90a3eaebd0600ff180d2fe5d..80ddb3e5cde708f535f71b75587171ed975f939c 100644 --- a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc @@ -308,9 +308,11 @@ llvm::Value* GpuElementalIrEmitter::EmitDeviceFunctionCall( false); // No variadic arguments. // Declares the callee if it is not declared already. - llvm::Function* callee = llvm::cast( - b_->GetInsertBlock()->getModule()->getOrInsertFunction( - llvm_ir::AsStringRef(callee_name), callee_type)); + llvm::Function* callee = llvm::dyn_cast( + b_->GetInsertBlock() + ->getModule() + ->getOrInsertFunction(llvm_ir::AsStringRef(callee_name), callee_type) + .getCallee()); for (auto attribute : attributes) { callee->addFnAttr(attribute); @@ -446,7 +448,7 @@ llvm_ir::ElementGenerator GpuElementalIrEmitter::MakeElementGenerator( return Load(accum_ptr); }; case HloOpcode::kReduce: - // TODO(b/112040122): This should be supported. + // TODO(b/118332391): This should be supported. CHECK_EQ(hlo->operand_count(), 2) << "Did not expect variadic reduce"; return [=, &operand_to_generator]( const IrArray::Index& output_index) -> StatusOr { diff --git a/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc b/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc index b8fbe7d2bcb6674ce0ca366e5d5baf5484729f6c..86c9bc6a345047fb5329af0be45c8981cc427f50 100644 --- a/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc @@ -206,6 +206,8 @@ auto GetGemmFn(PrimitiveType type) -> decltype(&DoGemm) { return &DoGemm; case C64: return &DoGemm>; + case C128: + return &DoGemm>; default: LOG(FATAL) << "Unsupported type."; } @@ -221,6 +223,8 @@ auto GetGemmWithAlgorithmFn(PrimitiveType type) return &DoGemmWithAlgorithm; case C64: return &DoGemmWithAlgorithm>; + case C128: + return &DoGemmWithAlgorithm>; default: LOG(FATAL) << "Unsupported type."; } @@ -235,6 +239,8 @@ auto GetGemmAutotuneFn(PrimitiveType type) -> decltype(&DoGemmAutotune) { return &DoGemmAutotune; case C64: return &DoGemmAutotune>; + case C128: + return &DoGemmAutotune>; default: LOG(FATAL) << "Unsupported type."; } @@ -255,6 +261,8 @@ se::blas::ComputationType GetBlasComputationType(PrimitiveType type) { return se::blas::ComputationType::kF64; case C64: return se::blas::ComputationType::kComplexF32; + case C128: + return se::blas::ComputationType::kComplexF64; default: LOG(FATAL) << "Unsupported type."; } diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc index 128cecd765bb1a746eabff1f7a3d53158eb5a6b9..434060ad89dac7ad65c790c8c0a7f3d6ad62a25a 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc @@ -310,12 +310,34 @@ StatusOr GpuExecutable::ExecuteOnStream( TF_ASSIGN_OR_RETURN( const BufferAllocation::Slice slice, this->assignment_->GetUniqueSlice(src_hlo, sources[0]->index())); - CHECK(!slice.allocation()->is_entry_computation_parameter()); se::DeviceMemoryBase src_base = buffer_allocations->GetDeviceAddress(slice.index()); CHECK(!src_base.is_null() || src_base.size() == 0); - *device_memory = src_base; + if (!slice.allocation()->is_entry_computation_parameter()) { + // If the buffer coming out of the result is from a parameter, it + // means the caller aliased some parameter buffer to an output one + // (via the HloInputOutputAliasConfig API). If that is the case, the + // caller will receive a partially complete scoped shaped buffer, + // which they will have to fill up on return. + // Unfortunately the interface to the execute APIs are ShapedBuffer + // pointer based, which assumes caller ownership, and hence a buffer + // coming from there cannot be part of the new ScopedShapedBuffer we + // create for the result (which assumes ownership). + *device_memory = src_base; + } else { + const HloInputOutputAliasConfig& input_output_alias = + module().input_output_alias_config(); + auto output_alias = input_output_alias.GetAliasedOutput( + slice.allocation()->parameter_number(), + slice.allocation()->param_shape_index()); + CHECK(output_alias) + << "Ouput buffer is coming from parameter " + << slice.allocation()->parameter_number() << " at index " + << slice.allocation()->param_shape_index() + << ", but no alias exists"; + CHECK_EQ(*output_alias, index); + } buffers_in_result.insert(src_base); return Status::OK(); })); diff --git a/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker.cc b/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker.cc index 4268fb2c7a813b3b53e4cd48746028a7b369f28e..4765f67c4b17e97419182e341573f75ad3d6ac30 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker.cc @@ -17,7 +17,6 @@ limitations under the License. #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/errors.h" namespace xla { diff --git a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.cc b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.cc index 58bdd4209a2315cdb7d29e920faded4d1a6a5876..a6d80f0b6dddb3d8d0fd00c639e11c71da6a9f09 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.cc @@ -240,6 +240,32 @@ Status GpuLayoutAssignment::AddBackendConstraints( TF_RETURN_IF_ERROR( constraints->SetBufferLayout(keys_layout, *output_buffer)); } + } else if (instruction->opcode() == HloOpcode::kTriangularSolve) { + // 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, + // although for the 'a' argument we could potentially accept row-major + // order and fold the transpose into the operator. + auto set_fortran_layout = [](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)); + }; + Shape op0_shape = instruction->operand(0)->shape(); + Shape op1_shape = instruction->operand(1)->shape(); + Shape output_shape = instruction->shape(); + set_fortran_layout(&op0_shape); + set_fortran_layout(&op1_shape); + set_fortran_layout(&output_shape); + TF_RETURN_IF_ERROR( + constraints->SetOperandLayout(op0_shape, instruction, 0)); + TF_RETURN_IF_ERROR( + constraints->SetOperandLayout(op1_shape, instruction, 1)); + TF_RETURN_IF_ERROR( + constraints->SetInstructionLayout(output_shape, instruction)); } } return Status::OK(); diff --git a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc index 29756d27260b0f41b2dd4b649ea9b1610ff90268..391029e574622925b2a7e801a7d41d95e49a1cfb 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc @@ -368,12 +368,21 @@ TEST_F(LayoutAssignmentTest, DotLayout) { TEST_F(LayoutAssignmentTest, SortLayout) { const char* hlo_text = R"( HloModule SortLayout + + compare { + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + p.1.lhs = f32[] parameter(2) + p.1.rhs = f32[] parameter(3) + ROOT lt = pred[] less-than(p.0.lhs, p.0.rhs) + } + ENTRY sort { keys = f32[3,2]{0,1} constant({{0,1},{0,1},{0,1}}) values = f32[2,3]{1,0} parameter(0) transpose = f32[3,2]{1,0} transpose(values), dimensions={1,0} ROOT sort = (f32[3,2]{1,0}, f32[3,2]{1,0}) sort(keys, transpose), - dimensions={1} + dimensions={1}, to_apply=compare })"; TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, diff --git a/tensorflow/compiler/xla/service/gpu/gpu_sanitize_constant_names.cc b/tensorflow/compiler/xla/service/gpu/gpu_sanitize_constant_names.cc new file mode 100644 index 0000000000000000000000000000000000000000..5e38ceca18de30e0e1fa75a7a4bd865e000b7d22 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/gpu_sanitize_constant_names.cc @@ -0,0 +1,70 @@ +/* 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/gpu_sanitize_constant_names.h" + +#include +#include +#include + +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/llvm_ir/buffer_assignment_util.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/logging.h" + +namespace xla { + +namespace gpu { + +StatusOr GpuSanitizeConstantNames::Run(HloModule* module) { + bool changed = false; + + NameUniquer instr_name_uniquer(/*separator=*/"_"); + // Collect the names used for the non-constant HLO instructions.+ + for (HloComputation* computation : module->computations()) { + for (HloInstruction* instr : computation->instructions()) { + if (instr->opcode() == HloOpcode::kConstant) { + continue; + } + + const string& old_name = instr->name(); + instr->UniquifyName(&instr_name_uniquer); + CHECK_EQ(old_name, instr->name()); + } + } + + // Sanitize the names for the constant HLO instructions and make them unique. + // This is not merged into the above loop because we don't want this pass to + // change the names of non-constant instructions, that is, if a constant HLO + // conflicts with a non-constant HLO, we change the name of the constant HLO + // even though the non-constant HLO comes after in the HLO module. + for (HloComputation* computation : module->computations()) { + for (HloInstruction* instr : computation->instructions()) { + if (instr->opcode() != HloOpcode::kConstant) { + continue; + } + string sanitized_name = llvm_ir::SanitizeConstantName(*instr); + instr->SetAndSanitizeName(sanitized_name); + instr->UniquifyName(&instr_name_uniquer); + changed = true; + } + } + + return changed; +} // namespace gpu + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/gpu_sanitize_constant_names.h b/tensorflow/compiler/xla/service/gpu/gpu_sanitize_constant_names.h new file mode 100644 index 0000000000000000000000000000000000000000..8d583d047e25698e86032020b7fc20df87f5ab68 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/gpu_sanitize_constant_names.h @@ -0,0 +1,38 @@ +/* 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_GPU_SANITIZE_CONSTANT_NAMES_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_SANITIZE_CONSTANT_NAMES_H_ + +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" + +namespace xla { +namespace gpu { + +// Sanitizes HLO instruction names for the GPU backend. Currently, it only +// replaces . and - in the HLO constant instruction names with _ to please the +// LLVM PTX backend. +class GpuSanitizeConstantNames : public HloModulePass { + public: + absl::string_view name() const override { return "sanitize-constant-names"; } + + StatusOr Run(HloModule* module) override; +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_SANITIZE_CONSTANT_NAMES_H_ diff --git a/tensorflow/compiler/xla/service/gpu/gpu_sanitize_constant_names_test.cc b/tensorflow/compiler/xla/service/gpu/gpu_sanitize_constant_names_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..f5adee8cc61f18f356406d8c089dd43565957739 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/gpu_sanitize_constant_names_test.cc @@ -0,0 +1,82 @@ +/* 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 + +#include "tensorflow/compiler/xla/service/gpu/gpu_sanitize_constant_names.h" + +#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/service/hlo_module_config.h" +#include "tensorflow/compiler/xla/service/hlo_parser.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/core/platform/test.h" + +namespace xla { +namespace gpu { +namespace { + +namespace op = xla::testing::opcode_matchers; +using SanitizeConstantNamesTest = HloTestBase; + +TEST_F(SanitizeConstantNamesTest, InstructionNameWithHyphenSanitized) { + const char *const kHloString = R"( + HloModule HyphenInInstructionName + ENTRY kernelEntry { + ROOT equal-to = s32[2]{0} constant({42, 73}) + })"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(kHloString)); + + EXPECT_TRUE(GpuSanitizeConstantNames().Run(module.get()).ValueOrDie()); + HloInstruction *root = module->entry_computation()->root_instruction(); + EXPECT_EQ(root->name(), "equal_to"); +} + +TEST_F(SanitizeConstantNamesTest, InstructionNameWithDotSanitized) { + const char *const kHloString = R"( + HloModule HyphenInInstructionName + ENTRY kernelEntry { + ROOT equal.to = s32[2]{0} constant({42, 73}) + })"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(kHloString)); + + EXPECT_TRUE(GpuSanitizeConstantNames().Run(module.get()).ValueOrDie()); + HloInstruction *root = module->entry_computation()->root_instruction(); + EXPECT_EQ(root->name(), "equal_to"); +} + +TEST_F(SanitizeConstantNamesTest, BufferSanitizedNameCollisionResolved) { + const char *const kHloString = R"( + HloModule BufferSanitizedName + ENTRY kernelEntry { + equal.to = s32[2]{0} constant({42, 73}) + equal-to = s32[2]{0} constant({67, 3}) + ROOT equal_to = s32[2]{0} add(equal.to, equal-to) + })"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(kHloString)); + + EXPECT_TRUE(GpuSanitizeConstantNames().Run(module.get()).ValueOrDie()); + EXPECT_THAT(FindInstruction(module.get(), "equal_to_1"), op::Constant()); + EXPECT_THAT(FindInstruction(module.get(), "equal_to_2"), op::Constant()); +} + +} // namespace +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc b/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc index 8c6a6914792a96ab517fa5f20ff2215e4785490e..e593f535642e15f28a4a1c1f321881ba3c694548 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc @@ -30,7 +30,6 @@ limitations under the License. #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/gtl/cleanup.h" #include "tensorflow/core/platform/logging.h" diff --git a/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.cc b/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.cc index 9634c92786ad4dd0d62bae9cbb087b4ef6b5866f..69aaaceca112364a4fd562f6a5eff1629fd3fc54 100644 --- a/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.cc +++ b/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h" +#include "absl/container/flat_hash_set.h" #include "absl/strings/str_cat.h" #include "llvm/IR/BasicBlock.h" #include "llvm/IR/Function.h" @@ -45,10 +46,10 @@ void HloToIrBindings::EmitBasePointersForHlos( // An HLO can have duplicated operands. This data structure remembers which // operand HLOs are already bound to avoid rebinding the same HLO. - std::set already_bound_for_this_function; + absl::flat_hash_set already_bound_for_this_function; auto arg_iter = function->arg_begin(); for (const HloInstruction* io_hlo : io_hlos) { - if (!already_bound_for_this_function.count(io_hlo)) { + if (!already_bound_for_this_function.contains(io_hlo)) { if (!is_nested_ && io_hlo->opcode() == HloOpcode::kGetTupleElement) { BindHloToIrValue(*io_hlo, EmitGetTupleElement(io_hlo, &*arg_iter)); } else { @@ -63,7 +64,7 @@ void HloToIrBindings::EmitBasePointersForHlos( temp_buffer_base_->setName("temp_buffer"); for (const HloInstruction* non_io_hlo : non_io_hlos) { - if (already_bound_for_this_function.count(non_io_hlo)) { + if (already_bound_for_this_function.contains(non_io_hlo)) { continue; } already_bound_for_this_function.insert(non_io_hlo); diff --git a/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h b/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h index c0edae530cedba45c897b07b7b9cc72eaaab397c..f57b594e9c18078a3bbbf4d2b4db7e989c4edfdd 100644 --- a/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h +++ b/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h @@ -18,6 +18,7 @@ limitations under the License. #include +#include "absl/container/flat_hash_map.h" #include "absl/types/span.h" #include "llvm/IR/IRBuilder.h" #include "llvm/IR/Value.h" @@ -61,7 +62,7 @@ class HloToIrBindings { // Returns whether `hlo` is bound to an LLVM IR value. bool BoundToIrValue(const HloInstruction& hlo) const { - return base_ptrs_.count(&hlo); + return base_ptrs_.contains(&hlo); } llvm::Value* GetTempBufferBase() const { return temp_buffer_base_; } @@ -110,7 +111,8 @@ class HloToIrBindings { // For an instruction that generates multiple outputs, the root will be a // tuple shape. The IrArray for each element output is stored in the subnode // in the ShapeTree. - std::unordered_map> base_ptrs_; + absl::flat_hash_map> + base_ptrs_; // The address of the memory block that contains all temporary buffers. llvm::Value* temp_buffer_base_ = nullptr; diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc index 6151dd8ff4c92bb81bd756c68cc9377633c8c9d5..f07141029cbf8b034b74548f6fca8f1628589f0c 100644 --- a/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc @@ -282,22 +282,7 @@ bool GpuInstructionFusion::ShouldFuse(HloInstruction* consumer, bool GpuInstructionFusion::ShouldFuseIntoMultiOutput(HloInstruction* consumer, int64 operand_index) { - const HloInstruction* producer = consumer->operand(operand_index); - // The IR emitter has limited support for non-loop fusions with multi output - // at present. - // TODO(tjoerg): Relax this constraint to allow for arbitraty kinds of fusion. - if (consumer->opcode() == HloOpcode::kFusion && - consumer->fusion_kind() != HloInstruction::FusionKind::kLoop) { - return false; - } - // Multi-output fusion requires instructions with compatible shapes. - if (!ShapeUtil::Compatible(producer->shape(), consumer->shape())) { - return false; - } - // TODO(tjoerg): Stop calling `ShouldFuse` to relax the criteria for - // multi-output fusion. In particular, do not check whether an instruction is - // expensive to duplicate, since this doesn't matter here. - return GpuInstructionFusion::ShouldFuse(consumer, operand_index); + return false; } HloInstruction::FusionKind GpuInstructionFusion::ChooseKind( diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc index 688604cd36e5a45debf855aacd29d05ecda92341..a05ab86cf77a134a1fc387d93cb482aa1ff5345b 100644 --- a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc +++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc @@ -506,202 +506,11 @@ TEST_F(InstructionFusionTest, MultiOutputFusion) { })") .ValueOrDie(); - ASSERT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) - .Run(module.get()) - .ValueOrDie()); - SCOPED_TRACE(module->ToString()); - - // Expect that there is one multi-output fusion and subtract has not been - // duplicated. - EXPECT_EQ(Count(*module, HloOpcode::kFusion), 1); - EXPECT_EQ(Count(*module, HloOpcode::kSubtract), 1); - TF_ASSERT_OK_AND_ASSIGN( - const HloInstruction* fusion, - FindHloInstruction(*module->entry_computation(), HloOpcode::kFusion)); - EXPECT_THAT( - fusion->fused_expression_root(), - op::Tuple(op::Add(op::Subtract(), op::Parameter()), op::Subtract())); -} - -TEST_F(InstructionFusionTest, MultiOutputFusionExpensiveOp) { - // tanh --> add --> tuple - // \---------------/ - auto module = ParseHloString(R"( - HloModule test_module - ENTRY OutputFusion { - p0 = f32[4,3]{1,0} parameter(0) - p1 = f32[4,3]{1,0} parameter(1) - tanh = f32[4,3]{1,0} tanh(p0) - add = f32[4,3]{1,0} add(tanh, p1) - ROOT tuple = (f32[4,3]{1,0}, f32[4,3]{1,0}) tuple(tanh, add) - })") - .ValueOrDie(); - - // TODO(tjoerg): Allow multi-output fusion for expensive operations like tanh. + // Multi-output fusion is disabled here and performed in the + // GpuMultiOutputFusion pass instead. ASSERT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) .Run(module.get()) - .ValueOrDie()) - << module->ToString(); -} - -TEST_F(InstructionFusionTest, MultiOutputFusion2) { - // sub --> add1 --\--------\ - // \----------> add2 --> tuple - auto module = ParseHloString(R"( - HloModule test_module - ENTRY OutputFusion { - p0 = f32[4,3]{1,0} parameter(0) - p1 = f32[4,3]{1,0} parameter(1) - p2 = f32[4,3]{1,0} parameter(2) - sub = f32[4,3]{1,0} subtract(p0, p2) - add1 = f32[4,3]{1,0} add(sub, p1) - add2 = f32[4,3]{1,0} add(sub, add1) - ROOT tuple = (f32[4,3]{1,0}) tuple(add1, add2) - })") - .ValueOrDie(); - - ASSERT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) - .Run(module.get()) - .ValueOrDie()); - SCOPED_TRACE(module->ToString()); - - // Expect that there is one multi-output fusion and subtract has not been - // duplicated. - EXPECT_EQ(Count(*module, HloOpcode::kFusion), 1); - EXPECT_EQ(Count(*module, HloOpcode::kSubtract), 1); - TF_ASSERT_OK_AND_ASSIGN( - const HloInstruction* fusion, - FindHloInstruction(*module->entry_computation(), HloOpcode::kFusion)); - EXPECT_THAT(fusion->fused_expression_root(), - op::Tuple(op::Add(op::Subtract(), op::Add()), - op::Add(op::Subtract(), op::Parameter()))); -} - -TEST_F(InstructionFusionTest, MultiOutputFusion3) { - // sub --> add1 ----\--------\ - // \ --> add2 --> add3 --> tuple - auto module = ParseHloString(R"( - HloModule test_module - ENTRY OutputFusion { - p0 = f32[4,3]{1,0} parameter(0) - p1 = f32[4,3]{1,0} parameter(1) - p2 = f32[4,3]{1,0} parameter(2) - p3 = f32[4,3]{1,0} parameter(3) - sub = f32[4,3]{1,0} subtract(p0, p2) - add1 = f32[4,3]{1,0} add(sub, p1) - add2 = f32[4,3]{1,0} add(p2, sub) - add3 = f32[4,3]{1,0} add(add1, add2) - ROOT tuple = (f32[4,3]{1,0}) tuple(add3, add2) - })") - .ValueOrDie(); - - ASSERT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) - .Run(module.get()) - .ValueOrDie()); - SCOPED_TRACE(module->ToString()); - - // Expect that there is one multi-output fusion and subtract has not been - // duplicated. - EXPECT_EQ(Count(*module, HloOpcode::kFusion), 1); - EXPECT_EQ(Count(*module, HloOpcode::kSubtract), 1); - TF_ASSERT_OK_AND_ASSIGN( - const HloInstruction* fusion, - FindHloInstruction(*module->entry_computation(), HloOpcode::kFusion)); - EXPECT_THAT(fusion->fused_expression_root(), - op::Tuple(op::Add(op::Add(), op::Add()), - op::Add(op::Parameter(), op::Subtract()))); -} - -TEST_F(InstructionFusionTest, NoCyclesDueToMultiOutputFusion) { - // sub --> mul ---\ - // \--> call --> add --> tuple - auto module = ParseHloString(R"( - HloModule test_module - ENTRY OutputFusion { - c = f32[] constant(42) - p0 = f32[4,3]{1,0} parameter(0) - p1 = f32[4,3]{1,0} parameter(1) - sub = f32[4,3]{1,0} subtract(p0, p1) - mul = f32[4,3]{1,0} multiply(sub, c) - call = f32[4,3]{1,0} custom-call(sub), custom_call_target="foo" - add = f32[4,3]{1,0} add(mul, call) - ROOT tuple = (f32[4,3]{1,0}) tuple(add) - })") - .ValueOrDie(); - - ASSERT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) - .Run(module.get()) - .ValueOrDie()); - // Visit instructions in post order to detect cycles. - // TODO(tjoerg): Add cycle detection to the HloVerifier. - class DummyVisitor : public DfsHloVisitorWithDefault { - public: - DummyVisitor() {} - Status DefaultAction(HloInstruction* /*hlo_instruction*/) override { - return Status::OK(); - } - } visitor; - for (const HloComputation* computation : module->MakeComputationPostOrder()) { - // Accept will return a FailedPrecondition when a cycle is detected. - EXPECT_TRUE(computation->root_instruction()->Accept(&visitor).ok()); - } -} - -TEST_F(InstructionFusionTest, NoMultiOutputFusionWithIncompatibleShapes) { - // sub[2,3] --> add[4,3] --> tuple([2,3], [4,3]) - // \-------------------------/ - auto module = ParseHloString(R"( - HloModule test_module - ENTRY OutputFusion { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[4,3]{1,0} parameter(1) - p2 = f32[2,3]{1,0} parameter(2) - sub = f32[2,3]{1,0} subtract(p0, p2) - add = f32[4,3]{1,0} add(sub, p1) - ROOT tuple = (f32[2,3]{1,0}, f32[4,3]{1,0}) tuple(sub, add) - })") - .ValueOrDie(); - - // Multi-output fusion requires shapes to be compatible. Since `sub` and `add` - // have incompatible shapes, expect that no multi-output fusion happens. - ASSERT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) - .Run(module.get()) - .ValueOrDie()) - << module->ToString(); -} - -TEST_F(InstructionFusionTest, FuseIntoInputFusionInstruction) { - auto module = ParseHloString(R"( - HloModule test_module - - add_computation { - add_lhs = f32[] parameter(0) - add_rhs = f32[] parameter(1) - ROOT add_root = f32[] add(add_lhs, add_rhs) - } - - fused_computation { - p1 = f32[10] parameter(0) - zero = f32[] constant(0) - ROOT f2_root = f32[] reduce(p1, zero), dimensions={0}, - to_apply=add_computation - } - - ENTRY entry { - p0 = f32[10] parameter(0) - mul = f32[10] multiply(p0, p0) - fusion = f32[] fusion(mul), kind=kInput, calls=fused_computation - ROOT tuple = (f32[10], f32[]) tuple(fusion, mul) - })") - .ValueOrDie(); - - // Multi-output fusion is not supported for non-loop fusions at present. Since - // `fused_computation` is a input fusion, expect no multi-output fusion to - // happen. - ASSERT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) - .Run(module.get()) - .ValueOrDie()) - << module->ToString(); + .ValueOrDie()); } TEST_F(InstructionFusionTest, FuseScalarConstant) { diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc index 5d25a032a99b2e477ba81e9d8ea31b2e1ddf93ec..3ed6553f9205803cfa17772b890c449cfb457c89 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc @@ -20,7 +20,6 @@ limitations under the License. #include "llvm/IR/Module.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/service/gpu/backend_configs.pb.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" @@ -54,7 +53,8 @@ bool AreValidGemmShapes(const Shape& lhs_shape, const Shape& rhs_shape, PrimitiveType output_primitive_type = output_shape.element_type(); bool type_is_allowed = (output_primitive_type == F16 || output_primitive_type == F32 || - output_primitive_type == F64 || output_primitive_type == C64); + output_primitive_type == F64 || output_primitive_type == C64 || + output_primitive_type == C128); return type_is_allowed && IsRank2(lhs_shape, batch_dimensions_size) && IsRank2(rhs_shape, batch_dimensions_size) && IsRank2(output_shape, batch_dimensions_size) && @@ -154,20 +154,17 @@ bool IsReductionToVector(const HloInstruction& reduce) { const HloInstruction* input = reduce.operand(0); std::vector dims_to_keep; for (int64 dim = 0; dim < input->shape().dimensions().size(); ++dim) { - if (!std::count(reduce.dimensions().begin(), reduce.dimensions().end(), - dim)) { + if (!absl::c_linear_search(reduce.dimensions(), dim)) { dims_to_keep.push_back(dim); } } return LayoutUtil::AreDimensionsConsecutive(input->shape().layout(), dims_to_keep) && - ShapeUtil::Equal(reduce.shape(), ShapeUtil::FilterDimensions( - [&dims_to_keep](int64 dim) { - return std::count( - dims_to_keep.begin(), - dims_to_keep.end(), dim); - }, - input->shape())); + ShapeUtil::Equal( + reduce.shape(), + ShapeUtil::FilterDimensions( + [&](int64 dim) { return absl::c_count(dims_to_keep, dim); }, + input->shape())); } // This emits a device-side call to diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter.cc index 0007a9a8a3369d8ac010640127e1561615a6d813..cb13e727a44166ec564b58106c18b5c7f28a4af2 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter.cc @@ -647,7 +647,7 @@ Status IrEmitter::HandleParameter(HloInstruction* parameter) { } Status IrEmitter::HandleReduce(HloInstruction* reduce) { - // TODO(b/112040122): Support variadic reduce. + // TODO(b/118332391): Support variadic reduce. if (!reduce->shape().IsArray()) { return Unimplemented("Variadic reduce is not supported on GPU"); } diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc index 063d493d90a2330c535ed4a16d063d6db9b49cd1..0cc65ebb52737aa9bb8866eb07278a2319aa797b 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc @@ -15,6 +15,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -37,7 +38,6 @@ limitations under the License. #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor.h" -#include "tensorflow/compiler/xla/service/gpu/backend_configs.pb.h" #include "tensorflow/compiler/xla/service/gpu/buffer_allocations.h" #include "tensorflow/compiler/xla/service/gpu/conditional_thunk.h" #include "tensorflow/compiler/xla/service/gpu/convolution_thunk.h" @@ -59,6 +59,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/partition_assignment.h" #include "tensorflow/compiler/xla/service/gpu/sequential_thunk.h" #include "tensorflow/compiler/xla/service/gpu/thunk.h" +#include "tensorflow/compiler/xla/service/gpu/triangular_solve_thunk.h" #include "tensorflow/compiler/xla/service/gpu/tuple_thunk.h" #include "tensorflow/compiler/xla/service/gpu/while_thunk.h" #include "tensorflow/compiler/xla/service/hlo_casting_utils.h" @@ -88,6 +89,9 @@ namespace xla { namespace gpu { using llvm_ir::KernelMappingScheme; +using EmitElementFunction = + std::function; namespace { @@ -483,6 +487,41 @@ Status IrEmitterUnnested::HandleFft(HloInstruction* fft) { return Status::OK(); } +Status IrEmitterUnnested::HandleTriangularSolve(HloInstruction* hlo) { + auto has_fortran_layout = [](const Layout& layout) { + int n = layout.minor_to_major_size(); + return layout.minor_to_major(0) == n - 2 && + layout.minor_to_major(1) == n - 1; + }; + TF_RET_CHECK(has_fortran_layout(hlo->operand(0)->shape().layout())); + TF_RET_CHECK(has_fortran_layout(hlo->operand(1)->shape().layout())); + TF_RET_CHECK(has_fortran_layout(hlo->shape().layout())); + + std::vector> thunks; + + // Triangular solve is in-place on 'b', so copy 'b' to the output if they + // aren't the same buffer. + auto operand_buffer = GetAllocationSlice(*hlo->operand(1)); + auto destination_buffer = GetAllocationSlice(*hlo); + if (operand_buffer != destination_buffer) { + thunks.push_back(absl::make_unique( + /*source_address=*/operand_buffer, + /*destination_buffer=*/destination_buffer, + /*mem_size=*/ShapeUtil::ByteSizeOf(hlo->operand(1)->shape()), hlo)); + } + + thunks.push_back(BuildTriangularSolveThunk(hlo)); + + // 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), hlo)); + } + return Status::OK(); +} + Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) { HloInstruction* root = fusion->fused_expression_root(); if (HloInstruction::FusionKind::kInput == fusion->fusion_kind()) { @@ -542,7 +581,7 @@ Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) { // a 1D array. The specialized version requires a initializer thunk that // initializes the output array to the initial value of the reduce. if (root->opcode() == HloOpcode::kReduce && root->shape().IsTuple()) { - // TODO(b/112040122): Support variadic reduce. + // TODO(b/118332391): Support variadic reduce. return Unimplemented("Variadic reduce is not supported on GPU"); } return EmitReductionToVector(fusion); @@ -631,7 +670,7 @@ Status IrEmitterUnnested::EmitExtraOutputsForReduce( } Status IrEmitterUnnested::HandleReduce(HloInstruction* reduce) { - // TODO(b/112040122): Support multi-output reduce. + // TODO(b/118332391): Support multi-output reduce. if (!reduce->shape().IsArray()) { return Unimplemented("Multi-output reduce is not supported on GPU"); } @@ -955,18 +994,18 @@ Status IrEmitterUnnested::HandleScatter(HloInstruction* scatter) { BuildKernelThunk(scatter, /*implements_whole_instruction=*/thunks.empty())); - TF_RETURN_IF_ERROR( - EmitScatter(thunks.back().get(), scatter, - /*scatter_indices_gen=*/ - [=](const IrArray::Index& index) { - return GetIrArray(*scatter_indices, *scatter) - .EmitReadArrayElement(index, &b_, "scatter_index"); - }, - /*updates_gen=*/ - [=](const IrArray::Index& index) { - return GetIrArray(*updates, *scatter) - .EmitReadArrayElement(index, &b_, "update"); - })); + TF_RETURN_IF_ERROR(EmitScatter( + thunks.back().get(), scatter, + /*scatter_indices_gen=*/ + [=](const IrArray::Index& index) { + return GetIrArray(*scatter_indices, *scatter) + .EmitReadArrayElement(index, &b_, "scatter_index"); + }, + /*updates_gen=*/ + [=](const IrArray::Index& index) { + return GetIrArray(*updates, *scatter) + .EmitReadArrayElement(index, &b_, "update"); + })); // Elide the sequential thunk if there's no copy. if (thunks.size() == 1) { @@ -1114,17 +1153,7 @@ Status IrEmitterUnnested::HandleSort(HloInstruction* sort) { std::vector> thunks; Shape keys_shape = sort->operand(0)->shape(); int64 dimension_to_sort = sort->dimensions(0); - // In case there is a 'values' parameter that is a iota, we take note and use - // it later to ensure a stable sort. Otherwise, we don't guarantee a stable - // sort. - int64 iota_values_parameter_index = -1; for (int64 i = 0; i < sort->operand_count(); ++i) { - if (i > 0 && sort->operand(i)->opcode() == HloOpcode::kIota && - ShapeUtil::ElementIsIntegral(sort->operand(i)->shape()) && - Cast(sort->operand(i))->iota_dimension() == - dimension_to_sort) { - iota_values_parameter_index = i; - } ShapeIndex shape_index = sort->operand_count() > 1 ? ShapeIndex({i}) : ShapeIndex({}); // We assume that the layout of all involved operands and outputs is the @@ -1237,25 +1266,23 @@ Status IrEmitterUnnested::HandleSort(HloInstruction* sort) { : standard_launch_dimensions; UpdateLaunchDimensions(launch_dimensions, thunks.back().get(), ir_emitter_context_->llvm_module()); - IrArray keys_array; std::vector values_arrays; - values_arrays.reserve(sort->operand_count() - 1); + values_arrays.reserve(sort->operand_count()); for (int64 i = 0; i < sort->operand_count(); ++i) { ShapeIndex shape_index = sort->operand_count() > 1 ? ShapeIndex({i}) : ShapeIndex({}); - if (i == 0) { - keys_array = GetIrArray(*sort, *sort, shape_index); - } else { - values_arrays.push_back(GetIrArray(*sort, *sort, shape_index)); - } + values_arrays.push_back(GetIrArray(*sort, *sort, shape_index)); } return llvm_ir::EmitSortInPlace( - dimension_to_sort, keys_array, values_arrays, - iota_values_parameter_index, IrName(sort), xor_masks, &b_, + dimension_to_sort, values_arrays, IrName(sort), xor_masks, &b_, launch_dimensions, xor_masks.size() > 1 ? num_iterations_in_sort_dim : standard_num_iterations_in_sort_dim, - kTileSize); + kTileSize, + [&](absl::Span operands, llvm::Value* output) { + return EmitCallToNestedComputation(*sort->to_apply(), operands, + output); + }); }; std::vector xor_masks; for (int64 stage = 0; stage < num_stages; ++stage) { @@ -1506,10 +1533,10 @@ std::unique_ptr IrEmitterUnnested::BuildKernelThunk( return !allocation->is_constant(); }); - std::sort(non_constant_buffers.begin(), non_constant_buffers.end(), - [](const BufferAllocation* a, const BufferAllocation* b) { - return a->index() < b->index(); - }); + absl::c_sort(non_constant_buffers, + [](const BufferAllocation* a, const BufferAllocation* b) { + return a->index() < b->index(); + }); llvm::Function* kernel = BuildKernelPrototype(*inst, non_constant_buffers); @@ -1754,6 +1781,29 @@ std::unique_ptr IrEmitterUnnested::BuildFftThunk( /*output_shape=*/inst->shape(), inst); } +std::unique_ptr IrEmitterUnnested::BuildTriangularSolveThunk( + const HloInstruction* inst) { + const HloInstruction* a = inst->operand(0); + const HloInstruction* b = inst->operand(1); + int64 m = b->shape().dimensions(b->shape().rank() - 2); + int64 n = b->shape().dimensions(b->shape().rank() - 1); + int64 batch_size = std::accumulate( + b->shape().dimensions().begin(), b->shape().dimensions().end() - 2, + int64{1}, [](int64 a, int64 b) { return a * b; }); + int64 elem_size = + ShapeUtil::ByteSizeOfPrimitiveType(inst->shape().element_type()); + int64 a_batch_stride = inst->triangular_solve_options().left_side() + ? m * m * elem_size + : n * n * elem_size; + int64 b_batch_stride = m * n * elem_size; + return absl::make_unique( + inst->triangular_solve_options(), + /*a_input_buffer=*/GetAllocationSlice(*a), + /*b_input_buffer=*/GetAllocationSlice(*inst), + inst->shape().element_type(), batch_size, m, n, a_batch_stride, + b_batch_stride, inst); +} + StatusOr> IrEmitterUnnested::BuildInitializerThunk( HloInstruction* hlo, const ShapeIndex& index) { bool fused = HloOpcode::kFusion == hlo->opcode(); @@ -2077,12 +2127,36 @@ Status IrEmitterUnnested::EmitTargetElementLoopInThunk( return Status::OK(); } +namespace { + +// Returns true if the fusion contains any instruction that is likely +// translated to complex LLVM IR, such as loops, and prevent vectorization. +bool MayPreventVectorization(const HloInstruction& fusion_hlo) { + CHECK_EQ(fusion_hlo.opcode(), HloOpcode::kFusion); + return absl::c_any_of( + fusion_hlo.fused_instructions_computation()->instructions(), + [&](const HloInstruction* instr) { + switch (instr->opcode()) { + case HloOpcode::kReduce: + case HloOpcode::kReduceWindow: + case HloOpcode::kSort: + case HloOpcode::kDot: + return true; + default: + return false; + } + }); +} + +} // namespace + Status IrEmitterUnnested::EmitTargetElementLoop( const HloInstruction& hlo, const llvm_ir::ElementGenerator& element_generator) { int unroll_factor = 1; // Unfused elementwise operations are usually memory bound, unroll them. - if (hlo.IsElementwise() || hlo.opcode() == HloOpcode::kFusion) { + if (hlo.IsElementwise() || + (hlo.opcode() == HloOpcode::kFusion && !MayPreventVectorization(hlo))) { unroll_factor = ComputeMaxUnrollFactor(&hlo); } @@ -2105,7 +2179,6 @@ std::vector IrEmitterUnnested::ConstructIrArrayForInputs( return param_arrays; } - int IrEmitterUnnested::ConstructInputReducedShapeAndCastInputIrArrayToShape( const HloInstruction& hlo, const std::vector& param_arrays, const std::vector& param_buffers, @@ -2133,53 +2206,86 @@ int IrEmitterUnnested::ConstructInputReducedShapeAndCastInputIrArrayToShape( namespace { -void EmitFullElementalTile( - const KernelMappingScheme* mapping_scheme, - const IrArray::Index& tile_origin_index, const string& loop_name, - KernelSupportLibrary* ksl, llvm::IRBuilder<>* builder, llvm::Value* y, - llvm::Value* x, llvm::Type* index_ty, - const std::function& emit_elem_function) { +std::tuple GetStartOffsetAndStepForX( + int64 tile_size_x, int64 num_threads_x, + const KernelMappingScheme* mapping_scheme, llvm::IRBuilder<>* builder, + llvm::Value* x, llvm::Type* index_ty) { + llvm::Value* start_offset_x; + int64 step_x; + if (mapping_scheme->DilatedX()) { + start_offset_x = x; + step_x = num_threads_x; + } else { + start_offset_x = builder->CreateMul( + x, llvm::ConstantInt::get(index_ty, tile_size_x / num_threads_x)); + step_x = 1; + } + return std::make_tuple(start_offset_x, step_x); +} + +void EmitFullElementalTile(const KernelMappingScheme* mapping_scheme, + const IrArray::Index& tile_origin_index, + const string& loop_name, KernelSupportLibrary* ksl, + llvm::IRBuilder<>* builder, llvm::Value* y, + llvm::Value* x, llvm::Type* index_ty, + const EmitElementFunction& emit_elem_function) { int64 num_threads_x = mapping_scheme->GetNumberOfThreadsForDimensionX(); int64 num_threads_y = mapping_scheme->GetNumberOfThreadsForDimensionY(); int64 tile_size_x = mapping_scheme->GetTileSizeForDimensionX(); int64 tile_size_y = mapping_scheme->GetTileSizeForDimensionY(); + + llvm::Value* start_offset_x; + int64 step_x; + std::tie(start_offset_x, step_x) = GetStartOffsetAndStepForX( + tile_size_x, num_threads_x, mapping_scheme, builder, x, index_ty); + IrArray::Index source_idx = + tile_origin_index.AddOffsetToDim(y, KernelMappingScheme::DimY, builder) + .AddOffsetToDim(start_offset_x, KernelMappingScheme::DimX, builder); ksl->For(loop_name + "_y", /*start=*/llvm::ConstantInt::get(index_ty, 0), /*end=*/llvm::ConstantInt::get(index_ty, tile_size_y), /*step=*/llvm::ConstantInt::get(index_ty, num_threads_y), [&](llvm::Value* y_indvar) { - IrArray::Index source_idx_y = tile_origin_index.AddOffsetToDim( + IrArray::Index source_idx_y = source_idx.AddOffsetToDim( y_indvar, KernelMappingScheme::DimY, builder); llvm::Value* y_loc = builder->CreateAdd(y_indvar, y); - for (int64 j = 0; j < tile_size_x; j += num_threads_x) { - IrArray::Index source_idx = source_idx_y.AddOffsetToDim( - llvm::ConstantInt::get(index_ty, j), + + for (int64 j = 0; j < tile_size_x / num_threads_x; j++) { + IrArray::Index source_idx_y_x = source_idx_y.AddOffsetToDim( + llvm::ConstantInt::get(index_ty, j * step_x), KernelMappingScheme::DimX, builder); - llvm::Value* x_loc = - builder->CreateAdd(llvm::ConstantInt::get(index_ty, j), x); - emit_elem_function(source_idx, y_loc, x_loc); + llvm::Value* x_loc = builder->CreateAdd( + llvm::ConstantInt::get(index_ty, j * step_x), + start_offset_x); + emit_elem_function(source_idx_y_x, y_loc, x_loc, j); } }); } -void EmitPartialElementalTile( - const KernelMappingScheme* mapping_scheme, - const IrArray::Index& tile_origin_index, const string& loop_name, - KernelSupportLibrary* ksl, llvm::IRBuilder<>* builder, llvm::Value* y, - llvm::Value* x, llvm::Value* tile_height, llvm::Value* tile_width, - llvm::Type* index_ty, - const std::function& emit_elem_function) { +void EmitPartialElementalTile(const KernelMappingScheme* mapping_scheme, + const IrArray::Index& tile_origin_index, + const string& loop_name, + KernelSupportLibrary* ksl, + llvm::IRBuilder<>* builder, llvm::Value* y, + llvm::Value* x, llvm::Value* tile_height, + llvm::Value* tile_width, llvm::Type* index_ty, + const EmitElementFunction& emit_elem_function) { int64 num_threads_x = mapping_scheme->GetNumberOfThreadsForDimensionX(); int64 num_threads_y = mapping_scheme->GetNumberOfThreadsForDimensionY(); int64 tile_size_x = mapping_scheme->GetTileSizeForDimensionX(); - for (int64 j = 0; j < tile_size_x; j += num_threads_x) { - IrArray::Index source_idx = - tile_origin_index.AddOffsetToDim(llvm::ConstantInt::get(index_ty, j), - KernelMappingScheme::DimX, builder); - llvm::Value* x_loc = - builder->CreateAdd(llvm::ConstantInt::get(index_ty, j), x); + llvm::Value* start_offset_x; + int64 step_x; + std::tie(start_offset_x, step_x) = GetStartOffsetAndStepForX( + tile_size_x, num_threads_x, mapping_scheme, builder, x, index_ty); + IrArray::Index source_idx = + tile_origin_index.AddOffsetToDim(y, KernelMappingScheme::DimY, builder) + .AddOffsetToDim(start_offset_x, KernelMappingScheme::DimX, builder); + for (int64 j = 0; j < tile_size_x / num_threads_x; j++) { + IrArray::Index source_idx_x = + source_idx.AddOffsetToDim(llvm::ConstantInt::get(index_ty, j * step_x), + KernelMappingScheme::DimX, builder); + llvm::Value* x_loc = builder->CreateAdd( + llvm::ConstantInt::get(index_ty, j * step_x), start_offset_x); ksl->If( loop_name + "_x_in_tile", builder->CreateICmpULT(x_loc, tile_width), @@ -2199,14 +2305,13 @@ void EmitPartialElementalTile( /*step=*/llvm::ConstantInt::get(index_ty, num_threads_y), [&](llvm::Value* y_indvar) { llvm::Value* y_loc = builder->CreateAdd(y_indvar, y); - ksl->If( - loop_name + "_y_in_tile", - builder->CreateICmpULT(y_loc, tile_height), [&] { - emit_elem_function( - source_idx.AddOffsetToDim( - y_indvar, KernelMappingScheme::DimY, builder), - y_loc, x_loc); - }); + ksl->If(loop_name + "_y_in_tile", + builder->CreateICmpULT(y_loc, tile_height), [&] { + emit_elem_function( + source_idx_x.AddOffsetToDim( + y_indvar, KernelMappingScheme::DimY, builder), + y_loc, x_loc, j); + }); }); }); } @@ -2225,8 +2330,7 @@ void EmitTiledElementalCodeWithBoundsCheck( const IrArray::Index& tile_origin_index, const string& loop_name, KernelSupportLibrary* ksl, llvm::IRBuilder<>* builder, llvm::Value* y, llvm::Value* x, llvm::Value* tile_height, llvm::Value* tile_width, - const std::function& emit_elem_function) { + const EmitElementFunction& emit_elem_function) { int64 tile_size_x = mapping_scheme->GetTileSizeForDimensionX(); int64 tile_size_y = mapping_scheme->GetTileSizeForDimensionY(); llvm::Type* index_ty = tile_width->getType(); @@ -2262,7 +2366,7 @@ void EmitTiledElementalCodeWithBoundsCheck( void IrEmitterUnnested::EmitTileElementForCopy( HloInstruction* hlo, const llvm_ir::IrArray::Index& index, const KernelCodegenInfo* kernel_info, llvm::Value* y_loc, - llvm::Value* x_loc) { + llvm::Value* x_loc, int64 /*x_iter_num*/) { llvm_ir::TiledParameterInfo* tiled_param_info = kernel_info->GetTiledParameterInfo(); // TODO(jlebar): Add AA metadata to this load. @@ -2292,7 +2396,7 @@ void IrEmitterUnnested::EmitTileElementForCopy( void IrEmitterUnnested::EmitTileElementForFusion( HloInstruction* hlo, const llvm_ir::IrArray::Index& index, const KernelCodegenInfo* kernel_info, llvm::Value* y_loc, - llvm::Value* x_loc) { + llvm::Value* x_loc, int64 /*x_iter_num*/) { llvm_ir::TiledParameterInfo* tiled_param_info = kernel_info->GetTiledParameterInfo(); std::vector output_arrays = ConstructIrArrayForOutputs(*hlo); @@ -2393,6 +2497,23 @@ class ReductionCodegenInfo : public IrEmitterUnnested::KernelCodegenInfo { : llvm_ir::KernelMappingScheme::DimX; } + int GetNumberOfPartialResults() const { + if (IsRowReduction()) { + return 1; + } + int64 num_thread = mapping_scheme_->GetNumberOfThreadsForDimensionX(); + int64 tile_size = mapping_scheme_->GetTileSizeForDimensionX(); + CHECK_EQ(tile_size % num_thread, 0); + return tile_size / num_thread; + } + + int GetPartialResultIndex(int64 x_iter_num) const { + if (IsRowReduction()) { + return 0; + } + return x_iter_num; + } + private: AddressVector partial_result_addresses_; AddressVector reduction_input_addresses_; @@ -2452,10 +2573,11 @@ void IrEmitterUnnested::EmitPrologueForOneReduction( llvm::AllocaInst* reduction_input_address = Alloca(element_type); reduction_input_addresses->push_back(reduction_input_address); + int num_partial_results = reduction_info->GetNumberOfPartialResults(); AddressVector* partial_result_addresses = reduction_info->GetMutablePartialResultAddresses(); llvm::AllocaInst* partial_result_address = - Alloca(element_type, /*ArraySize=*/nullptr, + Alloca(element_type, /*ArraySize=*/b_.getInt32(num_partial_results), "partial_reduction_result." + llvm::Twine(reduce_idx)); partial_result_addresses->push_back(partial_result_address); @@ -2478,7 +2600,9 @@ void IrEmitterUnnested::EmitPrologueForOneReduction( .EmitReadArrayElement(IrArray::Index(b_.getInt32Ty()), &b_); } - Store(init_ir_value, partial_result_address); + for (int i = 0; i < num_partial_results; ++i) { + Store(init_ir_value, InBoundsGEP(partial_result_address, {b_.getInt32(i)})); + } } void IrEmitterUnnested::EmitPrologueForReduction( @@ -2516,10 +2640,14 @@ void IrEmitterUnnested::EmitPrologueForReduction( std::move(output_shape_index)); } - // Allocate stack storage to store the current output linear index and record - // the address of the storage. + int num_partial_results = reduction_info->GetNumberOfPartialResults(); + + // Allocate stack storage to store the linear indices for the current output, + // and record the address of the storage. reduction_info->SetCurrentOutputLinearIndexAddress( - Alloca(reduction_info->GetIndexType())); + Alloca(reduction_info->GetIndexType(), + /*ArraySize=*/b_.getInt32(num_partial_results), + "current_output_linear_index_address")); if (!reduction_info->IsRowReduction()) { llvm::Type* bool_ty = b_.getInt1Ty(); @@ -2589,36 +2717,45 @@ void IrEmitterUnnested::EmitEpilogueForReduction( llvm_ir::SetToFirstInsertPoint(if_output_inbound_data.true_block, &b_); } + int num_partial_results = reduction_info->GetNumberOfPartialResults(); + // Emit an atomic operation that accumulates the partial reduction to the // output element. For row reduction, this is only for lane 0 due to the // if-statement emitted above. for (int i = 0; i != num_reduces; ++i) { - IrArray::Index element_index( - /*linear=*/Load(reduction_info->GetCurrentOutputLinearIndexAddress(), - "output_linear_addr"), - ShapeUtil::GetSubshape(unnested_hlo->shape(), - reduction_output_shape_indices[i]), - &b_); - llvm::Value* output_address = - GetIrArray(*unnested_hlo, *unnested_hlo, - reduction_output_shape_indices[i]) - .EmitArrayElementAddress(element_index, &b_, - "output_element_address"); - // Do not emit atomic operations if each element in the reduction result is - // computed by one block, that is the dimension being reduced has only one - // block. - const llvm_ir::KernelMappingScheme* mapping_scheme = - reduction_info->GetKernelMappingScheme(); - if (mapping_scheme->GetTileBlockSizeForDimension( - llvm_ir::KernelMappingScheme::DimZ) == 1 && - mapping_scheme->GetTileBlockSizeForDimension( - reduction_info->GetReducedDimensionEnum()) == 1) { - TF_CHECK_OK(EmitCallToNestedComputation( - *reducers[i], {output_address, partial_result_addresses[i]}, - output_address)); - } else { - TF_CHECK_OK(EmitAtomicOperationForNestedComputation( - *reducers[i], output_address, partial_result_addresses[i])); + for (int j = 0; j < num_partial_results; ++j) { + IrArray::Index element_index( + /*linear=*/Load( + InBoundsGEP(reduction_info->GetCurrentOutputLinearIndexAddress(), + {b_.getInt32(j)}), + "output_linear_addr"), + ShapeUtil::GetSubshape(unnested_hlo->shape(), + reduction_output_shape_indices[i]), + &b_); + llvm::Value* output_address = + GetIrArray(*unnested_hlo, *unnested_hlo, + reduction_output_shape_indices[i]) + .EmitArrayElementAddress(element_index, &b_, + "output_element_address"); + // Do not emit atomic operations if each element in the reduction result + // is computed by one block, that is the dimension being reduced has only + // one block. + const llvm_ir::KernelMappingScheme* mapping_scheme = + reduction_info->GetKernelMappingScheme(); + if (mapping_scheme->GetTileBlockSizeForDimension( + llvm_ir::KernelMappingScheme::DimZ) == 1 && + mapping_scheme->GetTileBlockSizeForDimension( + reduction_info->GetReducedDimensionEnum()) == 1) { + TF_CHECK_OK(EmitCallToNestedComputation( + *reducers[i], + {output_address, + InBoundsGEP(partial_result_addresses[i], {b_.getInt32(j)})}, + output_address)); + } else { + TF_CHECK_OK(EmitAtomicOperationForNestedComputation( + *reducers[i], output_address, + InBoundsGEP(partial_result_addresses[i], {b_.getInt32(j)}))); + } } } } @@ -2626,7 +2763,7 @@ void IrEmitterUnnested::EmitEpilogueForReduction( void IrEmitterUnnested::EmitTileElementForReduction( HloInstruction* unnested_hlo, const llvm_ir::IrArray::Index& index, const KernelCodegenInfo* kernel_info, llvm::Value* y_loc, - llvm::Value* x_loc) { + llvm::Value* x_loc, int64 x_iter_num) { VLOG(10) << "Emit tile element for reduce " << unnested_hlo->ToString(); HloInstruction* reduce_or_tuple = unnested_hlo->opcode() == HloOpcode::kFusion ? unnested_hlo->fused_expression_root() @@ -2639,8 +2776,11 @@ void IrEmitterUnnested::EmitTileElementForReduction( // Record the linear address for the current reduction. const ReductionCodegenInfo* reduction_info = dynamic_cast(kernel_info); + int partial_result_index = reduction_info->IsRowReduction() ? 0 : x_iter_num; + Store(index[reduction_info->GetKeptDimensionEnum()], - reduction_info->GetCurrentOutputLinearIndexAddress()); + InBoundsGEP(reduction_info->GetCurrentOutputLinearIndexAddress(), + {b_.getInt32(partial_result_index)})); if (!reduction_info->IsRowReduction()) { llvm::Type* bool_ty = b_.getInt1Ty(); llvm::AllocaInst* output_inbound_addr = @@ -2687,6 +2827,13 @@ void IrEmitterUnnested::EmitTileElementForReduction( reduction_info->GetKernelMappingScheme()->GetUnnormalizedIndex( index, GetFirstReduceInstruction(output_instructions)->operand(0)->shape()); + int num_partial_results = reduction_info->GetNumberOfPartialResults(); + if (num_partial_results > 1) { + // Clear the linear index field of the IrArray::Index to enable the use of + // GetElementPointer with array types. This enables the vectorization of + // the computation for different partial results. + input_index.ClearLinearIndex(); + } absl::Span partial_reduction_result_addresses = reduction_info->GetPartialResultAddresses(); absl::Span reduction_input_addresses = @@ -2699,10 +2846,12 @@ void IrEmitterUnnested::EmitTileElementForReduction( for (int i = 0; i != reducers.size(); ++i) { llvm::Value* const input_ir_value = input_gens[i](input_index).ValueOrDie(); Store(input_ir_value, reduction_input_addresses[i]); + llvm::Value* partial_result_address = + InBoundsGEP(partial_reduction_result_addresses[i], + {b_.getInt32(partial_result_index)}); TF_CHECK_OK(EmitCallToNestedComputation( - *reducers[i], - {partial_reduction_result_addresses[i], reduction_input_addresses[i]}, - partial_reduction_result_addresses[i])); + *reducers[i], {partial_result_address, reduction_input_addresses[i]}, + partial_result_address)); } // Emit code to generate the output for the non-reduction instructions in the @@ -2713,8 +2862,8 @@ void IrEmitterUnnested::EmitTileElementForReduction( // Emits a kernel for the hlo instruction using the given tiling scheme. void IrEmitterUnnested::EmitBlock(const TileGenerator& emit_one_tile, - const KernelCodegenInfo* kernel_info, - KernelSupportLibrary& ksl, + KernelCodegenInfo* kernel_info, + KernelSupportLibrary* ksl, llvm::Type* index_ty) { KernelMappingScheme* mapping_scheme = kernel_info->GetKernelMappingScheme(); absl::Span dims_in_tile = mapping_scheme->GetDimensionsInTiles(); @@ -2747,15 +2896,14 @@ void IrEmitterUnnested::EmitBlock(const TileGenerator& emit_one_tile, llvm::Value* num_tiles_in_block = Select(ICmpEQ(last_block_for_dim, block_id_for_dim), last_block_size_for_dim, block_size_for_dim); - - ksl.For(loop_name, - /*start=*/index_typed_constant(0), - /*end=*/num_tiles_in_block, - /*step=*/1, [&](llvm::Value* block_dim_induction_var) { - IrArray::Index tile_index = starting_tile.AddOffsetToDim( - block_dim_induction_var, dim_id, &b_); - emit_next_block_dim(tile_index); - }); + ksl->For(loop_name, + /*start=*/index_typed_constant(0), + /*end=*/num_tiles_in_block, + /*step=*/1, [&](llvm::Value* block_dim_induction_var) { + IrArray::Index tile_index = starting_tile.AddOffsetToDim( + block_dim_induction_var, dim_id, &b_); + emit_next_block_dim(tile_index); + }); } }; @@ -2810,7 +2958,8 @@ void IrEmitterUnnested::EmitBlock(const TileGenerator& emit_one_tile, // unnested_hlo: The unnested hlo instruction for which the kernel is generated. // Currently, these hlo instructions are supported: kLoop fusion, kCopy. // tiled_param_ids: The IDs for the parameters that are 0-2-1 transpose of -// other tensors with the same dimensions and need to be tiled and tranposed. +// other tensors with the same dimensions and are safe to be tranposed via +// the shared memory tranpose implementation. // mapping_scheme: The tiling scheme to use. // kernel_generator: Contains function objects for code generation, such as // element generator, block prologue and epilogue generators. @@ -2863,12 +3012,11 @@ LaunchDimensions IrEmitterUnnested::EmitKernel( // since we touch threadIdx.x and blockIdx.x at the beginning of the kernel // *anyway*. if (!reduction_info && unnested_hlo->IsMultiOutputFusion()) { - KernelSupportLibrary{&b_}.If( - "emit_mof_tuple", IsBlock0Thread0(&b_), [&] { - llvm_ir::EmitTuple(GetIrArray(*unnested_hlo, *unnested_hlo), - ConstructIrArrayForOutputs(*unnested_hlo), &b_, - module_); - }); + KernelSupportLibrary{&b_}.If("emit_mof_tuple", IsBlock0Thread0(&b_), [&] { + llvm_ir::EmitTuple(GetIrArray(*unnested_hlo, *unnested_hlo), + ConstructIrArrayForOutputs(*unnested_hlo), &b_, + module_); + }); } // For each tiled parameter, cast its input IrArray to the corresponding @@ -2898,8 +3046,7 @@ LaunchDimensions IrEmitterUnnested::EmitKernel( auto emit_tiled_elemental_code_with_bounds_check = [&](const IrArray::Index& index, const string& loop_name, llvm::Value* tile_height, llvm::Value* tile_width, - const std::function& emit_elem_function) { + const EmitElementFunction& emit_elem_function) { EmitTiledElementalCodeWithBoundsCheck(mapping_scheme, index, loop_name, &ksl, &b_, y, x, tile_height, tile_width, emit_elem_function); @@ -2912,10 +3059,6 @@ LaunchDimensions IrEmitterUnnested::EmitKernel( const IrArray::Index input_tile_origin( Permute({0, 2, 1}, output_tile_origin.multidim())); - const IrArray::Index input_index = - input_tile_origin.AddOffsetToDim(x, KernelMappingScheme::DimX, &b_) - .AddOffsetToDim(y, KernelMappingScheme::DimY, &b_); - // If shared memory transpose is needed, wait for all threads to reach this // point, lest we copy a value from tile to output before the other thread // copies it from input to tile. This is `__syncthreads` in CUDA. @@ -2925,9 +3068,10 @@ LaunchDimensions IrEmitterUnnested::EmitKernel( // Note that tile_width and tile_height are flipped here because we are // reading a transposed tile. emit_tiled_elemental_code_with_bounds_check( - input_index, "input", output_tile_bounds[2], output_tile_bounds[1], + input_tile_origin, "input", output_tile_bounds[2], + output_tile_bounds[1], [&](const IrArray::Index& index, llvm::Value* y_loc, - llvm::Value* x_loc) { + llvm::Value* x_loc, int64 /*x_iter_num*/) { for (int64 id : tiled_param_ids) { IrArray& input_in_logical_shape = param_in_reduced_shape_arrays[id]; @@ -2947,18 +3091,15 @@ LaunchDimensions IrEmitterUnnested::EmitKernel( llvm_ir::TiledParameterInfo tiled_param_info(param_shmem_buffers, y, x); kernel_info->SetTiledParamInfo(&tiled_param_info); - const IrArray::Index output_index = - output_tile_origin.AddOffsetToDim(x, KernelMappingScheme::DimX, &b_) - .AddOffsetToDim(y, KernelMappingScheme::DimY, &b_); - // Write to output[index] by emitting code like normal, except that values // for the tiled parameters are read from the shmem buffers. emit_tiled_elemental_code_with_bounds_check( - output_index, "output", output_tile_bounds[1], output_tile_bounds[2], - [&](const IrArray::Index& index, llvm::Value* y_loc, - llvm::Value* x_loc) { - kernel_generator.GetTileElementGenerator()(unnested_hlo, index, - kernel_info, y_loc, x_loc); + output_tile_origin, "output", output_tile_bounds[1], + output_tile_bounds[2], + [&](const IrArray::Index& index, llvm::Value* y_loc, llvm::Value* x_loc, + int64 x_iter_num) { + kernel_generator.GetTileElementGenerator()( + unnested_hlo, index, kernel_info, y_loc, x_loc, x_iter_num); }); // If a tile block contains multiple tiles and shared memory buffers are @@ -2976,7 +3117,7 @@ LaunchDimensions IrEmitterUnnested::EmitKernel( block_prologue_generator(unnested_hlo, kernel_info); } - EmitBlock(std::move(emit_one_tile), kernel_info, ksl, index_ty); + EmitBlock(std::move(emit_one_tile), kernel_info, &ksl, index_ty); const BlockEpilogueGenerator& block_epilogue_generator = kernel_generator.GetBlockEpilogueGenerator(); @@ -2989,7 +3130,10 @@ LaunchDimensions IrEmitterUnnested::EmitKernel( // Emits a kernel for the given hlo instruction using a tiled 0-2-1 transpose // algorithm to improve the memory access patterns for the input parameters -// with a shape that is a 0-2-1 transpose of the output tensor shape. +// with a shape that is a 0-2-1 transpose of the output tensor shape. The caller +// is responsible for making sure that it is safe to apply the shared memory +// tranpose on the input parameters. +// // // For the purpose of tiling, the output tensors have a logical shape of three // components 0-2-1 while the relevant input parameters have a logical shape @@ -3022,17 +3166,19 @@ LaunchDimensions IrEmitterUnnested::EmitHlo021Tile( element_generator = [&](HloInstruction* hlo, const llvm_ir::IrArray::Index& index, const KernelCodegenInfo* kernel_info, - llvm::Value* y_loc, llvm::Value* x_loc) { - EmitTileElementForCopy(hlo, index, kernel_info, y_loc, x_loc); + llvm::Value* y_loc, llvm::Value* x_loc, + int64 x_iter_num) { + EmitTileElementForCopy(hlo, index, kernel_info, y_loc, x_loc, x_iter_num); }; } else { DCHECK_EQ(hlo->opcode(), HloOpcode::kFusion); - element_generator = [&](HloInstruction* hlo, - const llvm_ir::IrArray::Index& index, - const KernelCodegenInfo* kernel_info, - llvm::Value* y_loc, llvm::Value* x_loc) { - EmitTileElementForFusion(hlo, index, kernel_info, y_loc, x_loc); - }; + element_generator = + [&](HloInstruction* hlo, const llvm_ir::IrArray::Index& index, + const KernelCodegenInfo* kernel_info, llvm::Value* y_loc, + llvm::Value* x_loc, int64 x_iter_num) { + EmitTileElementForFusion(hlo, index, kernel_info, y_loc, x_loc, + x_iter_num); + }; } KernelCodegenInfo kernel_info(&mapping_scheme); KernelCodeGenerator kernel_generator(std::move(element_generator)); @@ -3040,26 +3186,99 @@ LaunchDimensions IrEmitterUnnested::EmitHlo021Tile( } namespace { -// Returns true to indicate it is safe to use the tile based shared memory -// transpose implementation to implement the kernel for the instruction. +// A recursive function to inspect the users of a parameter to determine +// whether it's safe for a parameter to participate in a shared-memory +// transpose. // -// An instruction is not safe for such an implementation if it can change the -// element order of a tensor without changing the dimension of the tensor, and -// the instruction has a corresponding elemental_ir_emitter. -bool IsInstructionSafeForTileBasedTranspose(const HloInstruction* hlo) { - auto is_safe_for_tile_based_transpose = [&](const HloInstruction* instr) { - HloOpcode opcode = instr->opcode(); - CHECK_NE(opcode, HloOpcode::kFusion); - return (opcode != HloOpcode::kReverse && opcode != HloOpcode::kGather); - }; +// Consider a fusion parameter P for which we might want to use a shmem +// transpose. If we do, we use a GPU thread block to preload a tile of P with +// indices [z, y..y+31, x..x+31] to compute an output tile with the same indices +// cooperatively, where z, y, x are the indices for the normalized input/output +// tensor (see the document for FindTranspose021 for the definition of +// normalized tensor for 0-2-1 transpose). This shmem transpose implementation +// requires that the computation of the output tile only read elements within +// the preload tile. If this is not true, we can't use a shmem transpose for P. +// +// If the computation of output element [z, y, x] only requires the element of +// P with the same indices, the shmem tranpose implementation can be applied +// to P safely. This is a sufficient but not necessary condition. We check all +// the transitive users of P to see if we can find a user that may cause an +// exception to the situation. If such a user is not found, we conclude that P +// is safe for shmem transpose. +// +// This is trivially true for elementwise operations and some "data-movement" +// ops like kTuple. However, it's not true for operations that can change the +// dimensions of the inputs (e.g. pad, slice) and bitcast operation. +// For example: +// +// fused_computation { +// param_0 = f32[64,64]{1,0} parameter(0) +// ROOT bitcast = f32[64,64]{0,1} bitcast(param_0) +// } +// The output element at logical address [0, 63] depends on the input element +// at logical address [63, 0], which would not be within the shared-memory +// block. +// +// TODO(bixia): In order to extend this for kInput fusion, that is reduction +// with tranpose, we only need to end the use-chain checking with the input of +// a reduce operations. In this case, the above description on "output" apply +// to the result of such a use-chain, which provides the input to the reduce +// operation. +bool IsInstructionSafeForShmemTranspose(const HloInstruction* hlo) { + if (hlo->IsElementwise()) { + return absl::c_all_of(hlo->users(), [&](const HloInstruction* user) { + return IsInstructionSafeForShmemTranspose(user); + }); + } - if (hlo->opcode() == HloOpcode::kFusion) { - return absl::c_all_of(hlo->fused_instructions_computation()->instructions(), - is_safe_for_tile_based_transpose); + switch (hlo->opcode()) { + // Non-elementwise instructions that don't cause the shmem transpose + // to be unsafe, including the instructions that don't currently fuse. + case HloOpcode::kGetDimensionSize: + // The result of the operation doesn't rely on the content of the + // tensor. As such, there is no need to further inspect its users. + return true; + case HloOpcode::kGetTupleElement: + case HloOpcode::kMap: + case HloOpcode::kParameter: + case HloOpcode::kTuple: + case HloOpcode::kTupleSelect: + return absl::c_all_of(hlo->users(), [&](const HloInstruction* user) { + return IsInstructionSafeForShmemTranspose(user); + }); + + default: + return false; } +} - return is_safe_for_tile_based_transpose(hlo); +// Given a group of input parameters that are 0-2-1 tranpose of the outputs of +// a fusion kernel, returns the input parameters that are safe for the shared +// memory tranpose implementation. +// +// When a tile based shared memory transpose is used to implement an input with +// 0-2-1 transpose, we preload a tile of the input elements +// [z, y..y+31, x..x+31] to compute the output tile elements of the same +// indices. Preloading the input tile this way is only safe when the computation +// of the output tile elements do not need any input element outside the +// preloaded tile. We inspect all the transitive users of the input parameter +// up to the fusion root instruction to see if we can find any instruction +// that can make preloading the input tile unsafe. +std::vector FilterInputsForShmemTranspose(const HloInstruction* fusion, + std::vector input_ids) { + std::vector filtered_input_ids; + for (int64 i = 0; i < input_ids.size(); ++i) { + const HloInstruction* input = fusion->fused_parameter(input_ids[i]); + if (IsInstructionSafeForShmemTranspose(input)) { + filtered_input_ids.push_back(input_ids[i]); + } else { + VLOG(10) << "Input not safe for shmem transpose " << input->ToString() + << "\n"; + } + } + return filtered_input_ids; } + } // namespace bool IrEmitterUnnested::CheckAndEmitHloWithTile021(HloInstruction* hlo) { @@ -3106,8 +3325,11 @@ bool IrEmitterUnnested::CheckAndEmitHloWithTile021(HloInstruction* hlo) { return false; } - if (!IsInstructionSafeForTileBasedTranspose(hlo)) { - return false; + if (opcode == HloOpcode::kFusion) { + params_012 = FilterInputsForShmemTranspose(hlo, params_012); + if (params_012.empty()) { + return false; + } } // Each of our shared memory tiles has 32*33 elements (so ~4kb, if the @@ -3250,11 +3472,101 @@ std::tuple GetReductionToVectorDimensions( return std::make_tuple(num_reduced_major, num_kept, num_reduced_minor); } +// Returns true if all the transitive users of hlo before hitting users in +// use_chain_endings are elementwise operations. +bool AreUsersElementwise(const HloInstruction* hlo, + const ConstHloInstructionSet& use_chain_endings) { + return absl::c_all_of(hlo->users(), [&](const HloInstruction* user) { + return use_chain_endings.count(user) || + (user->IsElementwise() && + AreUsersElementwise(user, use_chain_endings)); + }); +} + +// Returns the number of fusion inputs that have the same dimension as the +// given shape, and involve in only elementwise operations. +int64 NumInputsInvolveInOnlyElementwiseOps( + const HloInstruction* unnested_hlo, const Shape& op_shape, + const ConstHloInstructionSet& use_chain_endings) { + return absl::c_count_if( + unnested_hlo->fused_parameters(), [&](const HloInstruction* parameter) { + const Shape& parameter_shape = parameter->shape(); + return ShapeUtil::SameDimensions(op_shape, parameter_shape) && + AreUsersElementwise(parameter, use_chain_endings); + }); +} + +// Returns the number of fusion inputs that have more elements than the given +// shape. +int64 NumInputsWithMoreElementsThan(const HloInstruction* unnested_hlo, + const Shape& shape) { + int64 num_elements = ShapeUtil::ElementsIn(shape); + return absl::c_count_if( + unnested_hlo->fused_parameters(), [&](const HloInstruction* parameter) { + return ShapeUtil::ElementsIn(parameter->shape()) > num_elements; + }); +} + +// The benefit of unrolling a kInput fusion that is a column reduction comes +// from the vectorization of non-reduction fusion outputs and fusion inputs. +// On the other hand, unrolling can also introduce factors that can cause +// the kernel to run slower. This routine uses a simple heuristic to estimate +// the benefit as well as the overhead of unrolling in order to decide whether +// unrolling is beneficial for the given kInput fusion. +bool IsUnrollingColumnReductionBeneficial(const HloInstruction* unnested_hlo, + const Shape& input_shape, + int64 num_kept) { + // TODO(b/122468062): Need further investigate to see whether we can + // remove the constraint on IsPowerOfTwo. + if (!IsPowerOfTwo(static_cast(num_kept))) { + return false; + } + + if (unnested_hlo->opcode() == HloOpcode::kReduce) { + return true; + } + + CHECK_EQ(unnested_hlo->opcode(), HloOpcode::kFusion); + int64 can_be_vectorized = 0; + int64 cannot_be_vectorized = 0; + const HloInstruction* fused_root = unnested_hlo->fused_expression_root(); + ConstHloInstructionSet use_chain_endings; + if (fused_root->opcode() == HloOpcode::kReduce) { + use_chain_endings.insert(fused_root); + // Atomic.add of the reduction result can't be vectorized. + cannot_be_vectorized++; + } else { + CHECK_EQ(fused_root->opcode(), HloOpcode::kTuple); + for (const HloInstruction* instr : fused_root->operands()) { + if (instr->opcode() == HloOpcode::kReduce) { + // Atomic.add of the reduction result can't be vectorized. + cannot_be_vectorized++; + } else { + // Write of the non-reduction result can be vectorized. + can_be_vectorized++; + } + use_chain_endings.insert(instr); + } + } + // Fusion inputs that have the same dimension as the reduce input and + // only involve in elementwise operations can be vectorized. + can_be_vectorized += NumInputsInvolveInOnlyElementwiseOps( + unnested_hlo, input_shape, use_chain_endings); + // Fusion inputs with more elements than the reduce op input must participate + // in non-elementwise operations and we assume that they are not vectorizable + // for the purpose of estimating the benefit of unrolling. If the kernel is + // unrolled even with such an assumption, and the accesses to those inputs + // turn out to be vectorizable, the compiler will still vectorize them. + cannot_be_vectorized += + NumInputsWithMoreElementsThan(unnested_hlo, input_shape); + return can_be_vectorized >= cannot_be_vectorized; +} + } // namespace std::tuple IrEmitterUnnested::ComputeMappingSchemeAndReductionKind( - const HloInstruction* first_reduce) { + const HloInstruction* unnested_hlo, const HloInstruction* first_reduce) { int64 depth = 1; int64 height = 1; int64 width = 1; @@ -3271,6 +3583,7 @@ IrEmitterUnnested::ComputeMappingSchemeAndReductionKind( std::tie(num_reduced_major, num_kept, num_reduced_minor) = GetReductionToVectorDimensions(input_shape, first_reduce->dimensions()); CHECK_EQ(num_output_elems, num_kept); + bool dilated_x = true; if (num_kept == 1) { // Scalar reduction is a special row reduction with depth = height = 1. @@ -3285,13 +3598,21 @@ IrEmitterUnnested::ComputeMappingSchemeAndReductionKind( is_row_reduction = false; // Column reduction without transpose doesn't require communication among // threads processing elements in the same tile. The current implementation - // only support the use of on hardware thread block to process one block of - // tiles in the KernelMappingScheme. We try to maximize the values of + // only support the use of one hardware thread block to process one block of + // tiles in the KernelMappingScheme. We try to use one thread to compute + // the partial results for two tensor elements and to maximize the values of // num_threads_x and tile_size_x to allow a bigger hardware thread block. int64 hw_threads_per_block_limit = ThreadsPerBlockLimit(ir_emitter_context_->device_description()); - tile_size_x = std::min(hw_threads_per_block_limit, num_kept); - num_threads_x = tile_size_x; + if (IsUnrollingColumnReductionBeneficial(unnested_hlo, input_shape, + num_kept)) { + tile_size_x = std::min(2 * hw_threads_per_block_limit, num_kept); + num_threads_x = tile_size_x / 2; + dilated_x = false; + } else { + tile_size_x = std::min(hw_threads_per_block_limit, num_kept); + num_threads_x = tile_size_x; + } int64 kNumElementsPerPartialSum = 128; tile_size_y = kNumElementsPerPartialSum; } else { @@ -3320,6 +3641,7 @@ IrEmitterUnnested::ComputeMappingSchemeAndReductionKind( llvm_ir::KernelMappingScheme mapping_scheme( dims_in_elem, tile_size_y, tile_size_x, req_block_sizes, num_threads_y, num_threads_x, &b_); + mapping_scheme.SetDilatedX(dilated_x); return std::make_tuple(mapping_scheme, is_row_reduction); } @@ -3368,14 +3690,15 @@ Status IrEmitterUnnested::EmitReductionToVector(HloInstruction* unnested_hlo) { bool is_row_reduction; llvm_ir::KernelMappingScheme mapping_scheme; std::tie(mapping_scheme, is_row_reduction) = - ComputeMappingSchemeAndReductionKind(first_reduce); + ComputeMappingSchemeAndReductionKind(unnested_hlo, first_reduce); ReductionCodegenInfo reduction_info(&mapping_scheme, is_row_reduction); KernelCodeGenerator kernel_generator( /*tile_element_generator=*/ [&](HloInstruction* hlo, const llvm_ir::IrArray::Index& index, const KernelCodegenInfo* kernel_info, llvm::Value* y_loc, - llvm::Value* x_loc) { - EmitTileElementForReduction(hlo, index, kernel_info, y_loc, x_loc); + llvm::Value* x_loc, int64 x_iter_num) { + EmitTileElementForReduction(hlo, index, kernel_info, y_loc, x_loc, + x_iter_num); }, /*block_prologue_generator=*/ [&](HloInstruction* hlo, KernelCodegenInfo* kernel_info) { diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h index d217ee36cf6e9b5278024a2f78513232328e7538..f85e18bbf0798ef3d5b87e81d287d8aed691dfc4 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h @@ -76,7 +76,6 @@ class IrEmitterUnnested : public IrEmitter { void SetLaneId(llvm::Value* v) { lane_id_ = v; } void SetIndexType(llvm::Type* t) { index_ty_ = t; } void SetTiledParamInfo(llvm_ir::TiledParameterInfo* tiled_param_info) { - CHECK_EQ(tiled_param_info_, nullptr); tiled_param_info_ = tiled_param_info; } @@ -89,7 +88,7 @@ class IrEmitterUnnested : public IrEmitter { } llvm::Type* GetIndexType() const { return index_ty_; } - private: + protected: llvm_ir::KernelMappingScheme* mapping_scheme_; llvm_ir::TiledParameterInfo* tiled_param_info_; llvm::Value* lane_id_; @@ -109,10 +108,12 @@ class IrEmitterUnnested : public IrEmitter { // y_loc: The y coordinate within a tile. // x_loc: The x coordinate within a tile. // kernel_info: Other information to support the kernel code generation. + // x_iter_num: When a thread process N elements in the X dimension, x_iter_num + // has a value of 0..N-1 to identify the element being process. using TileElementGenerator = std::function; + llvm::Value* x_loc, int64 x_iter_num)>; // KernelCodeGenerator records the code generator objects that generate code // for tile elements or tile block prologue/epilogue. @@ -175,6 +176,7 @@ class IrEmitterUnnested : public IrEmitter { Status HandleScatter(HloInstruction* scatter) override; Status HandleSelect(HloInstruction* select) override; Status HandleSort(HloInstruction* sort) override; + Status HandleTriangularSolve(HloInstruction* hlo) override; Status HandleTupleSelect(HloInstruction* tuple_select) override; Status HandleAllReduce(HloInstruction* crs) override; Status HandleAfterAll(HloInstruction* after_all) override; @@ -216,9 +218,13 @@ class IrEmitterUnnested : public IrEmitter { Status EmitReductionToVector(HloInstruction* unnested_hlo); // Computes the KernelMappingScheme for the reduce HLO and indicates whether - // the reduction is a row reduction. + // the reduction is a row reduction. For an un-fused reduce op, unnested_hlo + // and first_reduce are the same instruction. For a kInput fusion, + // unnested_hlo is the fusion instruction while first_reduce is the first + // reduce op. std::tuple - ComputeMappingSchemeAndReductionKind(const HloInstruction* first_reduce); + ComputeMappingSchemeAndReductionKind(const HloInstruction* unnested_hlo, + const HloInstruction* first_reduce); // Emits code for an in-place scatter, modifying `thunk`s launch dimensions in // the process. `scatter` may be fused, scatter indices are taken from @@ -243,26 +249,29 @@ class IrEmitterUnnested : public IrEmitter { const KernelCodeGenerator& kernel_generator, KernelCodegenInfo* kernel_info); void EmitBlock(const TileGenerator& emit_one_tile, - const KernelCodegenInfo* kernel_info, - KernelSupportLibrary& ksl, llvm::Type* index_ty); + KernelCodegenInfo* kernel_info, KernelSupportLibrary* ksl, + llvm::Type* index_ty); // Emits code to process a tensor element in a tile for the given kCopy HLO // that performs a 0-2-1 transpose. void EmitTileElementForCopy(HloInstruction* hlo, const llvm_ir::IrArray::Index& index, const KernelCodegenInfo* kernel_info, - llvm::Value* y_loc, llvm::Value* x_loc); + llvm::Value* y_loc, llvm::Value* x_loc, + int64 x_iter_num); // Emits code to process a tensor element in a tile for the given kLoop fusion // HLO containing parameters that are 0-2-1 transpose of its outputs. void EmitTileElementForFusion(HloInstruction* hlo, const llvm_ir::IrArray::Index& index, const KernelCodegenInfo* kernel_info, - llvm::Value* y_loc, llvm::Value* x_loc); + llvm::Value* y_loc, llvm::Value* x_loc, + int64 x_iter_num); // Emits code to process a tensor element in a tile for the given input hlo // that is either a unnested kReduce or a kInput fusion. void EmitTileElementForReduction(HloInstruction* unnested_hlo, const llvm_ir::IrArray::Index& index, const KernelCodegenInfo* kernel_info, - llvm::Value* y_loc, llvm::Value* x_loc); + llvm::Value* y_loc, llvm::Value* x_loc, + int64 x_iter_num); // Prepares for the code generation for a tile block of a reduction kernel. void EmitPrologueForReduction(HloInstruction* unnested_hlo, KernelCodegenInfo* kernel_info); @@ -311,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 TriangularSolveThunk that calls cuBlas to implement `inst`. + std::unique_ptr BuildTriangularSolveThunk(const HloInstruction* inst); + // Returns a GemmThunk that calls gemm to implement `inst`. The caller needs // to make sure `inst` outlives the lifetime of the returned Thunk object. std::unique_ptr BuildGemmThunk(const HloInstruction* inst); diff --git a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/nvptx_backend_lib.cc b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/nvptx_backend_lib.cc index eddaa877f28cb7bd32e924cd179814581eb97b12..153aab97d9eb971734c5ea95564895631bc2a9fa 100644 --- a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/nvptx_backend_lib.cc +++ b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/nvptx_backend_lib.cc @@ -110,11 +110,9 @@ static string GetLibdeviceFilename(const string& libdevice_dir_path, } // Gets the GPU name as it's known to LLVM for a given compute capability. If -// we see an unrecognized compute capability, we return "sm_30". +// we see an unrecognized compute capability, we return "sm_35". static string GetSmName(std::pair compute_capability) { static auto* m = new std::map, int>({ - {{3, 0}, 30}, - {{3, 2}, 32}, {{3, 5}, 35}, {{3, 7}, 37}, {{5, 0}, 50}, @@ -127,7 +125,7 @@ static string GetSmName(std::pair compute_capability) { {{7, 2}, 72}, {{7, 5}, 75}, }); - int sm_version = 30; + int sm_version = 35; auto it = m->find(compute_capability); if (it != m->end()) { sm_version = it->second; diff --git a/tensorflow/compiler/xla/service/gpu/multi_output_fusion.cc b/tensorflow/compiler/xla/service/gpu/multi_output_fusion.cc index 01fddcede64d1bb02ab89db5fc9524893c2d47a4..02e1207f377b8c28bf2566bee8cf3bcbc66794fb 100644 --- a/tensorflow/compiler/xla/service/gpu/multi_output_fusion.cc +++ b/tensorflow/compiler/xla/service/gpu/multi_output_fusion.cc @@ -67,7 +67,7 @@ int64 GpuMultiOutputFusion::GetProfit(HloInstruction* instr1, } int64 profit = 0; for (auto instr : instr2->operands()) { - if (!IsProfitableOperand(instr) || in_list.count(instr) == 0) { + if (!IsProfitableOperand(instr) || !in_list.contains(instr)) { continue; } profit += ShapeUtil::ByteSizeOf(instr->shape()); diff --git a/tensorflow/compiler/xla/service/gpu/multi_output_fusion_test.cc b/tensorflow/compiler/xla/service/gpu/multi_output_fusion_test.cc index d16c87ba5c63aa582753fe949e9e39ee2d8b81e5..40b87b16a195564c9b98497f79a70f1db0539d87 100644 --- a/tensorflow/compiler/xla/service/gpu/multi_output_fusion_test.cc +++ b/tensorflow/compiler/xla/service/gpu/multi_output_fusion_test.cc @@ -628,8 +628,7 @@ TEST_F(MultiOutputFusionTest, MultiOutputFusionDUS) { p.1 = s32[1]{0} parameter(1) p.2 = f16[1,96,1024]{2,1,0} parameter(2) c.0 = s32[] constant(0) - pad = s32[3]{0} pad(p.1, c.0), padding=0_2 - ROOT %dynamic-update-slice = f16[50,96,1024]{2,1,0} dynamic-update-slice(p.0, p.2, pad) + ROOT %dynamic-update-slice = f16[50,96,1024]{2,1,0} dynamic-update-slice(p.0, p.2, p.1, c.0, c.0) } fusion.2 { @@ -638,7 +637,7 @@ TEST_F(MultiOutputFusionTest, MultiOutputFusionDUS) { p.2 = f16[1,96,1024]{2,1,0} parameter(2) c.0 = s32[] constant(0) pad = s32[3]{0} pad(p.1, c.0), padding=0_2 - ROOT %dynamic-update-slice = f16[50,96,1024]{2,1,0} dynamic-update-slice(p.0, p.2, pad) + ROOT %dynamic-update-slice = f16[50,96,1024]{2,1,0} dynamic-update-slice(p.0, p.2, p.1, c.0, c.0) } ENTRY entry { diff --git a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc index cd369d55987b96eed2efb64ae0df6b3a76acb672..9c8a1816040d99bd20af111e8e930149287ed146 100644 --- a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc +++ b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc @@ -37,6 +37,8 @@ limitations under the License. #include "tensorflow/compiler/xla/service/call_inliner.h" #include "tensorflow/compiler/xla/service/conditional_simplifier.h" #include "tensorflow/compiler/xla/service/convolution_group_converter.h" +#include "tensorflow/compiler/xla/service/dot_decomposer.h" +#include "tensorflow/compiler/xla/service/dynamic_index_splitter.h" #include "tensorflow/compiler/xla/service/flatten_call_graph.h" #include "tensorflow/compiler/xla/service/gpu/cudnn_batchnorm_rewriter.h" #include "tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.h" @@ -51,6 +53,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.h" #include "tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker.h" #include "tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.h" +#include "tensorflow/compiler/xla/service/gpu/gpu_sanitize_constant_names.h" #include "tensorflow/compiler/xla/service/gpu/instruction_fusion.h" #include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h" #include "tensorflow/compiler/xla/service/gpu/ir_emitter_context.h" @@ -78,6 +81,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/service/reduce_precision_insertion.h" #include "tensorflow/compiler/xla/service/reshape_mover.h" +#include "tensorflow/compiler/xla/service/sort_simplifier.h" #include "tensorflow/compiler/xla/service/transpose_folding.h" #include "tensorflow/compiler/xla/service/tuple_simplifier.h" #include "tensorflow/compiler/xla/service/while_loop_constant_sinking.h" @@ -114,6 +118,9 @@ std::vector GetCudaRootCandidates( const HloModuleConfig& hlo_module_config) { std::vector potential_cuda_roots = tensorflow::CandidateCudaRoots(); + // "." is our last resort, even though it probably won't work. + potential_cuda_roots.push_back("."); + // CUDA location explicitly specified by user via --xla_gpu_cuda_data_dir has // highest priority. string xla_gpu_cuda_data_dir = @@ -125,9 +132,23 @@ std::vector GetCudaRootCandidates( return potential_cuda_roots; } +void PrintCantFindCudaMessage(absl::string_view msg, + const HloModuleConfig& hlo_module_config) { + LOG(WARNING) << msg; + LOG(WARNING) << "Searched in the following directories:"; + for (const auto& dir : GetCudaRootCandidates(hlo_module_config)) { + LOG(WARNING) << " " << dir; + } + LOG(WARNING) + << "You can choose the search directory by setting xla_gpu_cuda_data_dir " + "in HloModule's DebugOptions. For most apps, setting the environment " + "variable XLA_FLAGS=--xla_gpu_cuda_data_dir=/path/to/cuda will work."; +} + // Returns the directory containing nvvm libdevice files. string GetLibdeviceDir(const HloModuleConfig& hlo_module_config) { - for (const string& cuda_root : GetCudaRootCandidates(hlo_module_config)) { + const auto& candidate_dirs = GetCudaRootCandidates(hlo_module_config); + for (const string& cuda_root : candidate_dirs) { string libdevice_dir = tensorflow::io::JoinPath(cuda_root, "nvvm", "libdevice"); VLOG(2) << "Looking for libdevice at " << libdevice_dir; @@ -136,8 +157,14 @@ string GetLibdeviceDir(const HloModuleConfig& hlo_module_config) { return libdevice_dir; } } - LOG(WARNING) << "Unable to find libdevice dir. Using '.'"; - // Last resort: maybe in the current folder. + PrintCantFindCudaMessage( + "Can't find directory containing CUDA libevice. This may result in " + "compilation or runtime failures, if the program we try to run uses " + "routines from libdevice.", + hlo_module_config); + + // GetCudaRotCandidates always inclues ".", but but if everything fails, we + // return it anyway. Better than returning the empty string. return "."; } @@ -152,6 +179,7 @@ Status OptimizeHloModule(HloModule* hlo_module, se::StreamExecutor* stream_exec, HloPassPipeline pipeline("optimization"); pipeline.AddInvariantChecker(/*layout_sensitive=*/false, /*allow_mixed_precision=*/false); + pipeline.AddPass(); pipeline.AddPass(); ReducePrecisionInsertion::AddPasses( &pipeline, hlo_module->config().debug_options(), @@ -163,6 +191,7 @@ Status OptimizeHloModule(HloModule* hlo_module, se::StreamExecutor* stream_exec, // We need a cost model for GPUs. Currently, do nothing. return false; }; + pipeline.AddPass(false); pipeline.AddPass( cost_model, /*convert_batch_groups_only=*/true); @@ -194,10 +223,9 @@ Status OptimizeHloModule(HloModule* hlo_module, se::StreamExecutor* stream_exec, // elimination has to come after that pass. pipeline.AddPass(); - AlgebraicSimplifierOptions options( - [](const Shape&, const Shape&) { return false; }); - options.set_enable_permutation_sort_replacement(true); + AlgebraicSimplifierOptions options; pass.AddPass(options); + pass.AddPass(); pass.AddPass(); pass.AddPass(); pass.AddPass(); @@ -266,12 +294,8 @@ Status OptimizeHloModule(HloModule* hlo_module, se::StreamExecutor* stream_exec, // The LayoutAssignment pass may leave behind kCopy instructions which are // duplicate or NOPs, so remove them with algebraic simplification and CSE. - AlgebraicSimplifierOptions options( - /*valid_bitcast_callback=*/[](const Shape&, const Shape&) { - return true; - }); + AlgebraicSimplifierOptions options; options.set_is_layout_sensitive(true); - options.set_enable_permutation_sort_replacement(true); pipeline.AddPass>(options); // Choose the fastest algorithm for each conv. @@ -375,6 +399,7 @@ Status PrepareHloModuleForIrEmitting(HloModule* hlo_module) { pipeline.AddPass(); pipeline.AddPass(); pipeline.AddPass(); + pipeline.AddPass(); return pipeline.Run(hlo_module).status(); } @@ -770,14 +795,19 @@ StatusOr> NVPTXCompiler::RunBackend( std::unique_ptr profile_index_map; std::unique_ptr profile_printer; - if (module->config().hlo_profiling_enabled()) { + if (module->config().hlo_profiling_enabled() || VLOG_IS_ON(1)) { HloCostAnalysis cost_analysis(ShapeSizeBytesFunction()); cost_analysis.set_bytes_per_second( stream_exec->GetDeviceDescription().memory_bandwidth()); TF_RETURN_IF_ERROR(module->entry_computation()->Accept(&cost_analysis)); - profile_index_map = absl::make_unique(*module); - profile_printer = CreateHloProfilePrinterData( - *profile_index_map, cost_analysis, entry_computation->name()); + VLOG(1) << "HLO memory read+written: " + << tensorflow::strings::HumanReadableNumBytes( + cost_analysis.bytes_accessed()); + if (module->config().hlo_profiling_enabled()) { + profile_index_map = absl::make_unique(*module); + profile_printer = CreateHloProfilePrinterData( + *profile_index_map, cost_analysis, entry_computation->name()); + } } auto* gpu_executable = new GpuExecutable( @@ -841,10 +871,11 @@ std::vector NVPTXCompiler::CompilePtxOrGetCachedResult( log_warning = !warning_done.exchange(true); } if (log_warning) { - LOG(WARNING) - << "Failed to compile ptx to cubin. Will attempt to let " - "GPU driver compile the ptx. " - << maybe_cubin.status(); + PrintCantFindCudaMessage( + "Can't find ptxas binary. Will back to the GPU driver " + "for PTX -> sass compilation. This is OK so long as you don't " + "see a warning below about an out-of-date driver version.", + hlo_module_config); } // We're going to use the driver to JIT our PTX->SASS, so warn if diff --git a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc index 8154d75d23a6d49153ccb6824402aff73f365617..cb012649200c6386d3ae25d088aa3b16bd40be82 100644 --- a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc @@ -25,7 +25,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" namespace xla { namespace gpu { diff --git a/tensorflow/compiler/xla/service/gpu/stream_assignment.cc b/tensorflow/compiler/xla/service/gpu/stream_assignment.cc index 4775baf44aecfe6adaf2bf0d2791595436635b16..1dedbd3befce6e2ceb06126d83a061207a90dd8f 100644 --- a/tensorflow/compiler/xla/service/gpu/stream_assignment.cc +++ b/tensorflow/compiler/xla/service/gpu/stream_assignment.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/stream_assignment.h" +#include "absl/container/flat_hash_set.h" #include "absl/memory/memory.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h" @@ -25,7 +26,7 @@ namespace xla { namespace gpu { bool StreamAssignment::HasStreamAssigned(const HloInstruction& hlo) const { - return hlo_to_stream_number_.count(&hlo); + return hlo_to_stream_number_.contains(&hlo); } int StreamAssignment::StreamNumberForHlo(const HloInstruction& hlo) const { @@ -98,10 +99,10 @@ int ComputeStreamToAssign( // greedy approach. First, we compute as forbidden_stream_numbers the // streams assigned to GEMMs that are concurrent with `hlo`. Then, we assign // `hlo` a different stream. - std::set forbidden_stream_numbers; + absl::flat_hash_set forbidden_stream_numbers; for (const auto* seen_gemm : seen_gemms) { int stream_num = stream_assignment.StreamNumberForHlo(*seen_gemm); - if (!forbidden_stream_numbers.count(stream_num) && + if (!forbidden_stream_numbers.contains(stream_num) && CanRunConcurrently(*seen_gemm, hlo, reachability)) { forbidden_stream_numbers.insert(stream_num); } @@ -109,7 +110,7 @@ int ComputeStreamToAssign( for (int stream_num = 0; stream_num < stream_assignment.StreamCount(); ++stream_num) { - if (!forbidden_stream_numbers.count(stream_num)) { + if (!forbidden_stream_numbers.contains(stream_num)) { return stream_num; } } diff --git a/tensorflow/compiler/xla/service/gpu/tests/gpu_copy_test.cc b/tensorflow/compiler/xla/service/gpu/tests/gpu_copy_test.cc index a1ed8499040359fe7265a7317b0577a990a2234c..d33e9cf714ee3810b1fb2fa8c05c3ed399d27bfb 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/gpu_copy_test.cc +++ b/tensorflow/compiler/xla/service/gpu/tests/gpu_copy_test.cc @@ -24,7 +24,6 @@ limitations under the License. #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/xla_data.pb.h" #include "tensorflow/core/platform/test.h" namespace xla { diff --git a/tensorflow/compiler/xla/service/gpu/tests/gpu_kernel_tiling_test.cc b/tensorflow/compiler/xla/service/gpu/tests/gpu_kernel_tiling_test.cc index a302b582ede3723acd118d2e4a4bb3efdf7a4d0b..869724db601b2d5e4ed6d3c7bf3e10a748433146 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/gpu_kernel_tiling_test.cc +++ b/tensorflow/compiler/xla/service/gpu/tests/gpu_kernel_tiling_test.cc @@ -65,7 +65,7 @@ TEST_F(GpuKernelTilingTest, UnnestedTransposeWithProperDimensionsTiled) { CompileAndVerifyIr(std::move(hlo_module), R"( ; CHECK-LABEL: define void @copy -; CHECK: tail call void @llvm.nvvm.barrier0() +; CHECK: call void @llvm.nvvm.barrier0() ; CHECK: } )", /*match_optimized_ir=*/true); @@ -91,7 +91,7 @@ TEST_F(GpuKernelTilingTest, UnnestedTransposeWithSmallDimensionsNotTiled) { CompileAndVerifyIr(std::move(hlo_module), R"( ; CHECK-LABEL: define void @copy -; CHECK-NOT: tail call void @llvm.nvvm.barrier0() +; CHECK-NOT: call void @llvm.nvvm.barrier0() ; CHECK: } )", /*match_optimized_ir=*/true); @@ -118,7 +118,7 @@ TEST_F(GpuKernelTilingTest, SimpleFusionWithTransposeTiled) { CompileAndVerifyIr(std::move(hlo_module), R"( ; CHECK-LABEL: define void @fusion -; CHECK: tail call void @llvm.nvvm.barrier0() +; CHECK: call void @llvm.nvvm.barrier0() ; CHECK: } )", /*match_optimized_ir=*/true); @@ -152,7 +152,7 @@ TEST_F(GpuKernelTilingTest, MultipleOutputFusionWithOnePossibleTransposeTiled) { CompileAndVerifyIr(std::move(hlo_module), R"( ; CHECK-LABEL: define void @fusion -; CHECK: tail call void @llvm.nvvm.barrier0() +; CHECK: call void @llvm.nvvm.barrier0() ; CHECK: } )", /*match_optimized_ir=*/true); @@ -187,13 +187,13 @@ TEST_F(GpuKernelTilingTest, CompileAndVerifyIr(std::move(hlo_module), R"( ; CHECK-LABEL: define void @fusion -; CHECK-NOT: tail call void @llvm.nvvm.barrier0() +; CHECK-NOT: call void @llvm.nvvm.barrier0() ; CHECK: } )", /*match_optimized_ir=*/true); } -TEST_F(GpuKernelTilingTest, FusionTransposeWithReverseNotTiled) { +TEST_F(GpuKernelTilingTest, TransposedInputWithUserReverseNotTiled) { const char *const kHloString = R"( HloModule FusionTransposeWithReverseNotTiled fused_computation.1 { @@ -214,12 +214,203 @@ TEST_F(GpuKernelTilingTest, FusionTransposeWithReverseNotTiled) { CompileAndVerifyIr(std::move(hlo_module), R"( ; CHECK-LABEL: define void @fusion -; CHECK-NOT: tail call void @llvm.nvvm.barrier0() +; CHECK-NOT: call void @llvm.nvvm.barrier0() ; CHECK: } )", /*match_optimized_ir=*/true); } +TEST_F(GpuKernelTilingTest, TransposedInputWithUserBitcastNotTiled) { + const char *const kHloString = R"( + HloModule TransposedInputWithUserBitcast + + fused_computation { + param_0 = f32[20,20]{1,0} parameter(0) + ROOT bitcast = f32[20,20]{0,1} bitcast(param_0) + } + + ENTRY kernel_entry { + parameter.0 = f32[20,20]{1,0} parameter(0) + ROOT fusion = f32[20,20]{0,1} fusion(parameter.0), + kind=kLoop, calls=fused_computation + })"; + + // Check that a call to llvm.nvvm.barrier0 is not generated. + auto hlo_module = + ParseHloString(kHloString, ConfigWithoutLayoutAssignment()).ValueOrDie(); + CompileAndVerifyIr(std::move(hlo_module), + R"( +; CHECK-LABEL: define void @fusion +; CHECK-NOT: call void @llvm.nvvm.barrier0() +; CHECK: } +)", + /*match_optimized_ir=*/true); + + // Check that the kernel runs correctly. + EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{0.0})); +} + +TEST_F(GpuKernelTilingTest, TransposedInputWithoutUnsafeUseTiled) { + const char *const kHloString = R"( + HloModule TwoTransposedInputs + + fused_computation { + param_0 = f32[64,64]{1,0} parameter(0) + param_1 = f32[64,64]{1,0} parameter(1) + bitcast = f32[64,64]{0,1} bitcast(param_0) + copy = f32[64,64]{0,1} copy(param_1) + ROOT tuple = (f32[64,64]{0,1}, f32[64,64]{0,1}) tuple(bitcast, copy) + } + + ENTRY kernel_entry { + parameter.0 = f32[64,64]{1,0} parameter(0) + parameter.1 = f32[64,64]{1,0} parameter(1) + ROOT fusion = (f32[64,64]{0,1}, f32[64,64]{0,1}) + fusion(parameter.0, parameter.1), + kind=kLoop, calls=fused_computation + })"; + + // Check that a call to llvm.nvvm.barrier0 is generated. + auto hlo_module = + ParseHloString(kHloString, ConfigWithoutLayoutAssignment()).ValueOrDie(); + CompileAndVerifyIr(std::move(hlo_module), + R"( +; CHECK-LABEL: define void @fusion +; CHECK: call void @llvm.nvvm.barrier0() +; CHECK: } +)", + /*match_optimized_ir=*/true); + // Check that the kernel runs correctly. + EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{0.0})); +} + +TEST_F(GpuKernelTilingTest, ColumnReductionWithPowerOf2OutputElementsUnrolled) { + const char *const kHloString = R"( + HloModule column_reduce_powerof2 + + reduction { + x = f32[] parameter(0) + y = f32[] parameter(1) + ROOT add = f32[] add(x, y) + } + + ENTRY kernel_entry { + constant0 = f32[] constant(0) + arg1 = f16[1024,512]{1,0} parameter(0) + arg1_conv = f32[1024,512]{1,0} convert(arg1) + ROOT reduce = f32[512]{0} reduce(arg1_conv, constant0), dimensions={0}, to_apply=reduction + })"; + + // Check that two calls to llvm.nvvm.atomic are generated. + auto hlo_module = + ParseHloString(kHloString, ConfigWithoutLayoutAssignment()).ValueOrDie(); + CompileAndVerifyIr(std::move(hlo_module), + R"( +; CHECK-LABEL: define void @fusion +; CHECK: call float @llvm.nvvm.atomic.load.add.f32.p0f32 +; CHECK: call float @llvm.nvvm.atomic.load.add.f32.p0f32 +; CHECK-NOT: call float @llvm.nvvm.atomic.load.add.f32.p0f32 +; CHECK: } +)", + /*match_optimized_ir=*/true); + // Check that the kernel runs correctly. + EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{1.0e-5, 1.0e-5})); +} + +TEST_F(GpuKernelTilingTest, + ColumnReductionWithInputLargerThenReduceInputNotUnrolled) { + const char *const kHloString = R"( + HloModule larger_than_reduce_input_parameter + + reduction22 { + x = f32[] parameter(0) + y = f32[] parameter(1) + ROOT add = f32[] add(x, y) + } + + fused_computation { + constant0 = f32[] constant(0) + arg.1 = f16[1024,512]{1,0} parameter(0) + arg.2 = f16[1027,513]{1,0} parameter(1) + arg1.conv = f32[1024,512]{1,0} convert(arg.1) + arg2.conv = f32[1027,513]{1,0} convert(arg.2) + slice2 = f32[1024,512]{1,0} slice(arg2.conv), slice={[2:1026], [1:513]} + add2 = f32[1024,512]{1,0} add(arg1.conv, slice2) + ROOT reduce = f32[512]{0} reduce(add2, constant0), dimensions={0}, + to_apply=reduction22 + } + + ENTRY kernel_entry { + arg1 = f16[1024,512]{1,0} parameter(0) + arg2 = f16[1027,513]{1,0} parameter(1) + ROOT fusion = f32[512]{0} fusion(arg1, arg2), kind=kInput, + calls=fused_computation + })"; + + // Check that one call to llvm.nvvm.atomic is generated. + auto hlo_module = + ParseHloString(kHloString, ConfigWithoutLayoutAssignment()).ValueOrDie(); + CompileAndVerifyIr(std::move(hlo_module), + R"( +; CHECK-LABEL: define void @fusion +; CHECK: call float @llvm.nvvm.atomic.load.add.f32.p0f32 +; CHECK-NOT: call float @llvm.nvvm.atomic.load.add.f32.p0f32 +; CHECK: } +)", + /*match_optimized_ir=*/true); + // Check that the kernel runs correctly. + EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{1.0e-5, 1.0e-5})); +} + +TEST_F(GpuKernelTilingTest, ColumnReductionMOFUnrolled) { + const char *const kHloString = R"( + HloModule column_reduce_powerof2_mof + + reduction22 { + x = f32[] parameter(0) + y = f32[] parameter(1) + ROOT add = f32[] add(x, y) + } + + fused_computation { + constant0 = f32[] constant(0) + arg.1 = f16[1024,512]{1,0} parameter(0) + arg.2 = f16[1024,512]{1,0} parameter(1) + arg1.conv = f32[1024,512]{1,0} convert(arg.1) + arg2.conv = f32[1024,512]{1,0} convert(arg.2) + reduce1 = f32[512]{0} reduce(arg1.conv, constant0), dimensions={0}, + to_apply=reduction22 + reduce2 = f32[512]{0} reduce(arg2.conv, constant0), dimensions={0}, + to_apply=reduction22 + add = f32[1024,512]{1,0} add(arg1.conv, arg2.conv) + ROOT tuple = (f32[512]{0}, f32[512]{0}, f32[1024,512]{1,0}) + tuple(reduce1, reduce2, add) + } + + ENTRY kernel_entry { + arg1 = f16[1024,512]{1,0} parameter(0) + arg2 = f16[1024,512]{1,0} parameter(1) + ROOT fusion = (f32[512]{0}, f32[512]{0}, f32[1024,512]{1,0}) + fusion(arg1, arg2), kind=kInput, calls=fused_computation + })"; + + // Check that four calls to llvm.nvvm.atomic are generated. + auto hlo_module = + ParseHloString(kHloString, ConfigWithoutLayoutAssignment()).ValueOrDie(); + CompileAndVerifyIr(std::move(hlo_module), + R"( +; CHECK-LABEL: define void @fusion +; CHECK: call float @llvm.nvvm.atomic.load.add.f32.p0f32 +; CHECK: call float @llvm.nvvm.atomic.load.add.f32.p0f32 +; CHECK: call float @llvm.nvvm.atomic.load.add.f32.p0f32 +; CHECK: call float @llvm.nvvm.atomic.load.add.f32.p0f32 +; CHECK-NOT: call float @llvm.nvvm.atomic.load.add.f32.p0f32 +; CHECK: } +)", + /*match_optimized_ir=*/true); + // Check that the kernel runs correctly. + EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{1.0e-5, 1.0e-5})); +} } // namespace } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/tests/infeed_test.cc b/tensorflow/compiler/xla/service/gpu/tests/infeed_test.cc index f91a22d482bc8bc046977870a7a4d18ca1acde68..06b06a5b1ee1fb9996be3ebe326893c4160a7e29 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/infeed_test.cc +++ b/tensorflow/compiler/xla/service/gpu/tests/infeed_test.cc @@ -25,7 +25,6 @@ limitations under the License. #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/math/math_util.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/types.h" diff --git a/tensorflow/compiler/xla/service/gpu/thunk.cc b/tensorflow/compiler/xla/service/gpu/thunk.cc index c78605cebbc671272b8df9faf0e0cc54be2f5b1c..a677617727c04811584cbaa295d164ed27273bb2 100644 --- a/tensorflow/compiler/xla/service/gpu/thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/thunk.cc @@ -48,6 +48,8 @@ std::ostream& operator<<(std::ostream& os, Thunk::Kind kind) { return os << "kOutfeed"; case Thunk::kSequential: return os << "kSequential"; + case Thunk::kTriangularSolve: + return os << "kTriangularSolve"; case Thunk::kTuple: return os << "kTuple"; case Thunk::kWhile: diff --git a/tensorflow/compiler/xla/service/gpu/thunk.h b/tensorflow/compiler/xla/service/gpu/thunk.h index e68bee035a029178844282995429eaa960cc4817..bc69af897a01775d2d33d46067464b10e049f3e1 100644 --- a/tensorflow/compiler/xla/service/gpu/thunk.h +++ b/tensorflow/compiler/xla/service/gpu/thunk.h @@ -56,6 +56,7 @@ class Thunk { kMemzero, kOutfeed, kSequential, + kTriangularSolve, kTuple, kWhile, }; diff --git a/tensorflow/compiler/xla/service/gpu/thunk_schedule.cc b/tensorflow/compiler/xla/service/gpu/thunk_schedule.cc index 6b2d76764a077dc6cfa3f9ddc6e525ab330323be..25bad67bab9375559c431466571c62acd0452b01 100644 --- a/tensorflow/compiler/xla/service/gpu/thunk_schedule.cc +++ b/tensorflow/compiler/xla/service/gpu/thunk_schedule.cc @@ -14,17 +14,19 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/service/gpu/thunk_schedule.h" +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/lib/gtl/map_util.h" namespace xla { namespace gpu { void ThunkSchedule::AddDependenciesOnTransitiveOperands( const Thunk& thunk, const HloInstruction& operand, - const std::unordered_map& hlo_to_thunk) { - if (hlo_to_thunk.count(&operand)) { + const absl::flat_hash_map& hlo_to_thunk) { + if (hlo_to_thunk.contains(&operand)) { // If `operand` is mapped to a thunk, adds `operand` to `thunk`'s dependency // list if `operand` is assigned to a different stream. As an optimization, // we skip `operand`'s operands because `operand` depends on them already. @@ -48,14 +50,14 @@ ThunkSchedule::ThunkSchedule( const std::vector& hlo_total_order) : thunks_(std::move(thunks)), stream_assignment_(std::move(stream_assignment)) { - std::unordered_map hlo_to_thunk; + absl::flat_hash_map hlo_to_thunk; for (const auto& thunk : *thunks_) { InsertOrDie(&hlo_to_thunk, thunk->hlo_instruction(), thunk.get()); } for (HloInstruction* hlo : hlo_total_order) { - if (hlo_to_thunk.count(hlo)) { - thunk_total_order_.push_back(FindOrDie(hlo_to_thunk, hlo)); + if (Thunk** thunk = tensorflow::gtl::FindOrNull(hlo_to_thunk, hlo)) { + thunk_total_order_.push_back(*thunk); } } @@ -106,7 +108,7 @@ void ThunkSchedule::RemoveRedundantDependencyEdges() { // redundant dependency edge. Array2D last_dependency(stream_count, stream_count, -1); for (const Thunk* dst : thunk_total_order_) { - if (!depends_on_.count(dst)) { + if (!depends_on_.contains(dst)) { continue; } @@ -134,7 +136,7 @@ void ThunkSchedule::RemoveRedundantDependencyEdges() { const std::list& ThunkSchedule::DependsOn( const Thunk* thunk) const { - if (depends_on_.count(thunk)) { + if (depends_on_.contains(thunk)) { return FindOrDie(depends_on_, thunk); } else { return empty_thunk_list_; diff --git a/tensorflow/compiler/xla/service/gpu/thunk_schedule.h b/tensorflow/compiler/xla/service/gpu/thunk_schedule.h index 43b628a1baf0e79a3197f3cfad3547991642eaed..549378debd52417252724a5d8a6f4d24f2ad0369 100644 --- a/tensorflow/compiler/xla/service/gpu/thunk_schedule.h +++ b/tensorflow/compiler/xla/service/gpu/thunk_schedule.h @@ -21,6 +21,8 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" #include "tensorflow/compiler/xla/service/gpu/stream_assignment.h" #include "tensorflow/compiler/xla/service/gpu/thunk.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -54,7 +56,9 @@ class ThunkSchedule { // Thunks that `thunk` depends on. const std::list& DependsOn(const Thunk* thunk) const; // Whether `thunk` is depended by another thunk. - bool Depended(const Thunk* thunk) const { return depended_by_.count(thunk); } + bool Depended(const Thunk* thunk) const { + return depended_by_.contains(thunk); + } // Delegates to StreamAssignment. int StreamCount() const { return stream_assignment_->StreamCount(); } @@ -75,13 +79,13 @@ class ThunkSchedule { // thunk.hlo_instruction(). void AddDependenciesOnTransitiveOperands( const Thunk& thunk, const HloInstruction& operand, - const std::unordered_map& hlo_to_thunk); + const absl::flat_hash_map& hlo_to_thunk); std::unique_ptr thunks_; std::vector thunk_total_order_; - std::unordered_map> depends_on_; - std::set depended_by_; + absl::flat_hash_map> depends_on_; + absl::flat_hash_set depended_by_; std::list empty_thunk_list_; std::unique_ptr stream_assignment_; diff --git a/tensorflow/compiler/xla/service/gpu/triangular_solve_thunk.cc b/tensorflow/compiler/xla/service/gpu/triangular_solve_thunk.cc new file mode 100644 index 0000000000000000000000000000000000000000..5200a2af412979c7e38d95c5a9bd5bc2ab64f086 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/triangular_solve_thunk.cc @@ -0,0 +1,149 @@ +/* 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/triangular_solve_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 { + +TriangularSolveThunk::TriangularSolveThunk( + const TriangularSolveOptions& options, + const BufferAllocation::Slice& a_buffer, + const BufferAllocation::Slice& b_buffer, PrimitiveType type, + int64 batch_size, int64 m, int64 n, int64 a_batch_stride, + int64 b_batch_stride, const HloInstruction* hlo) + : Thunk(Kind::kTriangularSolve, hlo), + uplo_(options.lower() ? se::blas::UpperLower::kLower + : se::blas::UpperLower::kUpper), + side_(options.left_side() ? se::blas::Side::kLeft + : se::blas::Side::kRight), + unit_diagonal_(options.unit_diagonal() ? se::blas::Diagonal::kUnit + : se::blas::Diagonal::kNonUnit), + a_buffer_(a_buffer), + b_buffer_(b_buffer), + type_(type), + batch_size_(batch_size), + m_(m), + n_(n), + a_batch_stride_(a_batch_stride), + b_batch_stride_(b_batch_stride) { + transpose_a_ = [&] { + switch (options.transpose_a()) { + case TriangularSolveOptions::NO_TRANSPOSE: + return se::blas::Transpose::kNoTranspose; + case TriangularSolveOptions::TRANSPOSE: + return se::blas::Transpose::kTranspose; + case TriangularSolveOptions::ADJOINT: + return se::blas::Transpose::kConjugateTranspose; + default: + LOG(ERROR) << "Invalid triangular solve transpose value " + << options.transpose_a(); + return se::blas::Transpose::kNoTranspose; + } + }(); +} + +Status TriangularSolveThunk::ExecuteOnStream( + const BufferAllocations& buffer_allocations, se::Stream* stream, + HloExecutionProfiler* profiler) { + VLOG(3) << "uplo=" << se::blas::UpperLowerString(uplo_) + << " side=" << se::blas::SideString(side_) + << " diagonal=" << se::blas::DiagonalString(unit_diagonal_) + << " batch_size=" << batch_size_ << " m=" << m_ << " n=" << n_ + << " a_batch_stride=" << a_batch_stride_ + << " b_batch_stride=" << b_batch_stride_; + + const int lda = side_ == se::blas::Side::kLeft ? m_ : n_; + const int ldb = m_; + + char* a_base = static_cast( + buffer_allocations.GetDeviceAddress(a_buffer_).opaque()); + char* b_base = static_cast( + buffer_allocations.GetDeviceAddress(b_buffer_).opaque()); + for (int64 i = 0; i < batch_size_; ++i) { + bool launch_ok; + se::DeviceMemoryBase a_data = + se::DeviceMemoryBase(a_base + i * a_batch_stride_, a_batch_stride_); + se::DeviceMemoryBase b_data = + se::DeviceMemoryBase(b_base + i * b_batch_stride_, b_batch_stride_); + switch (type_) { + case F32: { + se::DeviceMemory b_data_typed(b_data); + launch_ok = stream + ->ThenBlasTrsm(side_, uplo_, transpose_a_, + unit_diagonal_, m_, n_, /*alpha=*/1.0f, + se::DeviceMemory(a_data), lda, + &b_data_typed, ldb) + .ok(); + break; + } + case F64: { + se::DeviceMemory b_data_typed(b_data); + launch_ok = stream + ->ThenBlasTrsm(side_, uplo_, transpose_a_, + unit_diagonal_, m_, n_, /*alpha=*/1.0, + se::DeviceMemory(a_data), lda, + &b_data_typed, ldb) + .ok(); + break; + } + case C64: { + se::DeviceMemory> b_data_typed(b_data); + launch_ok = + stream + ->ThenBlasTrsm(side_, uplo_, transpose_a_, unit_diagonal_, m_, + n_, /*alpha=*/1.0f, + se::DeviceMemory>(a_data), + lda, &b_data_typed, ldb) + .ok(); + break; + } + case C128: { + se::DeviceMemory> b_data_typed(b_data); + launch_ok = + stream + ->ThenBlasTrsm(side_, uplo_, transpose_a_, unit_diagonal_, m_, + n_, /*alpha=*/1.0, + se::DeviceMemory>(a_data), + lda, &b_data_typed, ldb) + .ok(); + break; + } + default: + return InvalidArgument("Invalid type for triangular solve %d", type_); + } + if (!launch_ok) { + return InternalError("Unable to launch triangular solve for thunk %p", + this); + } + } + return Status::OK(); +} + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/triangular_solve_thunk.h b/tensorflow/compiler/xla/service/gpu/triangular_solve_thunk.h new file mode 100644 index 0000000000000000000000000000000000000000..c947162ea32f197f808d099859eadbbc55a65ab1 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/triangular_solve_thunk.h @@ -0,0 +1,75 @@ +/* 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_TRIANGULAR_SOLVE_THUNK_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_TRIANGULAR_SOLVE_THUNK_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/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 triangular +// solve (BlasTrsm). It is generated by IrEmitter. +// +// Thread-compatible. +class TriangularSolveThunk : public Thunk { + public: + TriangularSolveThunk(const TriangularSolveOptions& options, + const BufferAllocation::Slice& a_buffer, + const BufferAllocation::Slice& b_buffer, + PrimitiveType type, int64 batch_size, int64 m, int64 n, + int64 a_batch_stride, int64 b_batch_stride, + const HloInstruction* hlo); + + TriangularSolveThunk(const TriangularSolveThunk&) = delete; + TriangularSolveThunk& operator=(const TriangularSolveThunk&) = delete; + + Status ExecuteOnStream(const BufferAllocations& buffer_allocations, + se::Stream* stream, + HloExecutionProfiler* profiler) override; + + private: + const se::blas::UpperLower uplo_; + const se::blas::Side side_; + const se::blas::Diagonal unit_diagonal_; + se::blas::Transpose transpose_a_; + + const BufferAllocation::Slice a_buffer_; + const BufferAllocation::Slice b_buffer_; + + const PrimitiveType type_; + const int64 batch_size_; + const int64 m_; + const int64 n_; + const int64 a_batch_stride_; + const int64 b_batch_stride_; +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_TRIANGULAR_SOLVE_THUNK_H_ diff --git a/tensorflow/compiler/xla/service/gpu/variadic_op_splitter.cc b/tensorflow/compiler/xla/service/gpu/variadic_op_splitter.cc index c552c2925497f1c4808d74a615d35cdbeeba1858..bbbcc2dbb0f71d08462a1aad6d97e7fd07b2a1fb 100644 --- a/tensorflow/compiler/xla/service/gpu/variadic_op_splitter.cc +++ b/tensorflow/compiler/xla/service/gpu/variadic_op_splitter.cc @@ -23,7 +23,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/errors.h" namespace xla { diff --git a/tensorflow/compiler/xla/service/gpu/xfeed_queue.h b/tensorflow/compiler/xla/service/gpu/xfeed_queue.h index dd46ff433ba0ad6bfa3999b96845fdaebe148aca..167c038420a64d9fa29746ed3fe349620e08e6ff 100644 --- a/tensorflow/compiler/xla/service/gpu/xfeed_queue.h +++ b/tensorflow/compiler/xla/service/gpu/xfeed_queue.h @@ -47,6 +47,10 @@ class XfeedQueue { // Blocks until the queue is non-empty, then returns the buffer at the head of // the queue. BufferType BlockingGetNextDestination() { + for (const auto& callback : before_get_next_dest_callbacks_) { + callback(); + } + bool became_empty; BufferType current_buffer; { @@ -69,6 +73,10 @@ class XfeedQueue { void RegisterOnEmptyCallback(std::function callback) { on_empty_callbacks_.push_back(std::move(callback)); } + void RegisterBeforeGetNextDestinationCallback( + std::function callback) { + before_get_next_dest_callbacks_.push_back(std::move(callback)); + } private: tensorflow::mutex mu_; @@ -82,6 +90,11 @@ class XfeedQueue { // List of callbacks which will be called when 'enqueued_buffers_' becomes // empty. std::vector> on_empty_callbacks_; + + // List of callbacks which will be called before BlockingGetNextDestination() + // is called. This lets you e.g. call EnqueueDestination() for each call to + // BlockingGetNextDestination(). + std::vector> before_get_next_dest_callbacks_; }; } // namespace gpu diff --git a/tensorflow/compiler/xla/service/heap_simulator.cc b/tensorflow/compiler/xla/service/heap_simulator.cc index 9220865867b770eebfb1ada8f31a5d24693a4b8d..4fca981c6a59cdb91a997e6a887fd26472c1a10a 100644 --- a/tensorflow/compiler/xla/service/heap_simulator.cc +++ b/tensorflow/compiler/xla/service/heap_simulator.cc @@ -199,7 +199,7 @@ Status HeapSimulator::RunComputation( // If the buffer has no users and isn't an entry parameter or output, it // must be a dead value. - if (live_buffers.count(buffer) == 0) { + if (!live_buffers.contains(buffer)) { dead_buffers_to_free.push_back(buffer); } } @@ -225,10 +225,10 @@ Status HeapSimulator::RunComputation( } } // Sort to get a deterministic iteration order. - std::sort(operand_buffers_to_free.begin(), operand_buffers_to_free.end(), - [](const BufferValue* x, const BufferValue* y) { - return x->id() < y->id(); - }); + absl::c_sort(operand_buffers_to_free, + [](const BufferValue* x, const BufferValue* y) { + return x->id() < y->id(); + }); // Allocate buffers defined by this instruction. This is the latest point // that we can allocate; right before the buffer is first used. This must @@ -253,7 +253,7 @@ Status HeapSimulator::RunComputation( bool shared = false; if (options_.may_reuse_operand_buffers) { for (const BufferValue* operand_buffer : operand_buffers_to_free) { - if (reused_buffers.count(operand_buffer) != 0) { + if (reused_buffers.contains(operand_buffer)) { continue; } if (buffer->instruction()->IsUserOf(operand_buffer->instruction()) && @@ -335,10 +335,9 @@ Status HeapSimulator::RunComputation( to_free.push_back(buffer); } - std::sort(to_free.begin(), to_free.end(), - [](const BufferValue* x, const BufferValue* y) { - return x->id() < y->id(); - }); + absl::c_sort(to_free, [](const BufferValue* x, const BufferValue* y) { + return x->id() < y->id(); + }); for (const BufferValue* buffer : to_free) { VLOG(3) << "Freeing pending: " << buffer->ToString(); Free(buffer, root); @@ -374,15 +373,15 @@ bool HeapSimulator::IgnoreBuffer(const BufferValue* buffer) const { return true; } return options_.buffers_to_assign != nullptr && - options_.buffers_to_assign->count(buffer) == 0; + !options_.buffers_to_assign->contains(buffer); } // Alloc always calls the underlying heap algorithm. void HeapSimulator::Alloc(const BufferValue* buffer, const HloInstruction* instruction) { - CHECK(allocated_buffers_.count(buffer) == 0) + CHECK(!allocated_buffers_.contains(buffer)) << "Alloc called on allocated buffer: " << *buffer; - CHECK(freed_buffers_.count(buffer) == 0) + CHECK(!freed_buffers_.contains(buffer)) << "Alloc called on freed buffer: " << *buffer; allocated_buffers_.insert(buffer); @@ -411,9 +410,9 @@ void HeapSimulator::Free(const BufferValue* buffer, buffer = group->canonical; } - CHECK(allocated_buffers_.count(buffer) > 0) + CHECK(allocated_buffers_.contains(buffer)) << "Free called on non-allocated buffer: " << *buffer; - CHECK(freed_buffers_.count(buffer) == 0) + CHECK(!freed_buffers_.contains(buffer)) << "Free called on freed buffer: " << *buffer; freed_buffers_.insert(buffer); @@ -433,11 +432,11 @@ void HeapSimulator::ShareBuffer(const BufferValue* buffer, const HloInstruction* instruction) { CHECK_LE(size_fn_(*buffer), size_fn_(*shared)) << "ShareBuffer oversized buffer" << *buffer << " shared: " << *shared; - CHECK(allocated_buffers_.count(buffer) == 0) + CHECK(!allocated_buffers_.contains(buffer)) << "ShareBuffer called on allocated buffer: " << *buffer; - CHECK(freed_buffers_.count(buffer) == 0) + CHECK(!freed_buffers_.contains(buffer)) << "ShareBuffer called on freed buffer: " << *buffer; - CHECK(freed_buffers_.count(shared) == 0) + CHECK(!freed_buffers_.contains(shared)) << "ShareBuffer called on freed shared buffer: " << *shared; const BufferValue* canonical = nullptr; @@ -452,7 +451,7 @@ void HeapSimulator::ShareBuffer(const BufferValue* buffer, } else { // The 'shared' buffer doesn't have a group; it must be the canonical. Add // both 'buffer' and 'shared' to a new group. - CHECK(allocated_buffers_.count(shared) > 0) + CHECK(allocated_buffers_.contains(shared)) << "ShareBuffer called on non-allocated shared buffer: " << *shared; auto group = std::make_shared(); canonical = shared; @@ -596,7 +595,7 @@ void DecreasingSizeRunsHeap::CallAndDrainRun() { } // Call ops in the run sorted by decreasing size, breaking ties by buffer id. - std::sort(run_.begin(), run_.end(), [](const Op& a, const Op& b) { + absl::c_sort(run_, [](const Op& a, const Op& b) { if (a.size != b.size) { return a.size > b.size; } @@ -866,23 +865,23 @@ HeapSimulator::Result GlobalDecreasingSizeBestFitHeap::Finish() { for (auto& entry : buffer_intervals_) { sorted_buffer_intervals.push_back(entry.second); } - std::sort(sorted_buffer_intervals.begin(), sorted_buffer_intervals.end(), - [](const BufferInterval& x, const BufferInterval& y) { - if (x.size != y.size) { - return x.size > y.size; - } - if (x.end - x.start != y.end - y.start) { - return x.end - x.start > y.end - y.start; - } - return x.buffer->id() < y.buffer->id(); - }); + absl::c_sort(sorted_buffer_intervals, + [](const BufferInterval& x, const BufferInterval& y) { + if (x.size != y.size) { + return x.size > y.size; + } + if (x.end - x.start != y.end - y.start) { + return x.end - x.start > y.end - y.start; + } + return x.buffer->id() < y.buffer->id(); + }); BufferIntervalTree interval_tree(sorted_buffer_intervals.size()); for (auto& buffer_interval : sorted_buffer_intervals) { auto chunks_overlapping_in_time = interval_tree.ChunksOverlappingInTime( buffer_interval.start, buffer_interval.end); - std::sort( - chunks_overlapping_in_time.begin(), chunks_overlapping_in_time.end(), + absl::c_sort( + chunks_overlapping_in_time, [](const Chunk& x, const Chunk& y) { return x.offset < y.offset; }); // Find the minimum free chunk that can hold this buffer. diff --git a/tensorflow/compiler/xla/service/heap_simulator.h b/tensorflow/compiler/xla/service/heap_simulator.h index dbbf43082f2c1d21f5ef42f53804bf0969903a58..3e0631aeb4aa374cb5748650e1c7529e26e10b34 100644 --- a/tensorflow/compiler/xla/service/heap_simulator.h +++ b/tensorflow/compiler/xla/service/heap_simulator.h @@ -158,7 +158,7 @@ class HeapSimulator { void FillDebugTrace(HeapSimulatorTrace::Event::Kind kind, const BufferValue* buffer, const HloInstruction* instruction, - const BufferValue* shared_with_canonical); + const BufferValue* share_with_canonical); // Counterintuitive: the algorithm_ itself can be a NoFragmentationStatsHeap, // in which case we are calculating the same allocs/frees twice in the diff --git a/tensorflow/compiler/xla/service/hlo.proto b/tensorflow/compiler/xla/service/hlo.proto index 9b50f1ca5b5365463f32106fc005ef2c63f2e37a..6e64549e7e1cad80b740452816028c730053623f 100644 --- a/tensorflow/compiler/xla/service/hlo.proto +++ b/tensorflow/compiler/xla/service/hlo.proto @@ -34,7 +34,7 @@ import "tensorflow/compiler/xla/xla_data.proto"; option cc_enable_arenas = true; // Serialization of HloInstruction. -// Next ID: 59 +// Next ID: 60 message HloInstructionProto { reserved 10; reserved "parameter_name"; @@ -193,6 +193,9 @@ message HloInstructionProto { // operand. bool constrain_layout = 56; repeated xla.ShapeProto operand_shapes_with_layout = 57; + + // Options for TriangularSolve + xla.TriangularSolveOptions triangular_solve_options = 59; } // Serialization of HloComputation. @@ -229,6 +232,18 @@ message HloScheduleProto { } message HloInputOutputAliasProto { + enum Kind { + // Define a UNDEFINED_ALIAS equal to zero to get around the default-0 proto3 + // behavior and missing has_*() APIs. + UNDEFINED_ALIAS = 0; + // An alias setup by the user as must alias. A use setting USER_ALIAS is + // expecting the designed output to be dropped over the given input + // parameter number+index. + USER_ALIAS = 1; + // An alias setup by the compiler as part of its optimizations. + SYSTEM_ALIAS = 2; + } + // The following proto describes a pair of aliased an input // (described by parameter number and a ShapeIndex of the parameter) // and an output (described by a ShapeIndex of the root @@ -249,6 +264,8 @@ message HloInputOutputAliasProto { int64 parameter_number = 2; // ShapeIndex of the parameter instruction. repeated int64 parameter_shape_index = 3; + // The kind of alias to be setup. + Kind kind = 4; } repeated AliasEntryProto entries = 1; diff --git a/tensorflow/compiler/xla/service/hlo_alias_analysis.cc b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc index 68094d8907f30202533672376b2199e7b77dc806..e511f1951c5dd07ebb64fa38fd5b7f6a0e87b429 100644 --- a/tensorflow/compiler/xla/service/hlo_alias_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc @@ -117,7 +117,7 @@ class BufferValueMap { for (const auto& pair : buffers_) { buffer_numbers.push_back(pair.first); } - std::sort(buffer_numbers.begin(), buffer_numbers.end()); + absl::c_sort(buffer_numbers); return buffer_numbers; } @@ -176,13 +176,12 @@ class BufferValueMap { const HloValue& value, std::vector* aliased_buffers) { // Get parameter value from an aliased_input object. const auto get_parameter_value = - [this](const std::pair& aliased_input) + [this](const HloInputOutputAliasConfig::Alias& aliased_input) -> const HloValue& { - int64 param_number = aliased_input.first; - const ShapeIndex& param_index = aliased_input.second; return dataflow_.GetUniqueValueAt( - module_->entry_computation()->parameter_instruction(param_number), - param_index); + module_->entry_computation()->parameter_instruction( + aliased_input.parameter_number), + aliased_input.parameter_index); }; // If the value shows up in a root instruction, alias it with parameter @@ -319,7 +318,7 @@ class BufferValueMap { ComputeWhileAliasedBuffers(value, &aliased_buffers); ComputeConditionalAliasedBuffers(value, &aliased_buffers); // Uniquify aliased buffers. - std::sort(aliased_buffers.begin(), aliased_buffers.end()); + absl::c_sort(aliased_buffers); aliased_buffers.erase( std::unique(aliased_buffers.begin(), aliased_buffers.end()), aliased_buffers.end()); @@ -367,7 +366,7 @@ std::vector HloAliasAnalysis::ComputeBuffersAt( } // Sort and uniquify vector before returning. - std::sort(buffers.begin(), buffers.end(), HloBuffer::IdLessThan); + absl::c_sort(buffers, HloBuffer::IdLessThan); buffers.erase(std::unique(buffers.begin(), buffers.end()), buffers.end()); return buffers; @@ -430,8 +429,7 @@ Status HloAliasAnalysis::Verify() const { for (const auto& pair : value_to_buffer_) { const HloValue* value = pair.first; const HloBuffer& buffer = *pair.second; - TF_RET_CHECK(std::find(buffer.values().begin(), buffer.values().end(), - value) != buffer.values().end()); + TF_RET_CHECK(absl::c_linear_search(buffer.values(), value)); } for (HloBuffer::Id id = 0; id < buffers_.size(); ++id) { @@ -515,7 +513,7 @@ StatusOr> HloAliasAnalysis::Run( auto& value_set = buffer_map.GetValuesInBuffer(buffer_number); std::vector sorted_values(value_set.begin(), value_set.end()); - std::sort(sorted_values.begin(), sorted_values.end(), HloValue::IdLessThan); + absl::c_sort(sorted_values, HloValue::IdLessThan); alias_analysis->buffers_.emplace_back(next_id++, sorted_values); for (const HloValue* value : sorted_values) { alias_analysis->value_to_buffer_[value] = @@ -547,16 +545,15 @@ bool HloAliasAnalysis::HasLiveRangeInterference( // tie-break using value ID. The tie-break is necessary because we need a // strict weak order for std::sort. std::vector values = buffer.values(); - std::sort(values.begin(), values.end(), - [&ordering](const HloValue* a, const HloValue* b) { - if (ordering.IsDefinedBefore(*a, *b)) { - return true; - } else if (ordering.IsDefinedBefore(*b, *a)) { - return false; - } else { - return a->id() < b->id(); - } - }); + absl::c_sort(values, [&ordering](const HloValue* a, const HloValue* b) { + if (ordering.IsDefinedBefore(*a, *b)) { + return true; + } else if (ordering.IsDefinedBefore(*b, *a)) { + return false; + } else { + return a->id() < b->id(); + } + }); // Walk through the ordered vector of values. First verify that the values // are totally ordered with respect to 'ordering', then check that no diff --git a/tensorflow/compiler/xla/service/hlo_alias_analysis_test.cc b/tensorflow/compiler/xla/service/hlo_alias_analysis_test.cc index 7e6150e94153cd15463725e862ce1b8593f2c991..b6dbf07959c541bceaa8eda5a0101503970ee832 100644 --- a/tensorflow/compiler/xla/service/hlo_alias_analysis_test.cc +++ b/tensorflow/compiler/xla/service/hlo_alias_analysis_test.cc @@ -238,13 +238,16 @@ TEST_F(HloAliasAnalysisTest, ParametersWithAliasing) { builder.AddInstruction(HloInstruction::CreateTuple({negate0, negate1})); module_->AddEntryComputation(builder.Build()); TF_ASSERT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); + /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); TF_ASSERT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1})); + /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); // Cannot alias an output twice. ASSERT_IS_NOT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{0})); + /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{0}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); const HloAliasAnalysis& analysis = RunAnalysis(); @@ -279,13 +282,16 @@ TEST_F(HloAliasAnalysisTest, ParametersWithCrossAliasing) { builder.AddInstruction(HloInstruction::CreateTuple({gte0, gte1})); module_->AddEntryComputation(builder.Build()); TF_ASSERT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{1})); + /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{1}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); TF_ASSERT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{0})); + /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{0}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); // Cannot alias an output twice. ASSERT_IS_NOT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1})); + /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); const HloAliasAnalysis& analysis = RunAnalysis(); @@ -365,9 +371,11 @@ TEST_F(HloAliasAnalysisTest, InputOutputAliasingWithWhile) { builder.AddInstruction(HloInstruction::CreateTuple({negate_1, negate_2})); module_->AddEntryComputation(builder.Build()); TF_ASSERT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0})); + /*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); TF_ASSERT_OK(module_->input_output_alias_config().SetUpAlias( - /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1})); + /*output_index=*/{1}, /*param_number=*/0, /*param_index=*/{1}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); const HloAliasAnalysis& analysis = RunAnalysis(); diff --git a/tensorflow/compiler/xla/service/hlo_buffer.cc b/tensorflow/compiler/xla/service/hlo_buffer.cc index 9c3aa0e64d119c2560f4955d0bcb492519fa52a2..32e48651b30bace4723169935d1f10dd7d7bfec3 100644 --- a/tensorflow/compiler/xla/service/hlo_buffer.cc +++ b/tensorflow/compiler/xla/service/hlo_buffer.cc @@ -49,7 +49,7 @@ std::vector HloBuffer::ComputePositions() const { value->positions().end()); } // Remove duplicates and sort positions. - std::sort(positions.begin(), positions.end()); + absl::c_sort(positions); positions.erase(std::unique(positions.begin(), positions.end()), positions.end()); return positions; diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc index 52ca67afb8eb270c490566ed51514b0b0f499b42..40fe91398be33f5681e1389e1b6fadcbd87487bb 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.cc +++ b/tensorflow/compiler/xla/service/hlo_computation.cc @@ -207,14 +207,14 @@ Status HloComputation::RemoveInstructionAndUnusedOperands( TF_RET_CHECK(instruction->user_count() == 0); TF_RET_CHECK(IsRemovable(instruction)) << "Cannot remove instruction: " << instruction->ToString(); - std::unordered_set removed; + absl::flat_hash_set removed; std::queue worklist; worklist.push(instruction); while (!worklist.empty()) { HloInstruction* item = worklist.front(); worklist.pop(); - if (removed.count(item) != 0 || item->user_count() != 0 || + if (removed.contains(item) || item->user_count() != 0 || item == root_instruction() || !IsRemovable(item) || (item->HasSideEffect() && item != instruction)) { continue; @@ -531,11 +531,10 @@ HloComputation::CreateFromProto( HloInstruction* root = instruction_map.at(proto.root_id()); // Sort the instructions in the proto id's order. - std::sort(instructions.begin(), instructions.end(), - [&](const std::unique_ptr& a, - const std::unique_ptr& b) { - return to_proto_id[a.get()] < to_proto_id[b.get()]; - }); + absl::c_sort(instructions, [&](const std::unique_ptr& a, + const std::unique_ptr& b) { + return to_proto_id[a.get()] < to_proto_id[b.get()]; + }); TF_RETURN_IF_ERROR([&]() -> Status { std::vector parameters_seen(parameter_count); @@ -694,22 +693,36 @@ bool HloComputation::operator==(const HloComputation& other) const { if (this == &other) { return true; } - std::set> visited; - std::function eq = - [&visited, &eq](const HloInstruction* a, const HloInstruction* b) { - // If are visited but not identical, the recursion should have - // been aborted. So, if are visited at this point, they must be - // identical. - if (visited.count(std::make_pair(a, b)) > 0) { - return true; - } - visited.emplace(a, b); - return a->Identical( - *b, eq, [](const HloComputation* a, const HloComputation* b) { - return *a == *b; - }); - }; - return eq(root_instruction(), other.root_instruction()); + absl::flat_hash_set> + visited; + std::vector> worklist; + + worklist.push_back({root_instruction(), other.root_instruction()}); + + while (!worklist.empty()) { + auto pair = worklist.back(); + worklist.pop_back(); + + if (visited.contains(pair)) { + continue; + } + visited.emplace(pair); + // TODO(b/123082518): Avoid recursively invoking == becasue it may + // cause a stack overflow with deeply nested subcomputations. + bool identical_ignoring_operands = pair.first->Identical( + *pair.second, + [](const HloInstruction*, const HloInstruction*) { return true; }, + [](const HloComputation* a, const HloComputation* b) { + return *a == *b; + }); + if (!identical_ignoring_operands) { + return false; + } + for (size_t i = 0; i < pair.first->operands().size(); ++i) { + worklist.push_back({pair.first->operand(i), pair.second->operand(i)}); + } + } + return true; } Status HloComputation::ReplaceWithNewInstruction( @@ -799,17 +812,16 @@ Status HloComputation::AcceptOrdered( absl::Span order) const { VLOG(3) << "Accepting visitor with order."; for (HloInstruction* root : CollectUnreachableRoots()) { - TF_RET_CHECK(std::find(order.begin(), order.end(), root) != order.end()) - << root->ToString(); + TF_RET_CHECK(absl::c_linear_search(order, root)) << root->ToString(); } TF_RET_CHECK(order.size() == instruction_count()); - std::unordered_set visited; + absl::flat_hash_set visited; for (const HloInstruction* instruction : order) { VLOG(3) << "Visiting ordered: " << instruction->ToString(); - TF_RET_CHECK(instruction_iterators_.count(instruction) == 1) + TF_RET_CHECK(instruction_iterators_.contains(instruction)) << "Instruction " << instruction->name() << " is not in computation " << name(); - TF_RET_CHECK(visited.count(instruction) == 0) + TF_RET_CHECK(!visited.contains(instruction)) << "Instruction " << instruction->name() << " appears more than once in order"; HloInstruction* mutable_instruction = @@ -845,29 +857,31 @@ Status HloComputation::Accept( std::unique_ptr HloComputation::Clone( const string& suffix, HloCloneContext* context) { return CloneWithReplacements( - /*replacements=*/std::unordered_map>(), - context, suffix); + /*replacements=*/absl::flat_hash_map>(), + /*extra_parameters=*/{}, context, suffix); } std::unique_ptr HloComputation::CloneWithReplacementPairs( std::pair> r1, HloCloneContext* context, const string& suffix) { - std::unordered_map> + absl::flat_hash_map> replacements; replacements.emplace(std::move(r1)); - return CloneWithReplacements(std::move(replacements), context, suffix); + return CloneWithReplacements(std::move(replacements), /*extra_parameters=*/{}, + context, suffix); } std::unique_ptr HloComputation::CloneWithReplacementPairs( std::pair> r1, std::pair> r2, HloCloneContext* context, const string& suffix) { - std::unordered_map> + absl::flat_hash_map> replacements; replacements.emplace(std::move(r1)); replacements.emplace(std::move(r2)); - return CloneWithReplacements(std::move(replacements), context, suffix); + return CloneWithReplacements(std::move(replacements), /*extra_parameters=*/{}, + context, suffix); } std::unique_ptr HloComputation::CloneWithReplacementPairs( @@ -875,17 +889,19 @@ std::unique_ptr HloComputation::CloneWithReplacementPairs( std::pair> r2, std::pair> r3, HloCloneContext* context, const string& suffix) { - std::unordered_map> + absl::flat_hash_map> replacements; replacements.emplace(std::move(r1)); replacements.emplace(std::move(r2)); replacements.emplace(std::move(r3)); - return CloneWithReplacements(std::move(replacements), context, suffix); + return CloneWithReplacements(std::move(replacements), /*extra_parameters=*/{}, + context, suffix); } std::unique_ptr HloComputation::CloneWithReplacements( - std::unordered_map> + absl::flat_hash_map> replacements, + absl::Span extra_parameters, HloCloneContext* context, const string& suffix) { std::unique_ptr context_ptr; if (context == nullptr) { @@ -951,6 +967,12 @@ std::unique_ptr HloComputation::CloneWithReplacements( } std::vector> instructions; + // First add the extra parameters to 'instructions'. + for (const auto& instr : extra_parameters) { + CHECK_EQ(instr->opcode(), HloOpcode::kParameter) + << "Only parameter instructions are allowed in 'extra_parameters'"; + instructions.emplace_back(instr->Clone()); + } for (auto instr : postorder) { std::vector new_operands; for (auto operand : instr->operands()) { diff --git a/tensorflow/compiler/xla/service/hlo_computation.h b/tensorflow/compiler/xla/service/hlo_computation.h index a0ccbc583f8c409f29d31756fcc1fa1b4af7dc35..fd1f990431a87ef27d3d7b0ae56ba73c444bc1cc 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.h +++ b/tensorflow/compiler/xla/service/hlo_computation.h @@ -20,7 +20,6 @@ limitations under the License. #include #include #include -#include #include #include #include @@ -323,11 +322,16 @@ class HloComputation { // that's not already in the computation, it's cloned and added to the new // computation. // + // 'extra_parameters' allows to specify additional parameters that should be + // added to the computation. + // // All relevant instructions are cloned, *including* unique_ptr in the // `replacements` map. std::unique_ptr CloneWithReplacements( - std::unordered_map> + absl::flat_hash_map> replacements, + absl::Span extra_parameters = {}, HloCloneContext* context = nullptr, const string& suffix = "clone"); // Convenience overloads for CloneWithReplacements. You want to do @@ -387,6 +391,10 @@ class HloComputation { fusion_instruction_ = fusion_instruction; } + // Clear the unique ID of the computation so that it can be re-assigned, such + // as for the purpose of compacting the unique IDs. + void ClearUniqueIdInternal() { unique_id_ = -1; } + // The id of this computation should be unique within the module. void SetUniqueId(int64 id) { CHECK_EQ(unique_id_, -1); diff --git a/tensorflow/compiler/xla/service/hlo_computation_test.cc b/tensorflow/compiler/xla/service/hlo_computation_test.cc index 0361c87428f6e4c031d95492a5bc782ad388e5b5..3b88e9745c27d6e1f2a46e5c83ac2e8bd8d05150 100644 --- a/tensorflow/compiler/xla/service/hlo_computation_test.cc +++ b/tensorflow/compiler/xla/service/hlo_computation_test.cc @@ -15,8 +15,12 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_computation.h" +#include #include +#include +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -226,7 +230,7 @@ TEST_F(HloComputationTest, VisitWithMultipleRoots) { : computation_(computation) {} Status DefaultAction(HloInstruction* hlo_instruction) override { - EXPECT_EQ(0, visited_set_.count(hlo_instruction)); + EXPECT_FALSE(visited_set_.contains(hlo_instruction)); visited_set_.insert(hlo_instruction); last_visited_ = hlo_instruction; return Status::OK(); @@ -239,7 +243,7 @@ TEST_F(HloComputationTest, VisitWithMultipleRoots) { } HloComputation* computation_; - std::set visited_set_; + absl::flat_hash_set visited_set_; int64 finish_visit_calls_ = 0; HloInstruction* last_visited_ = nullptr; }; @@ -491,6 +495,41 @@ TEST_F(HloComputationTest, CloneWithControlDependency) { EXPECT_THAT(successors, ::testing::ElementsAre(cloned_add)); } +TEST_F(HloComputationTest, CloneWithReplacements) { + auto builder = HloComputation::Builder(TestName()); + Shape r0s64 = ShapeUtil::MakeShape(S64, {}); + Shape r0s32 = ShapeUtil::MakeShape(S32, {}); + Shape r0u32 = ShapeUtil::MakeShape(U32, {}); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32_, "p.0.lhs")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32_, "p.0.rhs")); + auto param2 = + builder.AddInstruction(HloInstruction::CreateParameter(2, r0s64, "p.1")); + auto lt = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(PRED, {}), HloOpcode::kLt, param0, param1)); + auto module = CreateNewVerifiedModule(); + auto computation = + module->AddEntryComputation(builder.Build(/*root_instruction=*/lt)); + absl::flat_hash_map> + replacements; + replacements.emplace(param2, + HloInstruction::CreateParameter(2, r0s32, "p.1")); + auto param3 = HloInstruction::CreateParameter(3, r0u32, "p.2"); + std::vector extra_parameters{param3.get()}; + auto clone = computation->CloneWithReplacements(std::move(replacements), + extra_parameters); + ASSERT_EQ(clone->num_parameters(), 4); + EXPECT_TRUE( + ShapeUtil::Equal(clone->parameter_instruction(0)->shape(), r0f32_)); + EXPECT_TRUE( + ShapeUtil::Equal(clone->parameter_instruction(1)->shape(), r0f32_)); + EXPECT_TRUE( + ShapeUtil::Equal(clone->parameter_instruction(2)->shape(), r0s32)); + EXPECT_TRUE( + ShapeUtil::Equal(clone->parameter_instruction(3)->shape(), r0u32)); +} + TEST_F(HloComputationTest, Stringification) { const Shape s1 = ShapeUtil::MakeShape(F32, {5, 10}); const Shape s2 = ShapeUtil::MakeShape(F32, {20, 10}); @@ -606,5 +645,28 @@ TEST_F(HloComputationTest, StringificationCanonical) { EXPECT_EQ(computation->ToString(options), expected_computation2); } +std::unique_ptr MakeAddNComputation(int n) { + auto builder = HloComputation::Builder("add_n"); + auto result = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {}), "x_value")); + auto one = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + for (int i = 0; i < n; ++i) { + result = builder.AddInstruction(HloInstruction::CreateBinary( + one->shape(), HloOpcode::kAdd, result, one)); + } + return builder.Build(); +} + +TEST_F(HloComputationTest, DeepEquality) { + auto computation_a = MakeAddNComputation(200000); + auto computation_b = MakeAddNComputation(200000); + EXPECT_TRUE(*computation_a == *computation_b); + + auto computation_c = MakeAddNComputation(199999); + EXPECT_FALSE(*computation_a == *computation_c); + EXPECT_FALSE(*computation_c == *computation_b); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_constant_folding.cc b/tensorflow/compiler/xla/service/hlo_constant_folding.cc index c58d00ff5084d060498f2d6fbbfa8e12207b810e..e7ed858e8c5af83d08863d64a0aba162c75ed5cb 100644 --- a/tensorflow/compiler/xla/service/hlo_constant_folding.cc +++ b/tensorflow/compiler/xla/service/hlo_constant_folding.cc @@ -35,6 +35,34 @@ limitations under the License. namespace xla { +// Checks whether instr is or transitively contains an instruction that we +// shouldn't fold. +// +// Specifically, we don't fold kRng or kAfterAll instructions: +// +// - kRng is already marked as side-effecting and so is skipped elsewhere, but +// we check for it here. Even kRng weren't side-effecting and took an +// explicit seed, we *still* wouldn't want to constant-fold it, because the +// evaluator's handling of rng is not guaranteed to be identical to any +// particular backend's rng. +// +// - kAfterAll needs to be skipped because a kAfterAll op with no args can +// currently materialize a token "out of thin air". TODO(b/110532604): +// Remove this check once AfterAll requires at least one operand, in which +// case constant folding will be impossible. +static bool IsOrContainsIllegalInstr(const HloInstruction* instr) { + if (instr->opcode() == HloOpcode::kAfterAll || + instr->opcode() == HloOpcode::kRng) { + return true; + } + for (const HloComputation* c : instr->called_computations()) { + if (absl::c_any_of(c->instructions(), IsOrContainsIllegalInstr)) { + return true; + } + } + return false; +} + StatusOr HloConstantFolding::Run(HloModule* module) { // Limit the constant folding to 0 iterations to skip folding loops. This // retains the behavior from before while loop support in HloEvaluator and may @@ -52,25 +80,24 @@ StatusOr HloConstantFolding::Run(HloModule* module) { computation->root_instruction() != instruction) { continue; } - // Skip Constant, Parameter, Tuple, AfterAll operation. - // Tuple constants are not directly supported by any backends, hence - // folding Tuple is not useful and would in fact be expanded back into - // kTuple by Algebraic Simplifier. - // TODO(b/110532604): Enable AfterAll once AfterAll requires at least one - // operand in which case constant folding will be impossible and this - // special case is not necessary. - if (instruction->opcode() == HloOpcode::kParameter || - instruction->opcode() == HloOpcode::kConstant || - instruction->opcode() == HloOpcode::kTuple || - instruction->opcode() == HloOpcode::kAfterAll) { - continue; - } // Skip instructions with non-constant operands. if (!hlo_query::AllOperandsAreConstants(*instruction)) { continue; } + // Don't fold Constant, Parameter, and Tuple instructions. Tuple + // constants are not directly supported by any backends, hence folding + // Tuple is not useful and would in fact be expanded back into kTuple by + // Algebraic Simplifier. + // + // (We do allow folding subcomputations that contain these instructions.) + if (instruction->opcode() == HloOpcode::kParameter || + instruction->opcode() == HloOpcode::kConstant || + instruction->opcode() == HloOpcode::kTuple) { + continue; + } + // Broadcasts dramatically increase the size of constants, which is often // detrimental to performance and memory capacity, so do not fold // broadcasts. @@ -79,6 +106,18 @@ StatusOr HloConstantFolding::Run(HloModule* module) { continue; } + // Check for instructions that we can't fold even if they appear inside of + // a subcomputation (e.g. a kCall). + if (IsOrContainsIllegalInstr(instruction)) { + continue; + } + + // Don't constant-fold side-effecting instructions or instructions which + // contain side-effecting instructions. + if (instruction->HasSideEffect()) { + continue; + } + // Don't constant fold unless it's a net positive or the output is small. if (instruction->shape().IsArray()) { int64 elements_in_removed_operands = 0; diff --git a/tensorflow/compiler/xla/service/hlo_constant_folding_test.cc b/tensorflow/compiler/xla/service/hlo_constant_folding_test.cc index 92b748d813c3efef83ef0155f1d5d3c637ce2c57..4bdc980c9ac4fb79cde0242f407aea7057474b27 100644 --- a/tensorflow/compiler/xla/service/hlo_constant_folding_test.cc +++ b/tensorflow/compiler/xla/service/hlo_constant_folding_test.cc @@ -268,5 +268,51 @@ TEST_F(HloConstantFoldingTest, DoesNotFoldLargePad) { GmockMatch(m::Pad(m::Constant(), m::Constant()))); } +TEST_F(HloConstantFoldingTest, DontFoldSubcomputationContainingAfterAll) { + const char* const kModuleStr = R"( + HloModule test + + Fn { + tok = token[] after-all() + ROOT root = f32[10] iota(), iota_dimension=0 + } + + ENTRY entry { + ROOT call = f32[10] call(), to_apply=Fn + })"; + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kModuleStr)); + HloConstantFolding constant_folding; + TF_ASSERT_OK_AND_ASSIGN(bool result, + RunHloPass(&constant_folding, module.get())); + EXPECT_FALSE(result); +} + +TEST_F(HloConstantFoldingTest, + DontFoldSubcomputationTransitivelyContainingRng) { + const char* const kModuleStr = R"( + HloModule test + + InnerFn { + c0 = f32[] constant(0) + c1 = f32[] constant(1) + ROOT rng = f32[10] rng(c0, c1), distribution=rng_uniform + } + + Fn { + ROOT fusion = f32[10] fusion(), kind=kLoop, calls=InnerFn + } + + ENTRY entry { + ROOT call = f32[10] call(), to_apply=Fn + })"; + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kModuleStr)); + HloConstantFolding constant_folding; + TF_ASSERT_OK_AND_ASSIGN(bool result, + RunHloPass(&constant_folding, module.get())); + EXPECT_FALSE(result); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc index c4b4fa62ddcb46b8ac46567da5ab32a6a1f4914c..29ac263c5f39b5ec1f02a232704adcd3e3f21f60 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc @@ -237,24 +237,17 @@ Status HloCostAnalysis::HandleDomain(const HloInstruction* domain) { Status HloCostAnalysis::HandleDot(const HloInstruction* dot) { const Shape& lhs_shape = dot->operand(0)->shape(); - const Shape& rhs_shape = dot->operand(1)->shape(); + const Shape& dot_shape = dot->shape(); const DotDimensionNumbers& dnums = dot->dot_dimension_numbers(); // Count of elements along the reduction dimension (last dimension for the // rhs). - int64 reduction_width = - lhs_shape.dimensions(dnums.lhs_contracting_dimensions(0)); - // First divide by reduction width before multiplying by rhs elements to avoid - // overflow. - int64 fma_count; - if (reduction_width == 0) { - fma_count = 0; - } else { - fma_count = (ShapeUtil::ElementsIn(lhs_shape) / reduction_width) * - ShapeUtil::ElementsIn(rhs_shape); + int64 reduction_width = 1; + for (auto dim : dnums.lhs_contracting_dimensions()) { + reduction_width *= lhs_shape.dimensions(dim); } - - // We count an FMA operation as 2 floating point operations. - current_properties_[kFlopsKey] = kFmaFlops * fma_count; + // Each output elment requires reduction_width FMA operations. + current_properties_[kFlopsKey] = + kFmaFlops * ShapeUtil::ElementsIn(dot_shape) * reduction_width; return Status::OK(); } @@ -531,7 +524,8 @@ Status HloCostAnalysis::HandleConvolution(const HloInstruction* convolution) { } const int64 fma_count = (input_feature / convolution->feature_group_count()) * - output_feature * batch * + output_feature * + (batch / convolution->batch_group_count()) * Product(valid_position_counts); current_properties_[kFlopsKey] = fma_count * kFmaFlops; return Status::OK(); @@ -552,6 +546,21 @@ Status HloCostAnalysis::HandleFft(const HloInstruction* fft) { return Status::OK(); } +Status HloCostAnalysis::HandleTriangularSolve(const HloInstruction* hlo) { + float bytes_accessed = GetShapeSize(hlo->operand(0)->shape()) / 2.0f; + bytes_accessed += GetShapeSize(hlo->operand(1)->shape()); + current_properties_[kBytesAccessedKey] = bytes_accessed; + + const Shape& a_shape = hlo->operand(0)->shape(); + const Shape& b_shape = hlo->operand(1)->shape(); + // 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::HandleAllReduce(const HloInstruction* crs) { // We assume 2 replicas, so that each output element is the sum of two input // elements. @@ -577,6 +586,10 @@ Status HloCostAnalysis::HandleCollectivePermute(const HloInstruction* /*hlo*/) { return Status::OK(); } +Status HloCostAnalysis::HandleReplicaId(const HloInstruction* /*hlo*/) { + return Status::OK(); +} + Status HloCostAnalysis::HandleRng(const HloInstruction* random) { // TODO(b/26346211): Implement better estimates for the RNG cost, since the // cost changes with the implementation and the distribution. For now, assume diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.h b/tensorflow/compiler/xla/service/hlo_cost_analysis.h index b52305626dd67336eb31098d086ad357f12d96c7..96357dec68e390251c43c2c3fc6f5a5612063fbd 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.h +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.h @@ -71,9 +71,11 @@ class HloCostAnalysis : public ConstDfsHloVisitor { Status HandleDot(const HloInstruction* dot) override; Status HandleConvolution(const HloInstruction* convolution) override; Status HandleFft(const HloInstruction* fft) override; + Status HandleTriangularSolve(const HloInstruction* hlo) override; Status HandleAllReduce(const HloInstruction* crs) override; Status HandleAllToAll(const HloInstruction* hlo) override; Status HandleCollectivePermute(const HloInstruction* hlo) override; + Status HandleReplicaId(const HloInstruction* hlo) override; Status HandleInfeed(const HloInstruction* infeed) override; Status HandleOutfeed(const HloInstruction* outfeed) override; Status HandleRng(const HloInstruction* random) override; diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc b/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc index ff32faf298dd1f04c5b769f2a88f76a7a1e18ae7..4d42770ba784ba15fae9518b40a75d8a2f038e66 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc @@ -29,6 +29,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/service.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/compiler/xla/statusor.h" @@ -157,6 +158,87 @@ TEST_F(HloCostAnalysisTest, MatrixMultiply) { sizeof(float) * (10 * 5 + 5 * 30 + 10 * 30)); } +TEST_F(HloCostAnalysisTest, DotGeneral) { + XlaBuilder builder("matrix_multiply"); + auto lhs = + Parameter(&builder, 0, ShapeUtil::MakeShape(F32, {10, 5, 5}), "lhs"); + auto rhs = + Parameter(&builder, 1, ShapeUtil::MakeShape(F32, {5, 5, 30}), "rhs"); + DotDimensionNumbers dnums; + dnums.add_lhs_contracting_dimensions(1); + dnums.add_lhs_contracting_dimensions(2); + dnums.add_rhs_contracting_dimensions(0); + dnums.add_rhs_contracting_dimensions(1); + DotGeneral(lhs, rhs, dnums); + + // Run HLO cost analysis. + auto hlo_module = BuildHloGraph(&builder); + HloCostAnalysis analysis(ShapeSize); + ASSERT_IS_OK( + hlo_module->entry_computation()->root_instruction()->Accept(&analysis)); + + // Check the number of computations returned from the analysis (1500 FMAs). + EXPECT_EQ(analysis.flop_count(), 2 * 10 * 30 * 5 * 5); + + EXPECT_EQ(analysis.transcendental_count(), 0); + + // Bytes accessed is sum of inputs and output. + EXPECT_EQ(analysis.bytes_accessed(), + sizeof(float) * (10 * 5 * 5 + 5 * 5 * 30 + 10 * 30)); +} + +TEST_F(HloCostAnalysisTest, DotGeneral2) { + XlaBuilder builder("matrix_multiply"); + auto lhs = + Parameter(&builder, 0, ShapeUtil::MakeShape(F32, {10, 5, 5}), "lhs"); + auto rhs = + Parameter(&builder, 1, ShapeUtil::MakeShape(F32, {5, 5, 30}), "rhs"); + DotDimensionNumbers dnums; + dnums.add_lhs_contracting_dimensions(1); + dnums.add_lhs_batch_dimensions(2); + dnums.add_rhs_contracting_dimensions(0); + dnums.add_rhs_batch_dimensions(1); + DotGeneral(lhs, rhs, dnums); + + // Run HLO cost analysis. + auto hlo_module = BuildHloGraph(&builder); + HloCostAnalysis analysis(ShapeSize); + ASSERT_IS_OK( + hlo_module->entry_computation()->root_instruction()->Accept(&analysis)); + + // Check the number of computations returned from the analysis (1500 FMAs). + EXPECT_EQ(analysis.flop_count(), 2 * 10 * 30 * 5 * 5); + + EXPECT_EQ(analysis.transcendental_count(), 0); + + // Bytes accessed is sum of inputs and output. + EXPECT_EQ(analysis.bytes_accessed(), + sizeof(float) * (10 * 5 * 5 + 5 * 5 * 30 + 5 * 10 * 30)); +} + +TEST_F(HloCostAnalysisTest, DotGeneral3) { + XlaBuilder builder("matrix_multiply"); + auto lhs = Parameter(&builder, 0, ShapeUtil::MakeShape(F32, {10, 5}), "lhs"); + auto rhs = Parameter(&builder, 1, ShapeUtil::MakeShape(F32, {5, 30}), "rhs"); + DotDimensionNumbers dnums; + DotGeneral(lhs, rhs, dnums); + + // Run HLO cost analysis. + auto hlo_module = BuildHloGraph(&builder); + HloCostAnalysis analysis(ShapeSize); + ASSERT_IS_OK( + hlo_module->entry_computation()->root_instruction()->Accept(&analysis)); + + // Check the number of computations returned from the analysis (1500 FMAs). + EXPECT_EQ(analysis.flop_count(), 2 * 10 * 30 * 5 * 5); + + EXPECT_EQ(analysis.transcendental_count(), 0); + + // Bytes accessed is sum of inputs and output. + EXPECT_EQ(analysis.bytes_accessed(), + sizeof(float) * (10 * 5 + 5 * 30 + 5 * 5 * 10 * 30)); +} + TEST_F(HloCostAnalysisTest, Map) { XlaBuilder builder("map"); auto input = Parameter(&builder, 0, ShapeUtil::MakeShape(F32, {10}), "in"); @@ -529,7 +611,8 @@ TEST_F(HloCostAnalysisTest, DynamicSlice) { // Test the analysis on a slice. XlaBuilder builder("dynamic-slice"); auto x = Parameter(&builder, 0, ShapeUtil::MakeShape(F32, {2}), "x"); - DynamicSlice(x, ConstantR1(&builder, {1}), {1}); + DynamicSlice(x, absl::Span({ConstantR0(&builder, 1)}), + {1}); auto hlo_module = BuildHloGraph(&builder); // Run HLO cost analysis. @@ -545,7 +628,7 @@ TEST_F(HloCostAnalysisTest, DynamicUpdateSlice) { XlaBuilder builder("dynamic-update-slice"); auto x = Parameter(&builder, 0, ShapeUtil::MakeShape(F32, {2}), "x"); DynamicUpdateSlice(x, ConstantR1(&builder, {1.0}), - ConstantR1(&builder, {1})); + absl::Span({ConstantR0(&builder, 1)})); auto hlo_module = BuildHloGraph(&builder); // Run HLO cost analysis. diff --git a/tensorflow/compiler/xla/service/hlo_creation_utils.cc b/tensorflow/compiler/xla/service/hlo_creation_utils.cc index 8cea95a73ed4f9c884c1846ce6418fb5b6cb406e..070115604ba46dfe2de92b592a31e831ca2e1c87 100644 --- a/tensorflow/compiler/xla/service/hlo_creation_utils.cc +++ b/tensorflow/compiler/xla/service/hlo_creation_utils.cc @@ -17,9 +17,15 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/memory/memory.h" #include "absl/strings/str_cat.h" +#include "tensorflow/compiler/xla/client/lib/comparators.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/client/xla_computation.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/hlo_clone_context.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/shape_inference.h" #include "tensorflow/compiler/xla/util.h" @@ -105,12 +111,26 @@ StatusOr MakeDynamicSliceHlo( absl::Span slice_sizes) { HloComputation* computation = operand->parent(); CHECK_EQ(computation, start_indices->parent()); + int64 rank = start_indices->shape().dimensions(0); + std::vector scalar_start_indices; + for (int i = 0; i < rank; ++i) { + // TODO(b/118437727): Update callers to provide scalars directly. + auto slice = computation->AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(start_indices->shape().element_type(), {1}), + start_indices, {i}, {i + 1}, {1})); + scalar_start_indices.push_back( + computation->AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(start_indices->shape().element_type(), {}), + slice))); + } + std::vector scalar_start_indices_shapes( + rank, ShapeUtil::MakeShape(start_indices->shape().element_type(), {})); TF_ASSIGN_OR_RETURN( Shape dynamic_slice_shape, ShapeInference::InferDynamicSliceShape( - operand->shape(), {start_indices->shape()}, slice_sizes)); + operand->shape(), scalar_start_indices_shapes, slice_sizes)); return computation->AddInstruction(HloInstruction::CreateDynamicSlice( - dynamic_slice_shape, operand, start_indices, slice_sizes)); + dynamic_slice_shape, operand, scalar_start_indices, slice_sizes)); } StatusOr MakeDynamicUpdateSliceHlo( @@ -119,17 +139,31 @@ StatusOr MakeDynamicUpdateSliceHlo( HloComputation* computation = operand->parent(); CHECK_EQ(computation, update->parent()); CHECK_EQ(computation, start_indices->parent()); + int64 rank = start_indices->shape().dimensions(0); + std::vector scalar_start_indices; + for (int i = 0; i < rank; ++i) { + // TODO(b/118437727): Update callers to provide scalars directly. + auto slice = computation->AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(start_indices->shape().element_type(), {1}), + start_indices, {i}, {i + 1}, {1})); + scalar_start_indices.push_back( + computation->AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(start_indices->shape().element_type(), {}), + slice))); + } + std::vector scalar_start_indices_shapes( + rank, ShapeUtil::MakeShape(start_indices->shape().element_type(), {})); TF_ASSIGN_OR_RETURN( Shape dynamic_update_slice_shape, ShapeInference::InferDynamicUpdateSliceShape( - operand->shape(), update->shape(), {start_indices->shape()})); + operand->shape(), update->shape(), scalar_start_indices_shapes)); return computation->AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - dynamic_update_slice_shape, operand, update, start_indices)); + dynamic_update_slice_shape, operand, update, scalar_start_indices)); } -StatusOr MakeBroadcastHlo( - HloInstruction* operand, absl::Span broadcast_dimensions, - absl::Span result_shape_bounds) { +HloInstruction* MakeBroadcastHlo(HloInstruction* operand, + absl::Span broadcast_dimensions, + absl::Span result_shape_bounds) { HloComputation* computation = operand->parent(); Shape broadcast_shape = ShapeUtil::MakeShape(operand->shape().element_type(), result_shape_bounds); @@ -239,6 +273,29 @@ StatusOr MakeSelectHlo(HloInstruction* pred, select_shape, HloOpcode::kSelect, pred, on_true, on_false)); } +StatusOr MakeSortHlo( + const Shape& sort_shape, absl::Span operands, + int64 dimension_to_sort, HloComputation::Builder* builder, + HloModule* module) { + CHECK(!operands.empty()) << "Sort Hlo requires at least one operand."; + HloComputation* compare_computation; + XlaBuilder b("Sort.Compare"); + std::vector operand_types(operands.size()); + for (int64 i = 0; i < operands.size(); ++i) { + operand_types[i] = operands[i]->shape().element_type(); + } + XlaComputation comparator = CreateScalarLtComputation(operand_types, &b); + TF_ASSIGN_OR_RETURN(ProgramShape program_shape, comparator.GetProgramShape()); + HloModuleConfig config(program_shape); + TF_ASSIGN_OR_RETURN(auto new_module, + HloModule::CreateFromProto(comparator.proto(), config)); + HloCloneContext context(module); + compare_computation = + module->DeepCloneComputation(new_module->entry_computation(), &context); + return builder->AddInstruction(HloInstruction::CreateSort( + sort_shape, dimension_to_sort, operands, compare_computation)); +} + StatusOr CollapseFirstNDims(HloInstruction* operand, int64 n) { CHECK_GT(n, 0); @@ -365,9 +422,9 @@ StatusOr PadVectorWithZeros(HloInstruction* operand, return MakePadHlo(operand, zero, padding_config); } -StatusOr BroadcastZeros( - HloComputation* computation, PrimitiveType element_type, - absl::Span broadcast_dimensions) { +HloInstruction* BroadcastZeros(HloComputation* computation, + PrimitiveType element_type, + absl::Span broadcast_dimensions) { HloInstruction* zero = computation->AddInstruction( HloInstruction::CreateConstant(LiteralUtil::Zero(element_type))); return MakeBroadcastHlo(zero, /*broadcast_dimensions=*/{}, diff --git a/tensorflow/compiler/xla/service/hlo_creation_utils.h b/tensorflow/compiler/xla/service/hlo_creation_utils.h index 8e5ddbbd503a501bd493aec43a2ccd4db883ef0c..36b8cdc7feff9143a041ad0beb3a0dda91589618 100644 --- a/tensorflow/compiler/xla/service/hlo_creation_utils.h +++ b/tensorflow/compiler/xla/service/hlo_creation_utils.h @@ -82,9 +82,9 @@ StatusOr MakeDynamicUpdateSliceHlo( // Creates a broadcast HLO instruction and adds it to the computation containing // `operand`. -StatusOr MakeBroadcastHlo( - HloInstruction* operand, absl::Span broadcast_dimensions, - absl::Span result_shape_bounds); +HloInstruction* MakeBroadcastHlo(HloInstruction* operand, + absl::Span broadcast_dimensions, + absl::Span result_shape_bounds); // Creates a GetTupleElement HLO instruction and adds it to the computation // containing `operand`. @@ -123,6 +123,15 @@ StatusOr MakeSelectHlo(HloInstruction* pred, HloInstruction* on_true, HloInstruction* on_false); +// Creates a Sort HLO instruction and adds it to the computation containing the +// operands. All operands must be in the same computation. Also creates a +// default compare sub-computation which sorts the first operand into ascending +// order. +StatusOr MakeSortHlo( + const Shape& sort_shape, absl::Span operands, + int64 dimension_to_sort, HloComputation::Builder* builder, + HloModule* module); + // Creates an R1 Constant HLO instruction of the given PrimitiveType with the // given values and adds it to the given computation. template @@ -198,9 +207,9 @@ StatusOr PadVectorWithZeros(HloInstruction* operand, // Broadcasts a zero value of type `element_type` into a tensor with element // type `element_type` and dimension bounds `broadcast_dimensions`. The // broadcast instruction is emitted into `computation`. -StatusOr BroadcastZeros( - HloComputation* computation, PrimitiveType element_type, - absl::Span broadcast_dimensions); +HloInstruction* BroadcastZeros(HloComputation* computation, + PrimitiveType element_type, + absl::Span broadcast_dimensions); // Creates a HLO computation that takes arguments of type `domain` and produces // a value of type `range`. diff --git a/tensorflow/compiler/xla/service/hlo_creation_utils_test.cc b/tensorflow/compiler/xla/service/hlo_creation_utils_test.cc index 3715e12b4e2baf7bc2149237457c16c3919c5083..6025e6a77941369f75ebaa98bdf0979669b3a03c 100644 --- a/tensorflow/compiler/xla/service/hlo_creation_utils_test.cc +++ b/tensorflow/compiler/xla/service/hlo_creation_utils_test.cc @@ -191,9 +191,8 @@ TEST_F(HloCreationUtilsTest, BroadcastZeros_S32) { /*output_shape_dims=*/{2, 2}, ¶m, &entry_computation); - TF_ASSERT_OK_AND_ASSIGN( - HloInstruction * zeros, - BroadcastZeros(module->entry_computation(), S32, {2, 2})); + HloInstruction* zeros = + BroadcastZeros(module->entry_computation(), S32, {2, 2}); entry_computation->set_root_instruction(zeros); HloEvaluator evaluator; @@ -211,9 +210,8 @@ TEST_F(HloCreationUtilsTest, BroadcastZeros_F32) { /*output_shape_dims=*/{2, 2}, ¶m, &entry_computation); - TF_ASSERT_OK_AND_ASSIGN( - HloInstruction * zeros, - BroadcastZeros(module->entry_computation(), F32, {2, 2})); + HloInstruction* zeros = + BroadcastZeros(module->entry_computation(), F32, {2, 2}); entry_computation->set_root_instruction(zeros); HloEvaluator evaluator; diff --git a/tensorflow/compiler/xla/service/hlo_cse.cc b/tensorflow/compiler/xla/service/hlo_cse.cc index e602107cbe64320a8e8e740168cb294ec6be9667..849cac278ee379122ba1ff9fade3bf003969b8a7 100644 --- a/tensorflow/compiler/xla/service/hlo_cse.cc +++ b/tensorflow/compiler/xla/service/hlo_cse.cc @@ -33,7 +33,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/hash/hash.h" diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc index 1924204df044faa6d55f11c1180e2ecc1f0e9e64..3144a84805454488f417391f40ed6b9e9facc752 100644 --- a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc @@ -107,7 +107,7 @@ bool HloDataflowAnalysis::AreTransitiveUsesElementwiseOrTuple( return false; } } - if (!visited.count(user)) { + if (!visited.contains(user)) { stack.push_back(user); } } @@ -256,7 +256,7 @@ bool HloDataflowAnalysis::Phi( input_value_ids.push_back(value->id()); } } - std::sort(input_value_ids.begin(), input_value_ids.end()); + absl::c_sort(input_value_ids); input_value_ids.erase( std::unique(input_value_ids.begin(), input_value_ids.end()), input_value_ids.end()); @@ -271,8 +271,7 @@ bool HloDataflowAnalysis::Phi( if (current_value_defined_here) { VLOG(5) << "current_value_defined_here: " << current_value->ToString(); CHECK(current_value->is_phi()); - auto it = std::find(input_value_ids.begin(), input_value_ids.end(), - current_value->id()); + auto it = absl::c_find(input_value_ids, current_value->id()); if (it != input_value_ids.end()) { input_value_ids.erase(it); } @@ -921,8 +920,7 @@ StatusOr> HloDataflowAnalysis::Run( for (auto& pair : dataflow_analysis->values_) { dataflow_analysis->values_vector_.push_back(&pair.second); } - std::sort(dataflow_analysis->values_vector_.begin(), - dataflow_analysis->values_vector_.end(), HloValue::IdLessThan); + absl::c_sort(dataflow_analysis->values_vector_, HloValue::IdLessThan); TF_DCHECK_OK(dataflow_analysis->Verify()); @@ -937,9 +935,7 @@ Status HloDataflowAnalysis::Verify() const { for (const HloValue* value : values()) { for (const HloPosition& position : value->positions()) { const HloValueSet& value_set = GetValueSet(position); - TF_RET_CHECK(std::find(value_set.values().begin(), - value_set.values().end(), - value) != value_set.values().end()) + TF_RET_CHECK(absl::c_linear_search(value_set.values(), value)) << "Value set at position " << position << " does not contain value " << value->ToShortString(); } @@ -954,9 +950,7 @@ Status HloDataflowAnalysis::Verify() const { const HloValueSet& value_set = pair.second; const HloPosition position{instruction, index}; for (const HloValue* value : value_set.values()) { - TF_RET_CHECK(std::find(value->positions().begin(), - value->positions().end(), - position) != value->positions().end()) + TF_RET_CHECK(absl::c_linear_search(value->positions(), position)) << "Value set at position " << position << " unexpectedly contains value " << value->ToShortString(); } @@ -1041,11 +1035,10 @@ bool HloDataflowAnalysis::CanShareOperandBufferWithUser( // Check if one operand of kAdd fused root is kDot or kConvolution. auto* add = user->fused_expression_root(); auto add_operand_it = - std::find_if(add->operands().begin(), add->operands().end(), - [&](HloInstruction* operand) { - return operand->opcode() == HloOpcode::kConvolution || - operand->opcode() == HloOpcode::kDot; - }); + absl::c_find_if(add->operands(), [&](HloInstruction* operand) { + return operand->opcode() == HloOpcode::kConvolution || + operand->opcode() == HloOpcode::kDot; + }); if (add_operand_it == add->operands().end()) { return false; } @@ -1100,16 +1093,15 @@ bool HloDataflowAnalysis::CanShareOperandBufferWithUser( // *) The root instruction of the called computation is element-wise on // 'operand'. const bool found_caller_use = - std::find_if(uses.begin(), uses.end(), [user](const HloUse& use) { + absl::c_find_if(uses, [user](const HloUse& use) { return use.instruction == user; }) != uses.end(); auto* callee_root = user->to_apply()->root_instruction(); const bool found_elementwise_callee_use = - std::find_if( - uses.begin(), uses.end(), [callee_root](const HloUse& use) { - return use.instruction == callee_root && - callee_root->IsElementwiseOnOperand(use.operand_number); - }) != uses.end(); + absl::c_find_if(uses, [callee_root](const HloUse& use) { + return use.instruction == callee_root && + callee_root->IsElementwiseOnOperand(use.operand_number); + }) != uses.end(); return uses.size() == 2 && found_caller_use && found_elementwise_callee_use; } diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc b/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc index 1d165ac35f766af857dc0984d3f7012ad945b3c2..e3059e02cf0e527522811920a09154afd32976f5 100644 --- a/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc +++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_creation_utils.h" #include "tensorflow/compiler/xla/service/hlo_graph_dumper.h" #include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" @@ -1901,9 +1902,9 @@ ENTRY %AddDependency (p: f32[3]) -> f32[3] { EXPECT_FALSE(analysis->ValueIsDefinedAt(root)); } -INSTANTIATE_TEST_CASE_P(HloDataflowAnalysisInstantiation, - HloDataflowAnalysisTest, - ::testing::Values(false, true)); +INSTANTIATE_TEST_SUITE_P(HloDataflowAnalysisInstantiation, + HloDataflowAnalysisTest, + ::testing::Values(false, true)); class HloDataflowAnalysisTestBase : public HloTestBase { protected: @@ -1970,12 +1971,13 @@ TEST_F(DoesNotUseOperandBufferTest, FusedDynamicUpdateSlice) { // Create a DynamicUpdateSlice instruction of tuple element 1. auto starts = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({2}))); + HloInstruction::CreateConstant(LiteralUtil::CreateR0(2))); auto update = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({2.f, 2.f, 2.f}))); auto dynamic_update_slice = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - data_shape, gte1, update, starts)); + data_shape, gte1, update, + std::initializer_list({starts}))); builder.AddInstruction( HloInstruction::CreateTuple({gte0, dynamic_update_slice})); @@ -2012,12 +2014,13 @@ TEST_F(DoesNotUseOperandBufferTest, IndirectUses) { // Create a DynamicUpdateSlice instruction of tuple element 1. auto starts = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({2}))); + HloInstruction::CreateConstant(LiteralUtil::CreateR0(2))); auto update = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({2.f, 2.f, 2.f}))); auto dynamic_update_slice = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - data_shape, gte1, update, starts)); + data_shape, gte1, update, + std::initializer_list({starts}))); builder.AddInstruction( HloInstruction::CreateTuple({gte0, dynamic_update_slice})); @@ -2150,17 +2153,17 @@ TEST_F(CanShareOperandBufferWithUserTest, auto param = builder.AddInstruction( HloInstruction::CreateParameter(0, data_shape, "param0")); - auto index = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({0, 0}))); - auto ds = builder.AddInstruction( - HloInstruction::CreateDynamicSlice(slice_shape, param, index, {1, 2, 2})); + auto zero = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); + auto ds = builder.AddInstruction(HloInstruction::CreateDynamicSlice( + slice_shape, param, {zero, zero}, {1, 2, 2})); - auto dus = builder.AddInstruction( - HloInstruction::CreateDynamicUpdateSlice(data_shape, param, ds, index)); + auto dus = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( + data_shape, param, ds, {zero, zero})); BuildModule(builder.Build()); auto fusion = computation_->CreateFusionInstruction( - {dus, ds, index}, HloInstruction::FusionKind::kLoop); + {dus, ds, zero}, HloInstruction::FusionKind::kLoop); RunAnalysis(); EXPECT_TRUE( @@ -2219,12 +2222,13 @@ TEST_F(CanShareOperandBufferWithUserTest, FusedDynamicUpdateSlice) { // Create a DynamicUpdateSlice instruction of tuple element 1. auto starts = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({2}))); + HloInstruction::CreateConstant(LiteralUtil::CreateR0(2))); auto update = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({2.f, 2.f, 2.f}))); auto dynamic_update_slice = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - data_shape, gte1, update, starts)); + data_shape, gte1, update, + std::initializer_list({starts}))); builder.AddInstruction( HloInstruction::CreateTuple({gte0, dynamic_update_slice})); @@ -2259,12 +2263,13 @@ TEST_F(CanShareOperandBufferWithUserTest, // Create a DynamicUpdateSlice instruction of tuple element 1. auto starts = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({2}))); + HloInstruction::CreateConstant(LiteralUtil::CreateR0(2))); auto update = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({2.f, 2.f, 2.f}))); auto dynamic_update_slice = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - data_shape_bf16, convert1, update, starts)); + data_shape_bf16, convert1, update, + std::initializer_list({starts}))); auto convert2 = builder.AddInstruction( HloInstruction::CreateConvert(data_shape, dynamic_update_slice)); @@ -2290,10 +2295,13 @@ TEST_F(CanShareOperandBufferWithUserTest, DynamicUpdateSliceCanShare) { HloInstruction::CreateParameter(0, data_shape, "data")); auto update = builder.AddInstruction( HloInstruction::CreateParameter(1, update_shape, "update")); - auto starts = builder.AddInstruction( - HloInstruction::CreateParameter(2, starts_shape, "starts")); + auto start0 = builder.AddInstruction( + HloInstruction::CreateParameter(2, starts_shape, "start0")); + auto start1 = builder.AddInstruction( + HloInstruction::CreateParameter(3, starts_shape, "start1")); + auto dus = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - data_shape, data, update, starts)); + data_shape, data, update, {start0, start1})); BuildModuleAndRunAnalysis(builder.Build()); @@ -2304,7 +2312,9 @@ TEST_F(CanShareOperandBufferWithUserTest, DynamicUpdateSliceCanShare) { EXPECT_FALSE( dataflow_analysis_->CanShareOperandBufferWithUser(update, {}, dus, {})); EXPECT_FALSE( - dataflow_analysis_->CanShareOperandBufferWithUser(starts, {}, dus, {})); + dataflow_analysis_->CanShareOperandBufferWithUser(start0, {}, dus, {})); + EXPECT_FALSE( + dataflow_analysis_->CanShareOperandBufferWithUser(start1, {}, dus, {})); } TEST_F(CanShareOperandBufferWithUserTest, ScatterCanShare) { @@ -2347,14 +2357,16 @@ TEST_F(CanShareOperandBufferWithUserTest, ScatterCanShare) { TEST_F(CanShareOperandBufferWithUserTest, SortCanShare) { auto builder = HloComputation::Builder(TestName()); + module_ = CreateNewVerifiedModule(); Shape keys_shape = ShapeUtil::MakeShape(F32, {8}); auto keys = builder.AddInstruction( HloInstruction::CreateParameter(0, keys_shape, "keys")); - auto sort = - builder.AddInstruction(HloInstruction::CreateSort(keys_shape, 0, keys)); + TF_ASSERT_OK_AND_ASSIGN( + auto* sort, MakeSortHlo(keys_shape, {keys}, -1, &builder, module_.get())); - BuildModuleAndRunAnalysis(builder.Build()); + computation_ = module_->AddEntryComputation(builder.Build()); + RunAnalysis(); EXPECT_TRUE( dataflow_analysis_->CanShareOperandBufferWithUser(keys, {}, sort, {})); @@ -2362,6 +2374,7 @@ TEST_F(CanShareOperandBufferWithUserTest, SortCanShare) { TEST_F(CanShareOperandBufferWithUserTest, SortCanShareWithTupleUser) { auto builder = HloComputation::Builder(TestName()); + module_ = CreateNewVerifiedModule(); Shape keys_shape = ShapeUtil::MakeShape(F32, {8}); Shape values_shape = ShapeUtil::MakeShape(F32, {8}); @@ -2369,11 +2382,13 @@ TEST_F(CanShareOperandBufferWithUserTest, SortCanShareWithTupleUser) { HloInstruction::CreateParameter(0, keys_shape, "keys")); auto values = builder.AddInstruction( HloInstruction::CreateParameter(1, values_shape, "values")); - auto sort = builder.AddInstruction(HloInstruction::CreateSort( - ShapeUtil::MakeTupleShape({keys_shape, values_shape}), 0, keys, - {values})); + TF_ASSERT_OK_AND_ASSIGN( + auto* sort, + MakeSortHlo(ShapeUtil::MakeTupleShape({keys_shape, values_shape}), + {keys, values}, 0, &builder, module_.get())); - BuildModuleAndRunAnalysis(builder.Build()); + computation_ = module_->AddEntryComputation(builder.Build()); + RunAnalysis(); // The buffer for the keys can be shared with the first tuple entry. EXPECT_TRUE( diff --git a/tensorflow/compiler/xla/service/hlo_dce.cc b/tensorflow/compiler/xla/service/hlo_dce.cc index 7d35e251ca21951036336ff1a1eb4aabc87bc5ca..a5a11f09cf4f857b992e5ede3a9dbc5a937ce722 100644 --- a/tensorflow/compiler/xla/service/hlo_dce.cc +++ b/tensorflow/compiler/xla/service/hlo_dce.cc @@ -20,6 +20,7 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_set.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" @@ -65,7 +66,7 @@ StatusOr HloDCE::Run(HloModule* module) { // Now DCE HloComputations. First, collect the computations that are // referenced by some remaining instruction. - std::unordered_set live_computations; + absl::flat_hash_set live_computations; if (HloComputation* entry_computation = module->entry_computation()) { live_computations.insert(entry_computation); } @@ -79,7 +80,7 @@ StatusOr HloDCE::Run(HloModule* module) { // Remove dead computations. for (auto* computation : module->MakeComputationPostOrder()) { - if (live_computations.count(computation) == 0) { + if (!live_computations.contains(computation)) { TF_RETURN_IF_ERROR(module->RemoveEmbeddedComputation(computation)); changed = true; } diff --git a/tensorflow/compiler/xla/service/hlo_dce_test.cc b/tensorflow/compiler/xla/service/hlo_dce_test.cc index 1fa4259a3e42286cbc911907eea563e6ca6f8611..b5d72b386f89568cc3066b2e497be98428d1ed0c 100644 --- a/tensorflow/compiler/xla/service/hlo_dce_test.cc +++ b/tensorflow/compiler/xla/service/hlo_dce_test.cc @@ -43,9 +43,7 @@ class HloDceTest : public HloTestBase { // Returns whether the given instruction exists in the given computation. bool HasInstruction(const HloComputation& computation, const HloInstruction* instruction) { - return std::find(computation.instructions().begin(), - computation.instructions().end(), - instruction) != computation.instructions().end(); + return absl::c_linear_search(computation.instructions(), instruction); } }; diff --git a/tensorflow/compiler/xla/service/hlo_domain_map.cc b/tensorflow/compiler/xla/service/hlo_domain_map.cc index c6d02f9f67bb599e496d20fc2acf2e627ed54438..7cdb7f6bdf26241cda4fabbb5ccaf6e6f7de39ce 100644 --- a/tensorflow/compiler/xla/service/hlo_domain_map.cc +++ b/tensorflow/compiler/xla/service/hlo_domain_map.cc @@ -230,10 +230,10 @@ HloDomainMap::MakeNonDomainInstructions( } } // sort instructions according to instructions_order - std::sort(instructions.begin(), instructions.end(), - [&instructions_order](HloInstruction* a, HloInstruction* b) { - return instructions_order.at(a) < instructions_order.at(b); - }); + absl::c_sort(instructions, + [&instructions_order](HloInstruction* a, HloInstruction* b) { + return instructions_order.at(a) < instructions_order.at(b); + }); return instructions; } diff --git a/tensorflow/compiler/xla/service/hlo_element_type_converter.cc b/tensorflow/compiler/xla/service/hlo_element_type_converter.cc index 9b0f2b2a0f4dd5d1d1191e9ab0637cc3034b50da..7d6b86056af3fc2128fe1642bbfa0ca6f9ef1da0 100644 --- a/tensorflow/compiler/xla/service/hlo_element_type_converter.cc +++ b/tensorflow/compiler/xla/service/hlo_element_type_converter.cc @@ -127,6 +127,7 @@ StatusOr HloElementTypeConverter::Run(HloModule* module) { // These are ops where it does not make sense to convert them. if (opcode == HloOpcode::kParameter || opcode == HloOpcode::kConstant || opcode == HloOpcode::kTuple || opcode == HloOpcode::kConvert || + opcode == HloOpcode::kBitcastConvert || opcode == HloOpcode::kGetTupleElement || opcode == HloOpcode::kInfeed || opcode == HloOpcode::kOutfeed) { continue; @@ -145,7 +146,7 @@ StatusOr HloElementTypeConverter::Run(HloModule* module) { opcode == HloOpcode::kMap || opcode == HloOpcode::kReduce || opcode == HloOpcode::kReduceWindow || opcode == HloOpcode::kScatter || opcode == HloOpcode::kSelectAndScatter || - opcode == HloOpcode::kConditional) { + opcode == HloOpcode::kSort || opcode == HloOpcode::kConditional) { continue; } TF_RET_CHECK(hlo->called_computations().empty()) << hlo->ToString(); diff --git a/tensorflow/compiler/xla/service/hlo_element_type_converter_test.cc b/tensorflow/compiler/xla/service/hlo_element_type_converter_test.cc index a3b56a44a0b02923585c1dcb69571479236188a3..4171f738620dbf545e5883b8c26169fae4b93643 100644 --- a/tensorflow/compiler/xla/service/hlo_element_type_converter_test.cc +++ b/tensorflow/compiler/xla/service/hlo_element_type_converter_test.cc @@ -28,15 +28,7 @@ using ::testing::Eq; using ::testing::Not; using ::testing::ResultOf; -class HloElementTypeConverterTest : public HloTestBase { - public: - std::unique_ptr CreateModuleFromHloString( - const string& hlo_string) { - return HloRunner::CreateModuleFromString(hlo_string, - GetDebugOptionsForTest()) - .ValueOrDie(); - } -}; +using HloElementTypeConverterTest = HloTestBase; TEST_F(HloElementTypeConverterTest, CustomCallsNotConverted) { const string& hlo_string = R"( @@ -47,7 +39,7 @@ TEST_F(HloElementTypeConverterTest, CustomCallsNotConverted) { custom_call_target="foo" } )"; - auto module = CreateModuleFromHloString(hlo_string); + auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); HloElementTypeConverter type_converter(BF16, F32); TF_ASSERT_OK_AND_ASSIGN(bool converted, type_converter.Run(module.get())); EXPECT_FALSE(converted); @@ -63,7 +55,7 @@ TEST_F(HloElementTypeConverterTest, InfeedsOutfeedsNotConverted) { outfeed = token[] outfeed(infeed.data, token0) } )"; - auto module = CreateModuleFromHloString(hlo_string); + auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); HloElementTypeConverter type_converter(BF16, F32); TF_ASSERT_OK_AND_ASSIGN(bool converted, type_converter.Run(module.get())); EXPECT_FALSE(converted); @@ -73,17 +65,16 @@ TEST_F(HloElementTypeConverterTest, OperationsInNestedTuplesConverted) { const string& hlo_string = R"( HloModule NestedTuples ENTRY NestedTuples.v5 { - constant.4 = bf16[] constant(42) constant.2 = f32[2]{0} constant({1, 2}) - constant.3 = bf16[] constant(42) - add = bf16[] add(constant.2, constant.3) - tuple = (f32[2]{0}, bf16[]) tuple(constant.2, add) + constant.3 = bf16[2]{0} constant({42, 42}) + add = bf16[2]{0} add(constant.2, constant.3) + tuple = (f32[2]{0}, bf16[2]{0}) tuple(constant.2, add) constant.5 = bf16[2]{0} constant({22, 44}) - ROOT tuple.1 = ((f32[2]{0}, bf16[]), bf16[2]{0}) tuple(tuple, constant.5) + ROOT tuple.1 = ((f32[2]{0}, bf16[2]{0}), bf16[2]{0}) tuple(tuple, constant.5) } )"; - auto module = CreateModuleFromHloString(hlo_string); + auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); HloElementTypeConverter type_converter(BF16, F32); TF_ASSERT_OK_AND_ASSIGN(bool converted, type_converter.Run(module.get())); EXPECT_TRUE(converted); @@ -111,7 +102,7 @@ TEST_F(HloElementTypeConverterTest, BatchNormGradBF16Converted) { } )"; - auto module = CreateModuleFromHloString(hlo_string); + auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); HloElementTypeConverter type_converter(BF16, F32); TF_ASSERT_OK_AND_ASSIGN(bool converted, type_converter.Run(module.get())); EXPECT_TRUE(converted); @@ -135,7 +126,7 @@ ENTRY main { ROOT rng = bf16[1,1000,20]{2,1,0} rng(constant.3, constant.4), distribution=rng_uniform } )"; - auto module = CreateModuleFromHloString(hlo_string); + auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); HloElementTypeConverter type_converter(BF16, F32); TF_ASSERT_OK_AND_ASSIGN(bool converted, type_converter.Run(module.get())); EXPECT_TRUE(converted); @@ -161,7 +152,7 @@ ENTRY main { ROOT rng1 = bf16[1,1000,20]{2,1,0} rng(constant.3, constant.4), control-predecessors={%rng0}, distribution=rng_uniform } )"; - auto module = CreateModuleFromHloString(hlo_string); + auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); HloElementTypeConverter type_converter(BF16, F32); TF_ASSERT_OK_AND_ASSIGN(bool converted, type_converter.Run(module.get())); @@ -185,5 +176,19 @@ ENTRY main { EXPECT_THAT(rng1->control_predecessors(), ElementsAre(rng0)); } +TEST_F(HloElementTypeConverterTest, BitcastConvertIsUnmodified) { + const string& hlo_string = R"( + HloModule test + + ENTRY test { + p = bf16[] parameter(0) + ROOT c = u16[] bitcast-convert(p) + })"; + auto module = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + HloElementTypeConverter converter(BF16, F32); + TF_ASSERT_OK_AND_ASSIGN(bool converted, RunHloPass(&converter, module.get())); + EXPECT_FALSE(converted); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.cc b/tensorflow/compiler/xla/service/hlo_evaluator.cc index 794b97a7e0728862123d91048b36bedf2b114cac..691d5c1bbc3edb0aa47acc52d5752020068c3515 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator.cc +++ b/tensorflow/compiler/xla/service/hlo_evaluator.cc @@ -18,9 +18,9 @@ limitations under the License. #include #include #include +#include #include #include -#include #include #include "absl/algorithm/container.h" @@ -29,7 +29,6 @@ limitations under the License. #include "absl/strings/string_view.h" #include "tensorflow/compiler/xla/index_util.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/primitive_util.h" @@ -136,8 +135,44 @@ StatusOr Compare(const Shape& shape, HloOpcode opcode, return std::move(result); } +template <> +StatusOr Compare(const Shape& shape, HloOpcode opcode, + LiteralSlice lhs_literal, + LiteralSlice rhs_literal) { + std::function compare_op; + switch (opcode) { + case HloOpcode::kEq: + compare_op = [](complex128 lhs_el, complex128 rhs_el) { + return lhs_el == rhs_el; + }; + break; + case HloOpcode::kNe: + compare_op = [](complex128 lhs_el, complex128 rhs_el) { + return lhs_el != rhs_el; + }; + break; + default: + LOG(FATAL) << "unhandled HLO opcode for conversion to Comparison: " + << HloOpcodeString(opcode); + } + + Literal result(shape); + TF_RETURN_IF_ERROR( + result.Populate([&](absl::Span multi_index) { + return compare_op(lhs_literal.Get(multi_index), + rhs_literal.Get(multi_index)); + })); + + return std::move(result); +} + } // namespace +// Note that unsupported types by the typed visitor does not necessarily imply +// the non-typed HloEvaluator (parent evaluator) would not support them either +// in the type-agnostic handler. For e.g., HandleGetTupleElement in the parent +// type-agnostic evaluator will be able to accept Tuple primitive type, whereas +// HloEvaluatorTypedVisitor cannot. HloEvaluator::HloEvaluator(int64 max_loop_iterations) : max_loop_iterations_(max_loop_iterations) { typed_visitors_[PRED] = @@ -145,22 +180,14 @@ HloEvaluator::HloEvaluator(int64 max_loop_iterations) typed_visitors_[U8] = absl::make_unique>(this); typed_visitors_[U16] = - absl::make_unique([](HloInstruction*) { - return Unimplemented( - "HloEvaluator::HloEvaluatorTypedVisitor: unhandled primitive type: " - "U16."); - }); + absl::make_unique>(this); typed_visitors_[U32] = absl::make_unique>(this); typed_visitors_[U64] = absl::make_unique>(this); typed_visitors_[S8] = absl::make_unique>(this); typed_visitors_[S16] = - absl::make_unique([](HloInstruction*) { - return Unimplemented( - "HloEvaluator::HloEvaluatorTypedVisitor: unhandled primitive type: " - "S16."); - }); + absl::make_unique>(this); typed_visitors_[S32] = absl::make_unique>(this); typed_visitors_[S64] = @@ -173,6 +200,8 @@ HloEvaluator::HloEvaluator(int64 max_loop_iterations) absl::make_unique>(this); typed_visitors_[C64] = absl::make_unique>(this); + typed_visitors_[C128] = + absl::make_unique>(this); // Most of the evaluator computations we use don't support BF16 (e.g., // std::ceil, std::tanh). To make evaluator work with BF16, we set all @@ -229,6 +258,18 @@ StatusOr HloEvaluator::Evaluate( arg_literals_.push_back(&*literal_ptr); } + // Re-seed RNG, either from the configuration's seed or a monotonic + // per-evaluator seed (which prevents two evaluators from returning the same + // random sequence). + if (computation.parent()->config().seed()) { + seed_ = computation.parent()->config().seed(); + } else { + // Start global_seed at a (true) random value. + static std::atomic global_seed{std::random_device()()}; + seed_ = global_seed.fetch_add(1); + } + engine_.seed(seed_); + TF_RETURN_IF_ERROR(computation.Accept(this)); return GetEvaluatedLiteralFor(computation.root_instruction()).Clone(); } @@ -348,16 +389,45 @@ Status HloEvaluator::HandleBitcast(HloInstruction* bitcast) { return Status::OK(); } +Status HloEvaluator::HandleGetDimensionSize( + HloInstruction* get_dimension_size) { + HloInstruction* operand = get_dimension_size->mutable_operand(0); + int64 dim = get_dimension_size->dimension(); + if (dynamic_dimension_inference_ == nullptr) { + return InvalidArgument( + "Evaluator cannot evaluate get_dimension_size without " + "set_dynamic_dimension_inference."); + } + HloInstruction* dynamic_size = + dynamic_dimension_inference_->GetDynamicSize(operand, {}, dim); + if (dynamic_size != nullptr) { + evaluated_[get_dimension_size] = + GetEvaluatedLiteralFor(dynamic_size).Clone(); + return Status::OK(); + } + + const Shape& shape = get_dimension_size->operand(0)->shape(); + Literal output(ShapeUtil::MakeShape(U32, {})); + output.PopulateWithValue( + static_cast(shape.dimensions(get_dimension_size->dimension()))); + evaluated_[get_dimension_size] = std::move(output); + return Status::OK(); +} + Status HloEvaluator::HandleParameter(HloInstruction* parameter) { + // Nothing to do other than sanity checks. Parameters' values are stored in + // arg_literals_. CHECK_LT(parameter->parameter_number(), arg_literals_.size()); + +#ifndef NDEBUG const Literal* input_literal = arg_literals_[parameter->parameter_number()]; VLOG(2) << "Parameter evaluated to: " << input_literal->ToString(); DCHECK(ShapeUtil::Equal(parameter->shape(), input_literal->shape())) << "parameter shape is: " << ShapeUtil::HumanString(parameter->shape()) << ", but input literal shape is: " << ShapeUtil::HumanString(input_literal->shape()); +#endif - evaluated_[parameter] = input_literal->Clone(); return Status::OK(); } @@ -420,15 +490,52 @@ Status HloEvaluator::HandleConcatenate(HloInstruction* concatenate) { Status HloEvaluator::HandleIsFinite(HloInstruction* is_finite) { auto operand = is_finite->operand(0); - if (!ShapeUtil::ElementIsFloating(operand->shape())) { - return InvalidArgument( - "expected element type in shape to be float for IsFinite op, got: %s", - PrimitiveType_Name(operand->shape().element_type())); - } + auto elem_ty = operand->shape().element_type(); + switch (elem_ty) { + case PRED: + case TUPLE: + case OPAQUE: + case TOKEN: + case S8: + case S16: + case S32: + case S64: + case U8: + case U16: + case U32: + case U64: + case C64: + case C128: + // Explicitly enumerate all types in this switch so that when we add a new + // type, we'll get a compile error here. + case PRIMITIVE_TYPE_INVALID: + case PrimitiveType_INT_MIN_SENTINEL_DO_NOT_USE_: + case PrimitiveType_INT_MAX_SENTINEL_DO_NOT_USE_: + return InvalidArgument( + "expected element type in shape to be floating point, but " + "got: %s", + PrimitiveType_Name(elem_ty)); - switch (operand->shape().element_type()) { - case F16: - return Unimplemented("unhandled primitive type: F16."); + case F16: { + auto result_or = ElementWiseUnaryOpImpl( + is_finite, + [](Eigen::half elem_operand) { + return std::isfinite(static_cast(elem_operand)); + }, + GetEvaluatedLiteralFor(operand)); + TF_ASSIGN_OR_RETURN(evaluated_[is_finite], std::move(result_or)); + break; + } + case BF16: { + auto result_or = ElementWiseUnaryOpImpl( + is_finite, + [](bfloat16 elem_operand) { + return std::isfinite(static_cast(elem_operand)); + }, + GetEvaluatedLiteralFor(operand)); + TF_ASSIGN_OR_RETURN(evaluated_[is_finite], std::move(result_or)); + break; + } case F32: { auto result_or = ElementWiseUnaryOpImpl( is_finite, @@ -445,9 +552,6 @@ Status HloEvaluator::HandleIsFinite(HloInstruction* is_finite) { TF_ASSIGN_OR_RETURN(evaluated_[is_finite], std::move(result_or)); break; } - default: - LOG(FATAL) << "HandleIsFinite: unknown/unhandled primitive type: " - << PrimitiveType_Name(operand->shape().element_type()); } return Status::OK(); @@ -470,6 +574,13 @@ Status HloEvaluator::HandleReal(HloInstruction* real) { TF_ASSIGN_OR_RETURN(evaluated_[real], std::move(result_or)); break; } + case C128: { + auto result_or = ElementWiseUnaryOpImpl( + real, [](complex128 elem_operand) { return std::real(elem_operand); }, + GetEvaluatedLiteralFor(operand)); + TF_ASSIGN_OR_RETURN(evaluated_[real], std::move(result_or)); + break; + } case F16: { auto result_or = ElementWiseUnaryOpImpl( real, [](Eigen::half elem_operand) { return elem_operand; }, @@ -500,11 +611,29 @@ Status HloEvaluator::HandleReal(HloInstruction* real) { } Status HloEvaluator::HandleImag(HloInstruction* imag) { - auto result_or = ElementWiseUnaryOpImpl( - imag, [](complex64 elem_operand) { return std::imag(elem_operand); }, - GetEvaluatedLiteralFor(imag->operand(0))); + auto operand = imag->operand(0); + switch (operand->shape().element_type()) { + case C64: { + auto result_or = ElementWiseUnaryOpImpl( + imag, [](complex64 elem_operand) { return std::imag(elem_operand); }, + GetEvaluatedLiteralFor(imag->operand(0))); + + TF_ASSIGN_OR_RETURN(evaluated_[imag], std::move(result_or)); + break; + } + case C128: { + auto result_or = ElementWiseUnaryOpImpl( + imag, [](complex128 elem_operand) { return std::imag(elem_operand); }, + GetEvaluatedLiteralFor(imag->operand(0))); + + TF_ASSIGN_OR_RETURN(evaluated_[imag], std::move(result_or)); + break; + } + default: + LOG(FATAL) << "HandleImag: unknown/unhandled primitive type: " + << PrimitiveType_Name(operand->shape().element_type()); + } - TF_ASSIGN_OR_RETURN(evaluated_[imag], std::move(result_or)); return Status::OK(); } @@ -514,11 +643,27 @@ Status HloEvaluator::HandleComplex(HloInstruction* complex) { TF_RET_CHECK(ShapeUtil::Compatible(real.shape(), imag.shape())); Literal result(complex->shape()); - TF_RETURN_IF_ERROR( - result.Populate([&](absl::Span multi_index) { - return std::complex(real.Get(multi_index), - imag.Get(multi_index)); - })); + switch (complex->shape().element_type()) { + case C64: { + TF_RETURN_IF_ERROR( + result.Populate([&](absl::Span multi_index) { + return std::complex(real.Get(multi_index), + imag.Get(multi_index)); + })); + break; + } + case C128: { + TF_RETURN_IF_ERROR( + result.Populate([&](absl::Span multi_index) { + return std::complex(real.Get(multi_index), + imag.Get(multi_index)); + })); + break; + } + default: + LOG(FATAL) << "HandleComplex: unknown/unhandled primitive type: " + << PrimitiveType_Name(complex->shape().element_type()); + } evaluated_[complex] = std::move(result); return Status::OK(); @@ -557,8 +702,11 @@ Status HloEvaluator::HandleCompare(HloInstruction* compare) { evaluated_[compare], Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); } break; - case U16: - return Unimplemented("unhandled primitive type: U16."); + case U16: { + TF_ASSIGN_OR_RETURN( + evaluated_[compare], + Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); + } break; case U32: { TF_ASSIGN_OR_RETURN( evaluated_[compare], @@ -574,8 +722,11 @@ Status HloEvaluator::HandleCompare(HloInstruction* compare) { evaluated_[compare], Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); } break; - case S16: - return Unimplemented("unhandled primitive type: S16."); + case S16: { + TF_ASSIGN_OR_RETURN( + evaluated_[compare], + Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); + } break; case S32: { TF_ASSIGN_OR_RETURN( evaluated_[compare], @@ -611,6 +762,11 @@ Status HloEvaluator::HandleCompare(HloInstruction* compare) { Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); } break; + case C128: { + TF_ASSIGN_OR_RETURN(evaluated_[compare], + Compare(compare->shape(), opcode, + lhs_literal, rhs_literal)); + } break; default: LOG(FATAL) << "HandleCompare: unknown primitive type: " << PrimitiveType_Name(lhs->shape().element_type()); @@ -1067,6 +1223,8 @@ Status HloEvaluator::HandleCall(HloInstruction* call) { } HloEvaluator embedded_evaluator; + embedded_evaluator.set_dynamic_dimension_inference( + dynamic_dimension_inference_); Literal result = embedded_evaluator.Evaluate(*computation, arg_literals) .ConsumeValueOrDie(); @@ -1084,7 +1242,9 @@ Status HloEvaluator::HandleFusion(HloInstruction* fusion) { fusion->fused_instructions_computation()->Clone( /*suffix=*/"clone_with_layout", &context); for (auto* instruction : cloned_fused_computation->instructions()) { - LayoutUtil::SetToDefaultLayout(instruction->mutable_shape()); + if (!LayoutUtil::HasLayout(instruction->shape())) { + LayoutUtil::SetToDefaultLayout(instruction->mutable_shape()); + } } auto readded_computation = empty_hlo_module.AddEntryComputation(std::move(cloned_fused_computation)); @@ -1098,6 +1258,8 @@ Status HloEvaluator::HandleFusion(HloInstruction* fusion) { } HloEvaluator embedded_evaluator; + embedded_evaluator.set_dynamic_dimension_inference( + dynamic_dimension_inference_); Literal result = embedded_evaluator.Evaluate(*readded_computation, arg_literals) .ConsumeValueOrDie(); @@ -1117,6 +1279,8 @@ Status HloEvaluator::HandleConditional(HloInstruction* conditional) { auto* false_computation = conditional->false_computation(); HloEvaluator embedded_evaluator; + embedded_evaluator.set_dynamic_dimension_inference( + dynamic_dimension_inference_); Literal result; if (pred.Get({})) { result = @@ -1171,7 +1335,10 @@ Status HloEvaluator::HandleWhile(HloInstruction* while_hlo) { bool keep_going = true; int64 iteration_count = 0; HloEvaluator cond_evaluator(max_loop_iterations_); + cond_evaluator.set_dynamic_dimension_inference(dynamic_dimension_inference_); HloEvaluator loop_body_evaluator(max_loop_iterations_); + loop_body_evaluator.set_dynamic_dimension_inference( + dynamic_dimension_inference_); while (keep_going) { if (max_loop_iterations_ >= 0 && iteration_count++ > max_loop_iterations_) { return InvalidArgument("Loop %s exceeded loop iteration limit (%d).", @@ -1193,169 +1360,209 @@ Status HloEvaluator::HandleWhile(HloInstruction* while_hlo) { return Status::OK(); } -// Key-value sort is a special snowflake: it's templated on two different -// element types, one for the keys, and one for the values. Jump through some -// hoops to make this work. namespace { -template -StatusOr EvaluateSortInternal(HloInstruction* sort, - const Literal& keys_literal, - const Literal& values_literal) { - auto rank = keys_literal.shape().rank(); - TF_RET_CHECK( - ShapeUtil::SameDimensions(keys_literal.shape(), values_literal.shape())) - << "Sort keys and values must have the same dimensions"; - TF_RET_CHECK(sort->operand_count() >= 2) << "Expected key-value sort"; - // We need to sort an array of keys and an array of values, where the - // sorted order of the values is determined by the keys. The simplest(?) - // way to do this is to go to an array-of-pairs representation, sort the - // array using the keys, and then go back to pair-of-arrays. - VLOG(3) << "HandleSort keys_literal: " << keys_literal.ToString(); - VLOG(3) << "HandleSort values_literal: " << values_literal.ToString(); - - if (rank == 0) { - // Nothing to sort. - return LiteralUtil::MakeTuple({&keys_literal, &values_literal}); +template +Literal ExtractLiteralFromIndexPositions(const Literal& from, + absl::Span indices, + bool extract_as_scalar) { + if (extract_as_scalar) { + return LiteralUtil::CreateR0(from.Get({indices[0]})); + } + // We use a InlinedVector here because we need to convert it to an + // absl::Span later, and this would not work with std::vector. + absl::InlinedVector values; + for (int64 index : indices) { + values.push_back(from.Get({index})); + } + return LiteralUtil::CreateR1(values); +} + +StatusOr ExtractFromIndexPositions(const Literal& from, + absl::Span indices, + bool extract_as_scalar = false) { + if (extract_as_scalar) { + CHECK_EQ(indices.size(), 1); + } + PrimitiveType type = from.shape().element_type(); + switch (type) { + case PRED: { + return ExtractLiteralFromIndexPositions(from, indices, + extract_as_scalar); + } + case U8: { + return ExtractLiteralFromIndexPositions(from, indices, + extract_as_scalar); + } + case S8: { + return ExtractLiteralFromIndexPositions(from, indices, + extract_as_scalar); + } + case BF16: { + return ExtractLiteralFromIndexPositions(from, indices, + extract_as_scalar); + } + case F16: { + return ExtractLiteralFromIndexPositions(from, indices, + extract_as_scalar); + } + case U16: { + return ExtractLiteralFromIndexPositions(from, indices, + extract_as_scalar); + } + case S16: { + return ExtractLiteralFromIndexPositions(from, indices, + extract_as_scalar); + } + case F32: { + return ExtractLiteralFromIndexPositions(from, indices, + extract_as_scalar); + } + case U32: { + return ExtractLiteralFromIndexPositions(from, indices, + extract_as_scalar); + } + case S32: { + return ExtractLiteralFromIndexPositions(from, indices, + extract_as_scalar); + } + case F64: { + return ExtractLiteralFromIndexPositions(from, indices, + extract_as_scalar); + } + case U64: { + return ExtractLiteralFromIndexPositions(from, indices, + extract_as_scalar); + } + case S64: { + return ExtractLiteralFromIndexPositions(from, indices, + extract_as_scalar); + } + default: + return InvalidArgument("Unsupported type for Sort: %s", + PrimitiveType_Name(type)); + } +} +} // namespace + +Status HloEvaluator::HandleSort(HloInstruction* sort) { + TF_RET_CHECK(sort->operand_count() >= 1) + << "Expected at least 1 operand for sort"; + for (int64 i = 1; i < sort->operand_count(); ++i) { + TF_RET_CHECK(ShapeUtil::SameDimensions(sort->operand(0)->shape(), + sort->operand(i)->shape())) + << "All Sort operands must have the same dimensions"; } - Literal keys_result_literal(keys_literal.shape()); - Literal values_result_literal(values_literal.shape()); + if (VLOG_IS_ON(3)) { + for (int64 i = 0; i < sort->operand_count(); ++i) { + VLOG(3) << "HandleSort operand " << i << " literal: " + << GetEvaluatedLiteralFor(sort->operand(i)).ToString(); + } + } + Shape key_shape = sort->operand(0)->shape(); + auto rank = key_shape.rank(); + std::vector result_literals; + result_literals.reserve(sort->operand_count()); + for (int64 i = 0; i < sort->operand_count(); ++i) { + result_literals.emplace_back(sort->operand(i)->shape()); + } std::vector zero_base(rank, 0); std::vector increment(rank, 1); int64 sort_dim = sort->dimensions(0); - int64 sort_dim_elements = keys_literal.shape().dimensions(sort_dim); + int64 sort_dim_elements = key_shape.dimensions(sort_dim); increment[sort_dim] = sort_dim_elements; + HloEvaluator embedded_evaluator(max_loop_iterations_); // Iterate through each dimension except 'sort_dim'. TF_RETURN_IF_ERROR(ShapeUtil::ForEachIndexWithStatus( - keys_literal.shape(), zero_base, - AsInt64Slice(keys_literal.shape().dimensions()), increment, + key_shape, zero_base, AsInt64Slice(key_shape.dimensions()), increment, [&](absl::Span indices) -> StatusOr { - // Extract a slice from the keys and values literals that correspond to + // Extract a slice from each operand literal that corresponds to // exactly the row in dimension 'sort_dim'. std::vector limit_indices(indices.begin(), indices.end()); - std::for_each(limit_indices.begin(), limit_indices.end(), - [](int64& index) { ++index; }); + absl::c_for_each(limit_indices, [](int64& index) { ++index; }); limit_indices[sort_dim] = sort_dim_elements; - TF_ASSIGN_OR_RETURN(auto keys_to_sort, - keys_literal.Slice(indices, limit_indices) - .Reshape({sort_dim_elements})); - const auto& keys_data = keys_to_sort.data(); - TF_ASSIGN_OR_RETURN(auto values_to_sort, - values_literal.Slice(indices, limit_indices) - .Reshape({sort_dim_elements})); - const auto& values_data = values_to_sort.data(); - using kv_pair = std::pair; - std::vector key_value_vector; - key_value_vector.reserve(keys_data.size()); - for (int i = 0; i < keys_data.size(); ++i) { - key_value_vector.push_back( - std::make_pair(keys_data[i], values_data[i])); + std::vector literals_to_sort; + literals_to_sort.reserve(sort->operand_count()); + for (int64 i = 0; i < sort->operand_count(); ++i) { + TF_ASSIGN_OR_RETURN(auto literal_to_sort, + GetEvaluatedLiteralFor(sort->operand(i)) + .Slice(indices, limit_indices) + .Reshape({sort_dim_elements})); + literals_to_sort.push_back(std::move(literal_to_sort)); } - std::stable_sort(key_value_vector.begin(), key_value_vector.end(), - [](const kv_pair& a, const kv_pair& b) { - return SafeLess(a.first, b.first); - }); - std::vector result_keys; - // We use a InlinedVector here because we need to convert it to an - // absl::Span later, and this would not work with std::vector. - absl::InlinedVector result_values; - for (const auto& key_value : key_value_vector) { - result_keys.push_back(key_value.first); - result_values.push_back(key_value.second); + std::vector indices_to_sort(sort_dim_elements); + std::iota(indices_to_sort.begin(), indices_to_sort.end(), 0); + Status compare_status = Status::OK(); + std::stable_sort( + indices_to_sort.begin(), indices_to_sort.end(), + [sort, &compare_status, &embedded_evaluator, &literals_to_sort]( + int64 a, int64 b) { + std::vector literals; + literals.reserve(2 * sort->operand_count()); + for (int64 i = 0; i < sort->operand_count(); ++i) { + auto lhs = ExtractFromIndexPositions( + literals_to_sort[i], {a}, /*extract_as_scalar=*/true); + if (!lhs.ok()) { + compare_status = lhs.status(); + return false; + } + literals.push_back(std::move(lhs.ValueOrDie())); + auto rhs = ExtractFromIndexPositions( + literals_to_sort[i], {b}, /*extract_as_scalar=*/true); + if (!rhs.ok()) { + compare_status = rhs.status(); + return false; + } + literals.push_back(std::move(rhs.ValueOrDie())); + } + std::vector literal_ptrs; + absl::c_transform( + literals, std::back_inserter(literal_ptrs), + [](const Literal& literal) { return &literal; }); + + auto computed_result = + embedded_evaluator.Evaluate(*sort->to_apply(), literal_ptrs); + // Clear visit states so that we can use the evaluator again + // on the same computation. + embedded_evaluator.ResetVisitStates(); + if (!computed_result.ok()) { + compare_status = computed_result.status(); + return false; + } + return computed_result.ValueOrDie().Get({}); + }); + if (!compare_status.ok()) { + return compare_status; } - Literal sorted_keys(ShapeUtil::MakeShape( - keys_literal.shape().element_type(), {sort_dim_elements})); - sorted_keys.PopulateR1(absl::Span(result_keys)); - Literal sorted_values(ShapeUtil::MakeShape( - values_literal.shape().element_type(), {sort_dim_elements})); - sorted_values.PopulateR1(absl::Span(result_values)); std::vector slice_dimensions(rank, 1); slice_dimensions[sort_dim] = sort_dim_elements; std::vector start_indices(rank, 0); - TF_ASSIGN_OR_RETURN(auto sorted_keys_reshaped, - sorted_keys.Reshape(slice_dimensions)); - TF_RETURN_IF_ERROR(keys_result_literal.CopySliceFrom( - sorted_keys_reshaped, start_indices, indices, slice_dimensions)); - TF_ASSIGN_OR_RETURN(auto sorted_values_reshaped, - sorted_values.Reshape(slice_dimensions)); - TF_RETURN_IF_ERROR(values_result_literal.CopySliceFrom( - sorted_values_reshaped, start_indices, indices, slice_dimensions)); + for (int64 i = 0; i < sort->operand_count(); ++i) { + TF_ASSIGN_OR_RETURN( + Literal sorted_literal, + ExtractFromIndexPositions(literals_to_sort[i], indices_to_sort)); + TF_ASSIGN_OR_RETURN(auto sorted_literal_reshaped, + sorted_literal.Reshape(slice_dimensions)); + TF_RETURN_IF_ERROR(result_literals[i].CopySliceFrom( + sorted_literal_reshaped, start_indices, indices, + slice_dimensions)); + } return true; })); - Literal result_tuple; - result_tuple = - LiteralUtil::MakeTuple({&keys_result_literal, &values_result_literal}); - VLOG(3) << "HandleSort result_tuple: " << result_tuple.ToString(); - return std::move(result_tuple); -} - -template -StatusOr EvaluateSortCurried(HloInstruction* sort, - const Literal& keys_literal, - const Literal& values_literal) { - switch (values_literal.shape().element_type()) { - case PRED: - return EvaluateSortInternal(sort, keys_literal, - values_literal); - case F32: - return EvaluateSortInternal(sort, keys_literal, - values_literal); - case U32: - return EvaluateSortInternal(sort, keys_literal, - values_literal); - case S32: - return EvaluateSortInternal(sort, keys_literal, - values_literal); - case BF16: - return EvaluateSortInternal(sort, keys_literal, - values_literal); - default: - return InvalidArgument("Unsupported type for Sort"); - } -} - -StatusOr EvaluateSort(HloInstruction* sort, - const Literal& keys_literal, - const Literal& values_literal) { - switch (sort->operand(0)->shape().element_type()) { - case F32: - return EvaluateSortCurried(sort, keys_literal, values_literal); - case U32: - return EvaluateSortCurried(sort, keys_literal, values_literal); - case S32: - return EvaluateSortCurried(sort, keys_literal, values_literal); - case BF16: - return EvaluateSortCurried(sort, keys_literal, values_literal); - default: - return InvalidArgument("Unsupported type for Sort"); - } -} -} // namespace - -Status HloEvaluator::HandleSort(HloInstruction* sort) { - if (!sort->shape().IsTuple()) { - return DefaultAction(sort); + if (sort->operand_count() == 1) { + evaluated_[sort] = std::move(result_literals[0]); } else { - // This is a really stupid work-around for the fact it's hard to support a - // multi-value sort directly, due to the fact we need to template the - // evaluation function on all of the value types. - std::vector sort_results_backing; - for (int64 i = 0; i < sort->operand_count(); ++i) { - auto result = EvaluateSort(sort, GetEvaluatedLiteralFor(sort->operand(0)), - GetEvaluatedLiteralFor(sort->operand(i))); - if (!result.ok()) { - return result.status(); - } - sort_results_backing.push_back( - std::move(result.ValueOrDie().DecomposeTuple()[1])); - } - std::vector sort_results; - absl::c_transform(sort_results_backing, std::back_inserter(sort_results), + std::vector literal_ptrs; + absl::c_transform(result_literals, std::back_inserter(literal_ptrs), [](const Literal& literal) { return &literal; }); - evaluated_[sort] = LiteralUtil::MakeTuple(sort_results); - return Status::OK(); + + Literal result_tuple = LiteralUtil::MakeTuple(literal_ptrs); + VLOG(3) << "HandleSort result_tuple: " << result_tuple.ToString(); + + evaluated_[sort] = std::move(result_tuple); } + return Status::OK(); } Status HloEvaluator::HandleReduce(HloInstruction* reduce) { @@ -1374,6 +1581,27 @@ Status HloEvaluator::HandleReduce(HloInstruction* reduce) { } } +Status HloEvaluator::HandleCustomCall(HloInstruction* custom_call) { + if (!custom_call_handler_) { + // No handler is registered; this means custom-calls are not allowed. + return DefaultAction(custom_call); + } + + // Evaluate input operands so the handler has access to the operand data. + std::vector operands; + operands.reserve(custom_call->operand_count()); + for (const HloInstruction* operand : custom_call->operands()) { + operands.push_back(&GetEvaluatedLiteralFor(operand)); + } + + // Synchronously issue the handler to populate the instruction output literal. + TF_ASSIGN_OR_RETURN( + auto output, custom_call_handler_(custom_call, absl::MakeSpan(operands))); + + evaluated_[custom_call] = std::move(output); + return Status::OK(); +} + Status HloEvaluator::Preprocess(HloInstruction* hlo) { VLOG(2) << "About to visit HLO: " << hlo->ToString(); return ShapeUtil::ValidateShape(hlo->shape()); diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.h b/tensorflow/compiler/xla/service/hlo_evaluator.h index 44d5e5c030fd43c2510b0c41cd43e97f0761eda5..357975a131d0c7e63c06e96852468b43d97a37f2 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator.h +++ b/tensorflow/compiler/xla/service/hlo_evaluator.h @@ -16,13 +16,16 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_H_ +#include #include #include "absl/container/node_hash_map.h" #include "absl/memory/memory.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/array2d.h" +#include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" +#include "tensorflow/compiler/xla/service/dynamic_dimension_inference.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" @@ -123,9 +126,31 @@ class HloEvaluator : public DfsHloVisitorWithDefault { const PrecisionConfig& precision_config, const Literal& lhs, const Literal& rhs); + void set_dynamic_dimension_inference( + DynamicDimensionInference* dynamic_dimension_inference) { + dynamic_dimension_inference_ = dynamic_dimension_inference; + } + // Enable the fast path for certain operations like dot or convolution. void set_use_fast_path(bool value) { use_fast_path_ = value; } + // Handles evaluation of a custom-call op. + // Operand literals are provided in |operands| and implementations must + // populate |output| before returning. + using CustomCallHandler = std::function( + HloInstruction* custom_call, absl::Span operands)>; + + // Sets a handler that is called during evaluation for custom-call ops. + // If no handler is defined the default error behavior will occur. The handler + // will be provided evaluated literals for all operands and is expected to + // return an output literal of the appropriate shape. + void set_custom_call_handler( + std::function(HloInstruction* custom_call, + absl::Span operands)> + handler) { + custom_call_handler_ = std::move(handler); + } + // Returns the result of a matrix multiply `lhs x rhs`. static std::unique_ptr> MatmulArray2D( const Array2D& lhs, const Array2D& rhs); @@ -161,6 +186,8 @@ class HloEvaluator : public DfsHloVisitorWithDefault { // Status HandleBitcast(HloInstruction* bitcast) override; + Status HandleGetDimensionSize(HloInstruction* get_dimension_size) override; + Status HandleParameter(HloInstruction* parameter) override; Status HandleConstant(HloInstruction* constant) override; @@ -211,14 +238,47 @@ class HloEvaluator : public DfsHloVisitorWithDefault { Status HandleReduce(HloInstruction* reduce) override; + Status HandleCustomCall(HloInstruction* custom_call) override; + + // Unsupported HLOs, note some of them (such as BatchNorm*) are typically + // expanded in a semantic-preserving way into other HLOs by adding exanpsion + // HLO pass to the HLO optimization pass during compilation, which can then be + // handled by the evaluator. + Status HandleBatchNormGrad(HloInstruction* batch_norm_grad) override { + return Unimplemented("BatchNormGrad HLO is unsupported by the evaluator."); + }; + Status HandleBatchNormInference( + HloInstruction* batch_norm_inference) override { + return Unimplemented( + "BatchNormInference HLO is unsupported by the evaluator."); + }; + Status HandleBatchNormTraining(HloInstruction* batch_norm_training) override { + return Unimplemented( + "BatchNormTraining HLO is unsupported by the evaluator."); + }; + Status HandleInfeed(HloInstruction* infeed) override { + return Unimplemented("Infeed HLO is unsupported by the evaluator."); + }; + Status HandleOutfeed(HloInstruction* outfeed) override { + return Unimplemented("Outfeed HLO is unsupported by the evaluator."); + }; + // Returns the already-evaluated literal result for the instruction. + // // A Constant instruction is considered evaluated and its literal will be // returned directly without looking up the cache. + // + // Similarly, a Parameter instruction is considered evaluated and its literal + // is looked up in arg_literals. + // // Crash with log if the given instruction has not been evaluated previously. const Literal& GetEvaluatedLiteralFor(const HloInstruction* hlo) { if (hlo->IsConstant()) { return hlo->literal(); } + if (hlo->opcode() == HloOpcode::kParameter) { + return *arg_literals_.at(hlo->parameter_number()); + } auto it = evaluated_.find(hlo); CHECK(it != evaluated_.end()) << "could not find evaluated value for: " << hlo->ToString(); @@ -226,12 +286,18 @@ class HloEvaluator : public DfsHloVisitorWithDefault { } // Tracks the HLO instruction and its evaluated literal result. + // + // Parameters and constants aren't stored here, see implementation of + // GetEvaluatedLiteralFor. + // // TODO(b/35950897): have better memory management here to free instructions // that are no longer a parent for any other subsequent instruction in // post-orderring. + // // Must be cleared for each evaluation. - // Storing Literal in place require the container to have pointer stability so - // we cannot use flat_hash_map any more. + // + // Storing Literal in place requires the container to have pointer stability + // so we cannot use flat_hash_map any more. absl::node_hash_map evaluated_; // Use fast path that uses eigen in the evaluator. @@ -265,7 +331,21 @@ class HloEvaluator : public DfsHloVisitorWithDefault { std::vector arg_literals_; // Max loop iterations to execute with no maximum if negative. - int64 max_loop_iterations_; + int64 max_loop_iterations_ = 0; + + // Module-level seed handle. + uint64 seed_ = 0; + // RNG engine. + std::minstd_rand0 engine_; + + // DynamicDimensionInference is used to evaluate GetDimensionSize, which + // returns the dynamic dimension size of its operand. + DynamicDimensionInference* dynamic_dimension_inference_ = nullptr; + + // Optional handler for custom_call ops. + std::function(HloInstruction* custom_call, + absl::Span operands)> + custom_call_handler_; TF_DISALLOW_COPY_AND_ASSIGN(HloEvaluator); }; diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_test.cc b/tensorflow/compiler/xla/service/hlo_evaluator_test.cc index 0298da21747f85372e120f3b21e852480531560a..fb8cd299cef06d549130cd56dd2c430c4c1a0387 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator_test.cc +++ b/tensorflow/compiler/xla/service/hlo_evaluator_test.cc @@ -62,8 +62,7 @@ class HloEvaluatorTest : public HloTestBase { if (use_bfloat16_) { HloElementTypeConverter(F32, BF16).Run(m_.get()).ValueOrDie(); } - return HloEvaluator() - .Evaluate(*m_->entry_computation(), arg_literals) + return evaluator_.Evaluate(*m_->entry_computation(), arg_literals) .ConsumeValueOrDie(); } @@ -75,8 +74,7 @@ class HloEvaluatorTest : public HloTestBase { if (use_bfloat16_) { HloElementTypeConverter(F32, BF16).Run(m_.get()).ValueOrDie(); } - return HloEvaluator() - .Evaluate(*module->entry_computation(), arg_literals) + return evaluator_.Evaluate(*module->entry_computation(), arg_literals) .ConsumeValueOrDie(); } @@ -115,6 +113,7 @@ class HloEvaluatorTest : public HloTestBase { protected: explicit HloEvaluatorTest(bool use_bfloat16) : use_bfloat16_(use_bfloat16) {} + HloEvaluator evaluator_; const bool use_bfloat16_; std::unique_ptr m_ = CreateNewVerifiedModule(); @@ -127,8 +126,8 @@ class HloEvaluatorBf16Test : public ::testing::WithParamInterface, HloEvaluatorBf16Test() : HloEvaluatorTest(/*use_bfloat16=*/GetParam()) {} }; -INSTANTIATE_TEST_CASE_P(HloEvaluatorTest_Instantiation, HloEvaluatorBf16Test, - ::testing::ValuesIn(use_bf16_params)); +INSTANTIATE_TEST_SUITE_P(HloEvaluatorTest_Instantiation, HloEvaluatorBf16Test, + ::testing::ValuesIn(use_bf16_params)); // Verifies that HloEvaluator evaluates a HLO instruction that performs clamp // with 3 operands. @@ -309,6 +308,19 @@ TEST_F(HloEvaluatorTest, DoesNotR2) { {0, std::numeric_limits::min()}}); TestUnaryOp(HloOpcode::kNot, std::move(expected), std::move(operand)); } + +TEST_F(HloEvaluatorTest, DoesRealC128) { + auto x = LiteralUtil::CreateR1({{1, 0}, {-100, 4}}); + auto expected_real = LiteralUtil::CreateR1({1, -100}); + TestUnaryOp(HloOpcode::kReal, std::move(expected_real), std::move(x)); +} + +TEST_F(HloEvaluatorTest, DoesImagC128) { + auto x = LiteralUtil::CreateR1({{1, 0}, {-100, 4}}); + auto expected_imag = LiteralUtil::CreateR1({0, 4}); + TestUnaryOp(HloOpcode::kImag, std::move(expected_imag), std::move(x)); +} + // Verifies that HloEvaluator evaluates a HLO Computation with non-parameter nor // constant operands. TEST_F(HloEvaluatorTest, DoesTraverseInstructions) { @@ -1740,12 +1752,14 @@ TEST_P(HloEvaluatorBf16Test, DynamicSlice) { HloInstruction* operand = b.AddInstruction( HloInstruction::CreateConstant(std::move(operand_literal))); - auto start_indices = b.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({0, 1}))); + auto zero = b.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); + auto one = b.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(1))); Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - b.AddInstruction(HloInstruction::CreateDynamicSlice(shape, operand, - start_indices, {2, 3})); + b.AddInstruction( + HloInstruction::CreateDynamicSlice(shape, operand, {zero, one}, {2, 3})); m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1776,12 +1790,14 @@ TEST_P(HloEvaluatorBf16Test, DynamicSliceModSlice) { HloInstruction* operand = b.AddInstruction( HloInstruction::CreateConstant(std::move(operand_literal))); - auto start_indices = b.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({2, 1}))); + auto two = b.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(2))); + auto one = b.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(1))); Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - b.AddInstruction(HloInstruction::CreateDynamicSlice(shape, operand, - start_indices, {2, 3})); + b.AddInstruction( + HloInstruction::CreateDynamicSlice(shape, operand, {two, one}, {2, 3})); m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1810,15 +1826,17 @@ TEST_P(HloEvaluatorBf16Test, DynamicSliceUpdate) { HloInstruction* operand = b.AddInstruction( HloInstruction::CreateConstant(std::move(operand_literal))); - auto start_indices = b.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({0, 1}))); + auto zero = b.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); + auto one = b.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(1))); auto update = b.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR2({{-2.0, -3.0}, {-6.0, -7.0}}))); Shape shape = ShapeUtil::MakeShape(F64, {2, 3}); b.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - shape, operand, update, start_indices)); + shape, operand, update, {zero, one})); m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -2874,6 +2892,39 @@ ENTRY main { static_cast(pow31 * pow31)); } +TEST_F(HloEvaluatorTest, GetDimensionSize) { + constexpr absl::string_view hlo_text = R"( +HloModule Test + +ENTRY main { + size = u32[] parameter(0) + + data = s32[4] parameter(1) + + sum = s32[4] add(data, data) + + ROOT dynamic_size = u32[] get-dimension-size(sum), dimensions={0} +} +)"; + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); + + // Set up dynamic parameter binding. + TF_CHECK_OK(m_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{0, {}}, + DynamicParameterBinding::DynamicDimension{1, {}, 0})); + + TF_ASSERT_OK_AND_ASSIGN(DynamicDimensionInference dynamic_dimension_inference, + DynamicDimensionInference::Run(m_.get())); + + evaluator_.set_dynamic_dimension_inference(&dynamic_dimension_inference); + Literal size_arg = LiteralUtil::CreateR0(3); + Literal data_arg = LiteralUtil::CreateR1({1, 2, 3, 4}); + + Literal actual = Evaluate({&size_arg, &data_arg}); + + EXPECT_EQ(actual.GetFirstElement(), static_cast(3)); +} + // Check that we get a useful error if we pass inputs of the wrong shape. TEST_F(HloEvaluatorTest, EvaluateWithWrongInputShapes) { constexpr absl::string_view hlo_text = R"( @@ -2924,5 +2975,183 @@ ENTRY main { "Expected 1 argument, but got 2."); } +TEST_F(HloEvaluatorTest, PreserveFusionInputLayout) { + constexpr absl::string_view hlo_text = R"( + HloModule FusionInputLayout + + fused_computation { + param_0 = f32[20,20]{0,1} parameter(0) + ROOT bitcast = f32[20,20]{1,0} bitcast(param_0) + } + + ENTRY kernel_entry { + parameter.0 = f32[20,20]{0,1} parameter(0) + ROOT fusion = f32[20,20]{1,0} fusion(parameter.0), + kind=kLoop, calls=fused_computation + })"; + + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); + auto args = MakeFakeArguments(m_.get()).ConsumeValueOrDie(); + Literal actual = Evaluate({&args[0]}); + EXPECT_TRUE(absl::c_equal(args[0].data(), actual.data())); +} + +TEST_F(HloEvaluatorTest, PreserveFusionOutputLayout) { + constexpr absl::string_view hlo_text = R"( + HloModule FusionOutputLayout + + fused_computation { + param_0 = f32[20,20]{1,0} parameter(0) + ROOT bitcast = f32[20,20]{0,1} bitcast(param_0) + } + + ENTRY kernel_entry { + parameter.0 = f32[20,20]{1,0} parameter(0) + ROOT fusion = f32[20,20]{0,1} fusion(parameter.0), + kind=kLoop, calls=fused_computation + })"; + + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); + auto args = MakeFakeArguments(m_.get()).ConsumeValueOrDie(); + Literal actual = Evaluate({&args[0]}); + EXPECT_TRUE(absl::c_equal(args[0].data(), actual.data())); +} + +TEST_F(HloEvaluatorTest, PreserveMOFusionOutputLayout) { + constexpr absl::string_view hlo_text = R"( + HloModule MOFusionOutputLayout + + fused_computation { + param_0 = f32[20,20]{1,0} parameter(0) + bitcast = f32[20,20]{0,1} bitcast(param_0) + ROOT tuple = (f32[20,20]{0,1}) tuple(bitcast) + } + + ENTRY kernel_entry { + parameter.0 = f32[20,20]{1,0} parameter(0) + ROOT fusion = (f32[20,20]{0,1}) fusion(parameter.0), + kind=kLoop, calls=fused_computation + })"; + + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); + auto args = MakeFakeArguments(m_.get()).ConsumeValueOrDie(); + Literal actual_tuple = Evaluate({&args[0]}); + std::vector actual_literals = actual_tuple.DecomposeTuple(); + EXPECT_TRUE( + absl::c_equal(args[0].data(), actual_literals[0].data())); +} + +// Tests that custom_calls fail to evaluate when no handler is specified. +TEST_F(HloEvaluatorTest, EvaluateCustomCall_NoHandler) { + constexpr absl::string_view hlo_text = R"( + HloModule EvaluateCustomCall_NoHandler + ENTRY kernel_entry { + parameter.0 = u32[2,2]{1,0} parameter(0) + ROOT test_root = (u32[2,2]{1,0}) custom-call(parameter.0), + custom_call_target="_my_custom_call" + } + )"; + + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); + auto args = MakeFakeArguments(m_.get()).ConsumeValueOrDie(); + EXPECT_EQ(HloEvaluator().Evaluate(*m_, {&args[0]}).status().code(), + ::tensorflow::error::UNIMPLEMENTED); +} + +// Tests when a custom_call handler returns an error. +TEST_F(HloEvaluatorTest, EvaluateCustomCall_HandlerError) { + constexpr absl::string_view hlo_text = R"( + HloModule EvaluateCustomCall_HandlerError + ENTRY kernel_entry { + parameter.0 = u32[2,2]{1,0} parameter(0) + ROOT test_root = (u32[2,2]{1,0}) custom-call(parameter.0), + custom_call_target="_my_custom_call" + } + )"; + + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); + auto args = MakeFakeArguments(m_.get()).ConsumeValueOrDie(); + HloEvaluator evaluator; + evaluator.set_custom_call_handler( + [](HloInstruction* custom_call, absl::Span operands) { + return InternalError("Test error"); + }); + EXPECT_EQ(evaluator.Evaluate(*m_, {&args[0]}).status().code(), + ::tensorflow::error::INTERNAL); +} + +// Tests the custom_call handler on calls with many inputs. +// We sum the operands so that we can verify the operand and output literals +// are properly mapped for access. +TEST_F(HloEvaluatorTest, EvaluateCustomCall_ManyInputs) { + constexpr absl::string_view hlo_text = R"( + HloModule EvaluateCustomCall_ManyInputs + ENTRY kernel_entry { + parameter.0 = u32[1]{0} parameter(0) + parameter.1 = u32[1]{0} parameter(1) + ROOT test_root = u32[1]{0} custom-call(parameter.0, parameter.1), + custom_call_target="_my_custom_call" + } + )"; + + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); + auto args = MakeFakeArguments(m_.get()).ConsumeValueOrDie(); + HloEvaluator evaluator; + evaluator.set_custom_call_handler( + [](HloInstruction* custom_call, absl::Span operands) { + EXPECT_EQ(HloOpcode::kCustomCall, custom_call->opcode()); + EXPECT_EQ("_my_custom_call", custom_call->custom_call_target()); + EXPECT_EQ(2, custom_call->operand_count()); + EXPECT_EQ(2, operands.size()); + auto output = Literal::CreateFromShape(custom_call->shape()); + auto operand0_data = operands[0]->data(); + auto operand1_data = operands[1]->data(); + auto output_data = output.data(); + output_data[0] = operand0_data[0] + operand1_data[0]; + return output; + }); + TF_ASSERT_OK_AND_ASSIGN( + Literal actual_literal, + evaluator.Evaluate(*m_->entry_computation(), {&args[0], &args[1]})); + auto arg0_data = args[0].data(); + auto arg1_data = args[1].data(); + std::vector expected_data = {arg0_data[0] + arg1_data[0]}; + EXPECT_TRUE(absl::c_equal(expected_data, actual_literal.data())); +} + +TEST_F(HloEvaluatorTest, IsFiniteF16) { + constexpr absl::string_view hlo_text = R"( + HloModule test + + ENTRY IsFiniteTest { + c = f16[6] constant({nan, 7, nan, -1, inf, -inf}) + ROOT is-finite = pred[6] is-finite(c) + })"; + + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); + TF_ASSERT_OK_AND_ASSIGN( + Literal actual_literal, + HloEvaluator().Evaluate(*m_->entry_computation(), {})); + EXPECT_THAT(actual_literal.data(), + ::testing::ElementsAre(false, true, false, true, false, false)); +} + +TEST_F(HloEvaluatorTest, IsFiniteBf16) { + constexpr absl::string_view hlo_text = R"( + HloModule test + + ENTRY IsFiniteTest { + c = bf16[6] constant({nan, 7, nan, -1, inf, -inf}) + ROOT is-finite = pred[6] is-finite(c) + })"; + + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); + TF_ASSERT_OK_AND_ASSIGN( + Literal actual_literal, + HloEvaluator().Evaluate(*m_->entry_computation(), {})); + EXPECT_THAT(actual_literal.data(), + ::testing::ElementsAre(false, true, false, true, false, false)); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h index 6f3208be13b7ae27de6d11b165f4dc921cb1d633..8def61dc63db2c55f926a04fe097988af4417c1a 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h +++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h @@ -22,6 +22,7 @@ limitations under the License. #include "absl/base/casts.h" #include "absl/container/inlined_vector.h" #include "absl/memory/memory.h" +#include "absl/meta/type_traits.h" #include "absl/types/optional.h" #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/literal_util.h" @@ -39,49 +40,8 @@ namespace xla { // Anyway this is relatively safe as-is because hlo_evaluator_typed_visitor.h is // a "private" header that's not exposed outside of hlo_evaluator.cc. template -using is_complex_t = std::is_same; -template -using is_complex64_t = std::is_same; - -// It's UB to use std::sort with std::less, because of NaNs. Define -// "safe" less functions which are actually strict weak orders. -NaN and NaN -// should appear at the beginning and end of the ordering, and -0.0 should -// appear before 0.0. -template < - typename NativeT, - typename std::enable_if::value>::type* = nullptr> -bool SafeLess(const NativeT& a, const NativeT& b) { - return a < b; -} - -template ::value>::type* = nullptr> -bool SafeLess(const NativeT& a, const NativeT& b) { - bool lhs_is_negative = std::signbit(a); - bool rhs_is_negative = std::signbit(b); - // If the signs are different, we can just compare the signs. - if (lhs_is_negative != rhs_is_negative) { - return lhs_is_negative && !rhs_is_negative; - } - bool lhs_nan = std::isnan(a); - bool rhs_nan = std::isnan(b); - // Exactly one number is nan? - if (lhs_nan != rhs_nan) { - if (lhs_nan) { - return lhs_is_negative; - } - return !rhs_is_negative; - } - return a < b; -} - -template ::value || - std::is_same::value>::type* = nullptr> -bool SafeLess(const NativeT& a, const NativeT& b) { - return SafeLess(static_cast(a), static_cast(b)); -} +using is_complex_t = + absl::disjunction, std::is_same>; // ToArithmeticSafeType(T t): // - converts `t` to the bitwise-equivalent `unsigned T` if T is a signed @@ -212,7 +172,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { template < typename NativeT, - typename std::enable_if::value>::type* = nullptr> + typename std::enable_if::value>::type* = nullptr> Status HandleAbs(HloInstruction* abs) { const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(abs->operand(0)); @@ -231,6 +191,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { // specifying the ElementwiseT explicitly as C64 is needed below. if (abs->operand(0)->shape().element_type() == C64) { return HandleAbs(abs); + } else if (abs->operand(0)->shape().element_type() == C128) { + return HandleAbs(abs); } return HandleAbs(abs); } @@ -460,9 +422,9 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { return HandleNegate(negate); } - template < - typename NativeT, - typename std::enable_if::value>::type* = nullptr> + template ::value>::type* = + nullptr> Status HandleSign(HloInstruction* sign) { TF_ASSIGN_OR_RETURN(parent_->evaluated_[sign], ElementWiseUnaryOp(sign, [](ElementwiseT elem_operand) { @@ -472,6 +434,23 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { return Status::OK(); } + template ::value || + std::is_same::value || + std::is_floating_point::value>::type* = nullptr> + Status HandleSign(HloInstruction* sign) { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[sign], + ElementWiseUnaryOp(sign, [](ElementwiseT elem_operand) { + return std::isnan(elem_operand) + ? elem_operand + : std::copysign( + elem_operand != ElementwiseT(0), + elem_operand); + })); + return Status::OK(); + } + template < typename NativeT, typename std::enable_if::value>::type* = nullptr> @@ -673,11 +652,14 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { } Status HandlePower(HloInstruction* power) override { - TF_ASSIGN_OR_RETURN(parent_->evaluated_[power], - ElementWiseBinaryOp(power, [](ElementwiseT lhs_el, - ElementwiseT rhs_el) { - return std::pow(lhs_el, rhs_el); - })); + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[power], + ElementWiseBinaryOp( + power, [](ElementwiseT lhs_el, ElementwiseT rhs_el) { + return lhs_el == ElementwiseT(0) && rhs_el == ElementwiseT(0) + ? static_cast(1) + : std::pow(lhs_el, rhs_el); + })); return Status::OK(); } @@ -917,7 +899,11 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { Status HandleClamp(HloInstruction* clamp) { std::function clamp_op = [](ElementwiseT low, ElementwiseT value, ElementwiseT high) { - return std::fmin(high, std::fmax(value, low)); + if (std::isnan(low) || std::isnan(high)) { + return static_cast(NAN); + } + return static_cast( + std::fmin(high, std::fmax(value, low))); }; TF_ASSIGN_OR_RETURN( parent_->evaluated_[clamp], @@ -1036,15 +1022,13 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { auto lhs_literal_data = lhs_literal.data(); auto rhs_literal_data = rhs_literal.data(); - int64 feature_group_count = conv->feature_group_count(); - int64 batch_group_count = conv->batch_group_count(); + const int64 feature_group_count = conv->feature_group_count(); + const int64 batch_group_count = conv->batch_group_count(); - // The batch count > 1 case is unimplemented in the HLO evaluator so far. - TF_RET_CHECK(batch_group_count == 1); auto func = [&window_shape, &dnums, &lhs_shape, &rhs_shape, &window, &lhs_dim_multipliers, &rhs_dim_multipliers, lhs_literal_data, - rhs_literal_data, - feature_group_count](const absl::Span out_index) { + rhs_literal_data, feature_group_count, + batch_group_count](const absl::Span out_index) { // Dimension number applicable for input (lhs). const int64 input_batch_dim = dnums.input_batch_dimension(); const int64 input_z_dim = dnums.input_feature_dimension(); @@ -1057,6 +1041,12 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { const int64 input_z_size = ShapeUtil::GetDimension(lhs_shape, input_z_dim); + + const int64 input_batch_size = + ShapeUtil::GetDimension(lhs_shape, input_batch_dim); + + const int64 batch_group_size = input_batch_size / batch_group_count; + // The size of an input feature group. const int64 input_feature_group_size = input_z_size / feature_group_count; @@ -1072,11 +1062,15 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { const int64 feature_group_index = out_index[output_z_dim] / output_feature_group_size; + const int64 batch_group_index = out_index[output_z_dim]; + ElementwiseT result_val = static_cast(0); DimensionVector rhs_spatial_index(dnums.kernel_spatial_dimensions_size(), 0); // Convolve input feature with kernel. + // The mechanism indexes into the correct LHS (input) and RHS (kernel) + // locations and accumulates multiplications for a given output index. do { // Find corresponding spatial dimension index for input (lhs). int64 lhs_linear_spatial_index = 0; @@ -1129,11 +1123,24 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { feature_group_index * input_feature_group_size + rhs_iz; int64 lhs_linear_index = lhs_linear_spatial_index; + lhs_linear_index += out_index[output_batch_dim] * lhs_dim_multipliers[input_batch_dim]; + + // We are scraping only the diagonal elements in the resultant + // convolution output when batch_group_count is greater than 1, + // where 1 is the default. No scraping is done in that case. + // This approach works out automatically for 'groups' in batches + // with group_size > 1, because we already descend down the batch + // dimension for the 'output_batch_dim' above. + lhs_linear_index += + ((batch_group_index * batch_group_size) % input_batch_size) * + lhs_dim_multipliers[input_batch_dim]; + lhs_linear_index += iz * lhs_dim_multipliers[input_z_dim]; int64 rhs_linear_index = rhs_linear_spatial_index; + rhs_linear_index += out_index[output_z_dim] * rhs_dim_multipliers[kernel_output_z_dim]; rhs_linear_index += rhs_iz * rhs_dim_multipliers[kernel_input_z_dim]; @@ -1157,7 +1164,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { } Status HandleDot(HloInstruction* dot) override { - if (parent_->use_fast_path_) { + if (dot->dot_dimension_numbers().rhs_contracting_dimensions_size() == 1 && + parent_->use_fast_path_) { return HandleDot(dot); } return HandleDotSlowPath(dot); @@ -1319,12 +1327,17 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { static_cast(lhs_literal.Get(lhs_index)) * static_cast(rhs_literal.Get(rhs_index)); - for (int64 i = accumulate_index_sizes.size() - 1; i >= 0; --i) { - int64 value = ++accumulate_index[i]; - if (value != accumulate_index_sizes[i]) { - break; + // If there are no contracting dimension accumulate_index_sizes is + // empty, do not try to count down from -1 to 0 since it is and + // infinite loop. + if (!accumulate_index_sizes.empty()) { + for (int64 i = accumulate_index_sizes.size() - 1; i >= 0; --i) { + int64 value = ++accumulate_index[i]; + if (value != accumulate_index_sizes[i]) { + break; + } + accumulate_index[i] = 0; } - accumulate_index[i] = 0; } } @@ -1406,22 +1419,12 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { auto operand = dynamic_slice->operand(0); auto start_indices = dynamic_slice->operand(1); auto result_shape = dynamic_slice->shape(); - // TODO(b/118437727): Remove all of this nonsense. - // We may get an instruction without a parent module. In this case, assume - // scalar indices are not allowed. - bool allow_scalar_index = false; - if (dynamic_slice->GetModule() != nullptr) { - allow_scalar_index = dynamic_slice->GetModule() - ->config() - .debug_options() - .xla_allow_scalar_index_dynamic_ops(); - } TF_ASSIGN_OR_RETURN( auto inferred_return_shape, ShapeInference::InferDynamicSliceShape( operand->shape(), Cast(dynamic_slice)->index_shapes(), - dynamic_slice->dynamic_slice_sizes(), allow_scalar_index)); + dynamic_slice->dynamic_slice_sizes())); TF_RET_CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape)) << "return shape is set to: " << ShapeUtil::HumanString(result_shape) << " but is inferred to be: " @@ -1479,20 +1482,12 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { auto update = dynamic_update_slice->operand(1); auto start_indices = dynamic_update_slice->operand(2); auto result_shape = dynamic_update_slice->shape(); - bool allow_scalar_index = false; - if (dynamic_update_slice->GetModule() != nullptr) { - allow_scalar_index = dynamic_update_slice->GetModule() - ->config() - .debug_options() - .xla_allow_scalar_index_dynamic_ops(); - } TF_ASSIGN_OR_RETURN( auto inferred_return_shape, ShapeInference::InferDynamicUpdateSliceShape( operand->shape(), update->shape(), Cast(dynamic_update_slice) - ->index_shapes(), - allow_scalar_index)); + ->index_shapes())); TF_RET_CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape)) << "return shape is set to: " << ShapeUtil::HumanString(result_shape) << " but is inferred to be: " @@ -1630,6 +1625,10 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl(map)); break; } + case C128: { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl(map)); + break; + } default: LOG(FATAL) << "HandleMap: unhandled primitive type for " "input operand: " @@ -1640,74 +1639,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { return Status::OK(); } - template ::value && - !std::is_same::value>::type* = nullptr> - Status HandleSort(HloInstruction* sort) { - auto keys = sort->operand(0); - TF_RET_CHECK(sort->operand_count() == 1) - << "Typed visitor does not support key-value sort"; - - const Literal& keys_literal = parent_->GetEvaluatedLiteralFor(keys); - int64 sort_dim = sort->dimensions(0); - int64 sort_dim_elements = keys->shape().dimensions(sort_dim); - int64 rank = keys->shape().rank(); - if (rank == 0) { - // Nothing to sort. - parent_->evaluated_[sort] = keys_literal.Clone(); - return Status::OK(); - } - Literal result_literal(keys_literal.shape()); - std::vector zero_base(rank, 0); - std::vector increment(rank, 1); - increment[sort_dim] = sort_dim_elements; - // Iterate through each dimension except 'sort_dim'. - TF_RETURN_IF_ERROR(ShapeUtil::ForEachIndexWithStatus( - keys->shape(), zero_base, AsInt64Slice(keys->shape().dimensions()), - increment, [&](absl::Span indices) -> StatusOr { - // Extract a slice from the literal that corresponds to exactly the - // row in dimension 'sort_dim'. - std::vector limit_indices(indices.begin(), indices.end()); - std::for_each(limit_indices.begin(), limit_indices.end(), - [](int64& index) { ++index; }); - limit_indices[sort_dim] = sort_dim_elements; - TF_ASSIGN_OR_RETURN(auto row_to_sort, - keys_literal.Slice(indices, limit_indices) - .Reshape({sort_dim_elements})); - const auto& row_data = row_to_sort.data(); - - std::vector result_data(row_data.begin(), row_data.end()); - std::stable_sort(result_data.begin(), result_data.end(), - [](const NativeT& a, const NativeT& b) { - return SafeLess(a, b); - }); - Literal sorted_row(ShapeUtil::MakeShape(keys->shape().element_type(), - {sort_dim_elements})); - sorted_row.PopulateR1(absl::Span(result_data)); - std::vector slice_dimensions(rank, 1); - slice_dimensions[sort_dim] = sort_dim_elements; - TF_ASSIGN_OR_RETURN(auto sorted_row_reshaped, - sorted_row.Reshape(slice_dimensions)); - std::vector start_indices(rank, 0); - TF_RETURN_IF_ERROR(result_literal.CopySliceFrom( - sorted_row_reshaped, start_indices, indices, slice_dimensions)); - return true; - })); - parent_->evaluated_[sort] = std::move(result_literal); - return Status::OK(); - } - - template ::value || - std::is_same::value>::type* = - nullptr> - Status HandleSort(HloInstruction* sort) { - return UnsupportedTypeError(sort); - } - Status HandleSort(HloInstruction* sort) override { - return HandleSort(sort); + return UnsupportedTypeError(sort); } Status HandleReduce(HloInstruction* hlo) override { @@ -2649,7 +2582,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { template ::value>::type* = nullptr> Status HandleReducePrecision(HloInstruction* reduce_precision) { - return InvalidArgument("Double not supported for reduce precision"); + return InvalidArgument("Double is not supported for reduce precision"); } template < @@ -2664,12 +2597,13 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { return HandleReducePrecision(reduce_precision); } - template ::value || - std::is_same::value || - std::is_integral::value || - std::is_floating_point::value>::type* = nullptr> + template < + typename NativeT, + typename std::enable_if< + std::is_same::value || + std::is_same::value || + std::is_integral::value || is_complex_t::value || + std::is_floating_point::value>::type* = nullptr> Status HandleIota(HloInstruction* instruction) { auto* iota = Cast(instruction); const int64 iota_size = iota->shape().dimensions(iota->iota_dimension()); @@ -2700,12 +2634,13 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { return Status::OK(); } - template ::value || - std::is_same::value || - std::is_integral::value || - std::is_floating_point::value)>::type* = nullptr> + template < + typename NativeT, + typename std::enable_if< + !(std::is_same::value || + std::is_same::value || + std::is_integral::value || is_complex_t::value || + std::is_floating_point::value)>::type* = nullptr> Status HandleIota(HloInstruction* iota) { return UnsupportedTypeError(iota); } @@ -2713,6 +2648,103 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { return HandleIota(iota); } + template ::value || + std::is_floating_point::value)>::type* = nullptr> + Status HandleRng(HloInstruction* random) { + return UnsupportedTypeError(random); + } + template ::value)>::type* = nullptr> + Status HandleRng(HloInstruction* random) { + RandomDistribution distribution = random->random_distribution(); + const auto result_shape = random->shape(); + Literal result(result_shape); + + switch (distribution) { + case RNG_UNIFORM: { + const Literal& low = + parent_->GetEvaluatedLiteralFor(random->operand(0)); + const Literal& high = + parent_->GetEvaluatedLiteralFor(random->operand(1)); + + std::uniform_real_distribution generator( + low.Get({}), high.Get({})); + + TF_RETURN_IF_ERROR( + result.Populate([&](absl::Span /*indexes*/) { + return generator(parent_->engine_); + })); + break; + } + case RNG_NORMAL: { + const Literal& mean = + parent_->GetEvaluatedLiteralFor(random->operand(0)); + const Literal& stddev = + parent_->GetEvaluatedLiteralFor(random->operand(1)); + + std::normal_distribution generator(mean.Get({}), + stddev.Get({})); + + TF_RETURN_IF_ERROR( + result.Populate([&](absl::Span /*indexes*/) { + return generator(parent_->engine_); + })); + break; + } + default: + return UnimplementedStrCat("The distribution ", + RandomDistribution_Name(distribution), + " is not implemented."); + } + parent_->evaluated_[random] = std::move(result); + return Status::OK(); + } + template ::value)>::type* = + nullptr> + Status HandleRng(HloInstruction* random) { + RandomDistribution distribution = random->random_distribution(); + const auto result_shape = random->shape(); + Literal result(result_shape); + + switch (distribution) { + case RNG_UNIFORM: { + const Literal& low = + parent_->GetEvaluatedLiteralFor(random->operand(0)); + const Literal& high = + parent_->GetEvaluatedLiteralFor(random->operand(1)); + + // Note std::uniform_int_distribution assumes interval is closed, i.e., + // [low, high], but we want [low, high) instead. Hence high-1 is used as + // the upper range. + std::uniform_int_distribution generator( + low.Get({}), high.Get({}) - 1); + + TF_RETURN_IF_ERROR( + result.Populate([&](absl::Span /*indexes*/) { + return static_cast(generator(parent_->engine_)); + })); + break; + } + case RNG_NORMAL: { + return Unimplemented( + "Normal distribution is not supported for integral types."); + } + default: + return UnimplementedStrCat("The distribution ", + RandomDistribution_Name(distribution), + " is not implemented."); + } + parent_->evaluated_[random] = std::move(result); + return Status::OK(); + } + Status HandleRng(HloInstruction* random) override { + return HandleRng(random); + } + private: // Creates a vector of multipliers which can be used to create a linear index // into shape. @@ -2777,21 +2809,10 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { absl::Span start_indices, const Shape& result_shape) { std::vector start; - // TODO(b/118437727): Remove the R1 code-path. Note that to distinguish - // between the cases, this currently assumes there is at least 1 index. That - // is wrong in the general case, because for scalar indices, if the operand - // is scalar, then there are no indices. This problem with resolve itself. - const HloInstruction* first_index = start_indices[0]; - if (first_index->shape().rank() == 1) { - auto start_indices_typed = - parent_->GetEvaluatedLiteralFor(first_index).data(); - start = std::vector(start_indices_typed.begin(), - start_indices_typed.end()); - } else { - for (HloInstruction* index : start_indices) { - start.push_back( - parent_->GetEvaluatedLiteralFor(index).GetFirstElement()); - } + + for (HloInstruction* index : start_indices) { + start.push_back( + parent_->GetEvaluatedLiteralFor(index).GetFirstElement()); } // Clamp the start indices so the slice is in-bounds w.r.t the operand. @@ -2824,22 +2845,11 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { auto result = operand_literal.Clone(); const auto rank = result.shape().rank(); std::vector start; - // TODO(b/118437727): Remove the R1 code-path. Note that to distinguish - // between the cases, this currently assumes there is at least 1 index. That - // is wrong in the general case, because for scalar indices, if the operand - // is scalar, then there are no indices. This problem with resolve itself. - const HloInstruction* first_index = start_indices[0]; - if (first_index->shape().rank() == 1) { - auto start_indices_typed = - parent_->GetEvaluatedLiteralFor(first_index).data(); - start = std::vector(start_indices_typed.begin(), - start_indices_typed.end()); - } else { - for (HloInstruction* index : start_indices) { - start.push_back( - parent_->GetEvaluatedLiteralFor(index).GetFirstElement()); - } + for (HloInstruction* index : start_indices) { + start.push_back( + parent_->GetEvaluatedLiteralFor(index).GetFirstElement()); } + // Clamp the update start indices so the slice is in-bounds w.r.t the // operand. for (int64 i = 0; i < rank; ++i) { @@ -2956,6 +2966,7 @@ extern template class HloEvaluatorTypedVisitor; extern template class HloEvaluatorTypedVisitor; extern template class HloEvaluatorTypedVisitor; extern template class HloEvaluatorTypedVisitor; +extern template class HloEvaluatorTypedVisitor; extern template class HloEvaluatorTypedVisitor; } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_complex128.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_complex128.cc new file mode 100644 index 0000000000000000000000000000000000000000..1f48140ee4f6ca9415bef80c83664213109dbf9f --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_complex128.cc @@ -0,0 +1,22 @@ +/* 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/hlo_evaluator_typed_visitor.h" + +#include "tensorflow/compiler/xla/service/hlo_evaluator.h" + +namespace xla { +template class HloEvaluatorTypedVisitor; +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int16.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int16.cc new file mode 100644 index 0000000000000000000000000000000000000000..e54285a1577a3f3c97fba5ba6c2f969299ab599e --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int16.cc @@ -0,0 +1,22 @@ +/* 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/hlo_evaluator_typed_visitor.h" + +#include "tensorflow/compiler/xla/service/hlo_evaluator.h" + +namespace xla { +template class HloEvaluatorTypedVisitor; +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint16.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint16.cc new file mode 100644 index 0000000000000000000000000000000000000000..cc708952d20a00429944c8388a84a0e610c2f38f --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint16.cc @@ -0,0 +1,22 @@ +/* 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/hlo_evaluator_typed_visitor.h" + +#include "tensorflow/compiler/xla/service/hlo_evaluator.h" + +namespace xla { +template class HloEvaluatorTypedVisitor; +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.cc b/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.cc index c919dbd82d3668c477bf37074f1d56f8cb7d9506..862b2029718bbd802b69d789b66683a4edfa2367 100644 --- a/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.cc +++ b/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.cc @@ -17,6 +17,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/dynamic_dimension_inference.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/shape_inference.h" @@ -25,7 +26,9 @@ namespace xla { namespace { -StatusOr ReplaceGetSize(HloInstruction* instr) { +StatusOr ReplaceGetSize( + HloInstruction* instr, + const DynamicDimensionInference* dynamic_dimension_inference) { if (instr->opcode() != HloOpcode::kGetDimensionSize) { return false; } @@ -36,10 +39,18 @@ StatusOr ReplaceGetSize(HloInstruction* instr) { instr->operand(0)->shape(), instr->dimension())); TF_RET_CHECK(ShapeUtil::Equal(instr->shape(), legal_shape)); TF_RET_CHECK(ShapeUtil::HasPrimitiveType(instr->shape(), U32)); - uint32 size = instr->operand(0)->shape().dimensions(instr->dimension()); - HloInstruction* new_instr = computation->AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(size))); - TF_RETURN_IF_ERROR(instr->ReplaceAllUsesWith(new_instr)); + HloInstruction* operand = instr->mutable_operand(0); + int64 dim = instr->dimension(); + HloInstruction* dynamic_size = + dynamic_dimension_inference->GetDynamicSize(operand, {}, dim); + if (dynamic_size != nullptr) { + TF_RETURN_IF_ERROR(instr->ReplaceAllUsesWith(dynamic_size)); + } else { + uint32 size = instr->operand(0)->shape().dimensions(dim); + HloInstruction* new_instr = computation->AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(size))); + TF_RETURN_IF_ERROR(instr->ReplaceAllUsesWith(new_instr)); + } return true; } @@ -48,10 +59,13 @@ StatusOr ReplaceGetSize(HloInstruction* instr) { StatusOr HloGetDimensionSizeRewriter::Run(HloModule* module) { bool changed = false; HloProto proto; + TF_ASSIGN_OR_RETURN(DynamicDimensionInference inference, + DynamicDimensionInference::Run(module)); *proto.mutable_hlo_module() = module->ToProto(); for (auto* computation : module->computations()) { for (auto instruction : computation->instructions()) { - TF_ASSIGN_OR_RETURN(bool replaced, ReplaceGetSize(instruction)); + TF_ASSIGN_OR_RETURN(bool replaced, + ReplaceGetSize(instruction, &inference)); changed = changed || replaced; } } diff --git a/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.h b/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.h index 30f44c23a835b3bcc935caaa917e040e07c4e703..9aa79fe66b665c48ec871c4188e44ba2056de3ad 100644 --- a/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.h +++ b/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.h @@ -21,7 +21,9 @@ limitations under the License. namespace xla { -// Pass to replace a kGetDimensionSize instruction with a constant instruction. +// Pass to replace a kGetDimensionSize instruction with a hlo instruction +// representing the dynamic size if the dimension is dynamic, otherwise a +// constant instruction representing the static size. class HloGetDimensionSizeRewriter : public HloModulePass { public: absl::string_view name() const override { diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc index c6eaead8dd9c3dbf19bce34989e8b1f0f60468cc..e6f446c92687d0b27fcf1cdc4f38919e64c1035b 100644 --- a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc +++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc @@ -24,9 +24,9 @@ limitations under the License. #include #include #include -#include #include +#include "absl/container/flat_hash_map.h" #include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" @@ -380,7 +380,7 @@ class HloDotDumper { // Each HloInstruction dumped gets a monotically-increasing node ID. This // must start at 1, because that's where graphviz's accounting starts. int64 next_node_id_ = 1; - std::unordered_map node_ids_; + absl::flat_hash_map node_ids_; // The "root" tag doesn't have an associated HloInstruction pointer, so we // need to store it outside the map. @@ -397,7 +397,7 @@ class HloDotDumper { // Each HloComputation that's emitted gets a monotonically-increasing ID. int64 next_cluster_id_ = 1; - std::unordered_map cluster_ids_; + absl::flat_hash_map cluster_ids_; // Edges to print from Footer(). Edges come at the end because graphviz is // unhappy if an edge from a subcomputation to a node in the outer computation @@ -407,7 +407,7 @@ class HloDotDumper { // When coloring by sharding information, we track the sharding string // representation to color association, by round-robin the color schemes. - std::unordered_map + absl::flat_hash_map sharding_colors_; int64 next_shard_color_ = 0; }; @@ -561,8 +561,8 @@ bool HloDotDumper::ShouldShowSubcomputation(const HloComputation* subcomp) { } // Show the subcomputation if we're showing any of its members. - return std::any_of( - subcomp->instructions().begin(), subcomp->instructions().end(), + return absl::c_any_of( + subcomp->instructions(), [&](const HloInstruction* instr) { return filter_.Show(instr); }); } @@ -735,15 +735,14 @@ bool HloDotDumper::ShouldMergeIntoUsers(const HloInstruction* instr) const { const int kMinUsersToOmit = 3; return instr->opcode() == HloOpcode::kParameter && instr->shape().IsTuple() && !instr->IsFused() && - std::count_if(instr->users().begin(), instr->users().end(), - [&](const HloInstruction* user) { - return filter_.Show(user); - }) > kMinUsersToOmit && - std::all_of(instr->users().begin(), instr->users().end(), - [&](const HloInstruction* user) { - return !filter_.Show(user) || - user->opcode() == HloOpcode::kGetTupleElement; - }); + absl::c_count_if(instr->users(), + [&](const HloInstruction* user) { + return filter_.Show(user); + }) > kMinUsersToOmit && + absl::c_all_of(instr->users(), [&](const HloInstruction* user) { + return !filter_.Show(user) || + user->opcode() == HloOpcode::kGetTupleElement; + }); } string HloDotDumper::DumpInstruction(const HloInstruction* instr) { @@ -900,12 +899,11 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { // the same color as a parameter. Unless the merged-in parameter is a // parameter to a fusion node that is bound to a constant -- these aren't // "real" parameters from the user's perspective. - if (std::any_of(instr->operands().begin(), instr->operands().end(), - [&](const HloInstruction* operand) { - return operand->opcode() == HloOpcode::kParameter && - ShouldMergeIntoUsers(operand) && - TryGetFusionParameterConstant(operand) == nullptr; - })) { + if (absl::c_any_of(instr->operands(), [&](const HloInstruction* operand) { + return operand->opcode() == HloOpcode::kParameter && + ShouldMergeIntoUsers(operand) && + TryGetFusionParameterConstant(operand) == nullptr; + })) { return parameter_color; } @@ -1013,6 +1011,7 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { case HloOpcode::kConvolution: case HloOpcode::kDot: case HloOpcode::kFft: + case HloOpcode::kTriangularSolve: return kDarkBlue; case HloOpcode::kReducePrecision: return kRed; @@ -1039,6 +1038,7 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { case HloOpcode::kRecvDone: case HloOpcode::kSend: case HloOpcode::kSendDone: + case HloOpcode::kReplicaId: return kBrown; case HloOpcode::kCall: case HloOpcode::kConditional: @@ -1282,11 +1282,12 @@ namespace { // Gets a NodeFilter that includes roughly all instructions whose distance from // root is <= radius. -NodeFilter MakeNodeRadiusAroundFilter(const HloInstruction* root, - int64 radius) { +NodeFilter MakeNodeRadiusAroundFilter( + const HloInstruction* root, int64 radius, + const absl::flat_hash_set& boundary) { // First, find the neighborhood of nodes with distance from root <= radius. // These nodes are our initial set of "normal" nodes. - std::unordered_map nodes; + absl::flat_hash_map nodes; std::deque> worklist; worklist.push_back({root, 0}); while (!worklist.empty()) { @@ -1299,6 +1300,9 @@ NodeFilter MakeNodeRadiusAroundFilter(const HloInstruction* root, if (depth == radius) { continue; } + if (boundary.contains(instr)) { + continue; + } // Traverse into instr's operands. // @@ -1307,7 +1311,7 @@ NodeFilter MakeNodeRadiusAroundFilter(const HloInstruction* root, // are not interesting to the graph at hand. if (instr == root || instr->opcode() != HloOpcode::kTuple) { for (const HloInstruction* operand : instr->operands()) { - if (!nodes.count(operand)) { + if (!nodes.contains(operand)) { worklist.push_back({operand, depth + 1}); } } @@ -1335,7 +1339,7 @@ NodeFilter MakeNodeRadiusAroundFilter(const HloInstruction* root, continue; } for (const HloInstruction* user : instr->users()) { - if (!nodes.count(user)) { + if (!nodes.contains(user)) { worklist.push_back({user, depth + 1}); } } @@ -1344,7 +1348,7 @@ NodeFilter MakeNodeRadiusAroundFilter(const HloInstruction* root, auto is_displayed = [&](const HloInstruction* instr) { // Constants are displayed inline with their users; they're never omitted. // Nodes in subcomputations are always shown. - return nodes.count(instr) > 0 || instr->opcode() == HloOpcode::kConstant || + return nodes.contains(instr) || instr->opcode() == HloOpcode::kConstant || instr->parent() != root->parent(); }; @@ -1355,12 +1359,11 @@ NodeFilter MakeNodeRadiusAroundFilter(const HloInstruction* root, NodeFilterResult& filter_result = kv.second; const auto& operands = instr->operands(); - if (std::any_of(operands.begin(), operands.end(), is_displayed) && - !std::all_of(operands.begin(), operands.end(), is_displayed)) { + if (absl::c_any_of(operands, is_displayed) && + !absl::c_all_of(operands, is_displayed)) { // Mark nodes with some operands omitted appropriately. filter_result = kSomeOperandsOmitted; - } else if (!operands.empty() && - std::none_of(operands.begin(), operands.end(), is_displayed)) { + } else if (!operands.empty() && absl::c_none_of(operands, is_displayed)) { // Mark nodes with *all* operands omitted appropriately. filter_result = kOmitNodeOperands; } @@ -1368,8 +1371,7 @@ NodeFilter MakeNodeRadiusAroundFilter(const HloInstruction* root, // Promote nodes with type kSomeUsersOmitted to kNormalNode if all of their // users made it into the graph. if (filter_result == kSomeUsersOmitted && - std::all_of(instr->users().begin(), instr->users().end(), - is_displayed)) { + absl::c_all_of(instr->users(), is_displayed)) { filter_result = kNormalNode; } } @@ -1515,12 +1517,13 @@ string DumpGraph(const HloComputation& computation, const string& label, return graph_url; } -string DumpNeighborhoodAround(const HloInstruction& node, int radius, - bool show_backend_config) { +string DumpNeighborhoodAround( + const HloInstruction& node, int radius, bool show_backend_config, + const absl::flat_hash_set& boundary) { auto debug_options = node.GetModule()->config().debug_options(); string label = StrCat("Neighborhood of ", radius, " nodes around ", node.name()); - NodeFilter filter = MakeNodeRadiusAroundFilter(&node, radius); + NodeFilter filter = MakeNodeRadiusAroundFilter(&node, radius, boundary); string graph = HloDotDumper(node.parent(), label, debug_options, show_backend_config, /*profile=*/nullptr, filter) diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.h b/tensorflow/compiler/xla/service/hlo_graph_dumper.h index 8e51454ef1cf992386cc7325e32705c08bf7712f..b5444a32b18bfe75d048009a49f170930befd12d 100644 --- a/tensorflow/compiler/xla/service/hlo_graph_dumper.h +++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.h @@ -63,8 +63,12 @@ string DumpGraph(const HloComputation& computation, const string& label, // The number of nodes dumped is controlled by the radius parameter, which // (roughly) corresponds to the max distance a node may be from the primary node // before it's omitted from the graph. -string DumpNeighborhoodAround(const HloInstruction& node, int radius, - bool show_backend_config = false); +// +// The optional boundary specifies a set of boundary nodes, beyond which nodes +// will be omitted even if they are within the radius. +string DumpNeighborhoodAround( + const HloInstruction& node, int radius, bool show_backend_config = false, + const absl::flat_hash_set& boundary = {}); // Dumps nodes on any of the paths from `from` to `to`. If there are more than // max_nodes on all paths, restricts to the max_nodes nodes on the shortest diff --git a/tensorflow/compiler/xla/service/hlo_input_output_alias_config.cc b/tensorflow/compiler/xla/service/hlo_input_output_alias_config.cc index 70d4df5d1ced88301ac2fc44ed72af084e2a5c06..b01c00121b3363630b83a1e49d0027a66f3a9e1a 100644 --- a/tensorflow/compiler/xla/service/hlo_input_output_alias_config.cc +++ b/tensorflow/compiler/xla/service/hlo_input_output_alias_config.cc @@ -20,28 +20,31 @@ namespace xla { bool HloInputOutputAliasConfig::OutputHasAlias( const ShapeIndex& output_index) const { - return aliased_output_indices_.count(output_index) > 0; + return alias_.element(output_index).has_value(); } Status HloInputOutputAliasConfig::SetUpAlias(const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& param_index) { - TF_RET_CHECK(!OutputHasAlias(output_index)) - << "Output index " << output_index << " already has an alias setup"; + const ShapeIndex& param_index, + AliasKind kind) { + TF_RET_CHECK(kind == AliasKind::kUserAlias || kind == AliasKind::kSystemAlias) + << kind; TF_RET_CHECK(ShapeUtil::IndexIsValid(alias_.shape(), output_index)) << absl::StrCat("Tring to set up alias at ", output_index.ToString(), " which is an invalid index for shape ", ShapeUtil::HumanString(alias_.shape())); + TF_RET_CHECK(param_number >= 0) << param_number; + TF_RET_CHECK(!OutputHasAlias(output_index)) + << "Output index " << output_index << " already has an alias setup"; // Output can't be aliased with multiple parameters. TF_RET_CHECK(!alias_.element(output_index)) << absl::StrFormat( "Trying to set up output alias for param %lld at %s but failed: output " "index %s is already aliased with param %lld at %s", param_number, param_index.ToString(), output_index.ToString(), - alias_.element(output_index)->first, - alias_.element(output_index)->second.ToString()); + alias_.element(output_index)->parameter_number, + alias_.element(output_index)->parameter_index.ToString()); (*alias_.mutable_element(output_index)) = - std::make_pair(param_number, param_index); - aliased_output_indices_.insert(output_index); + Alias(kind, param_number, param_index); VLOG(4) << "Set up alias between output index " << output_index.ToString() << " and parameter " << param_index << " at index " << param_index.ToString(); @@ -51,15 +54,24 @@ Status HloInputOutputAliasConfig::SetUpAlias(const ShapeIndex& output_index, HloInputOutputAliasProto HloInputOutputAliasConfig::ToProto() const { HloInputOutputAliasProto result; alias_.ForEachElement( - [&](const ShapeIndex& index, - const absl::optional>& data) { + [&](const ShapeIndex& index, const absl::optional& data) { if (data) { HloInputOutputAliasProto::AliasEntryProto entry; + switch (data->kind) { + case AliasKind::kUserAlias: + entry.set_kind(HloInputOutputAliasProto::USER_ALIAS); + break; + case AliasKind::kSystemAlias: + entry.set_kind(HloInputOutputAliasProto::SYSTEM_ALIAS); + break; + default: + LOG(FATAL) << "Unknown alias kind " << data->kind; + } for (int64 i : index) { entry.add_output_shape_index(i); } - entry.set_parameter_number(data->first); - for (int64 i : data->second) { + entry.set_parameter_number(data->parameter_number); + for (int64 i : data->parameter_index) { entry.add_parameter_shape_index(i); } result.add_entries()->Swap(&entry); @@ -75,14 +87,18 @@ StatusOr HloInputOutputAliasConfig::CreateFromProto( proto.entries()) { ShapeIndex output_index(entry.output_shape_index().begin(), entry.output_shape_index().end()); - int64 param_number = entry.parameter_number(); ShapeIndex param_index(entry.parameter_shape_index().begin(), entry.parameter_shape_index().end()); + // Handle backward compatibility with existing protos, which only knew of + // system aliases. + AliasKind kind = AliasKind::kSystemAlias; + if (entry.kind() == HloInputOutputAliasProto::USER_ALIAS) { + kind = AliasKind::kUserAlias; + } TF_RETURN_IF_ERROR( - result.SetUpAlias(output_index, param_number, param_index)); + result.SetUpAlias(output_index, param_number, param_index, kind)); } - return result; } @@ -90,45 +106,44 @@ string HloInputOutputAliasConfig::ToString() const { std::vector pieces; pieces.push_back("HloInputOutputAliasConfig"); - ForEachAlias([&](const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& param_index) { + ForEachAlias([&](const ShapeIndex& output_index, const Alias& alias) { + const char* kind = alias.kind == AliasKind::kUserAlias ? "USER" : "SYSTEM"; pieces.push_back(absl::StrFormat( - " OutputIndex %s is aliased with parameter %lld at %s:", - output_index.ToString(), param_number, param_index.ToString())); + " OutputIndex %s is aliased (kind=%s) with parameter %lld at %s:", + output_index.ToString(), kind, alias.parameter_number, + alias.parameter_index.ToString())); }); - return absl::StrJoin(pieces, "\n"); } -bool HloInputOutputAliasConfig::ParameterHasAlias( +HloInputOutputAliasConfig::AliasKind +HloInputOutputAliasConfig::ParameterAliasKind( int64 param_number, const ShapeIndex& param_index) const { - bool output = false; + AliasKind kind = AliasKind::kNoAlias; alias_.ForEachElement( - [&](const xla::ShapeIndex&, - absl::optional> alias) { - if (alias && alias->first == param_number && - alias->second == param_index) { - output = true; + [&](const xla::ShapeIndex&, absl::optional alias) { + if (alias && alias->parameter_number == param_number && + alias->parameter_index == param_index) { + kind = alias->kind; } }); - return output; + return kind; } absl::optional HloInputOutputAliasConfig::GetAliasedOutput( int64 param_number, const ShapeIndex& param_index) const { absl::optional output; alias_.ForEachElement( - [&](const xla::ShapeIndex& output_index, - absl::optional> alias) { - if (alias && alias->first == param_number && - alias->second == param_index) { + [&](const xla::ShapeIndex& output_index, absl::optional alias) { + if (alias && alias->parameter_number == param_number && + alias->parameter_index == param_index) { output = output_index; } }); return output; } -absl::optional> +absl::optional HloInputOutputAliasConfig::GetAliasedParameter( const ShapeIndex& output_index) const { CHECK(ShapeUtil::IndexIsValid(alias_.shape(), output_index)); @@ -137,10 +152,9 @@ HloInputOutputAliasConfig::GetAliasedParameter( void HloInputOutputAliasConfig::ForEachAlias(AliasFn fn) const { alias_.ForEachElement( - [&](const ShapeIndex& output_index, - absl::optional> aliased) { + [&](const ShapeIndex& output_index, absl::optional aliased) { if (aliased) { - fn(output_index, aliased->first, aliased->second); + fn(output_index, *aliased); } }); } @@ -148,10 +162,9 @@ void HloInputOutputAliasConfig::ForEachAlias(AliasFn fn) const { Status HloInputOutputAliasConfig::ForEachAliasWithStatus( AliasFnWithStatus fn) const { return alias_.ForEachElementWithStatus( - [&](const ShapeIndex& output_index, - absl::optional> aliased) { + [&](const ShapeIndex& output_index, absl::optional aliased) { if (aliased) { - TF_RETURN_IF_ERROR(fn(output_index, aliased->first, aliased->second)); + TF_RETURN_IF_ERROR(fn(output_index, *aliased)); } return Status::OK(); }); @@ -167,20 +180,19 @@ Status HloInputOutputAliasConfig::Verify( param_has_seen.emplace_back(param->shape()); } return ForEachAliasWithStatus([&](const ShapeIndex& output_index, - int64 param_number, - const ShapeIndex& param_index) -> Status { + const Alias& alias) -> Status { const HloInstruction* root = entry->root_instruction(); - TF_RET_CHECK(0 <= param_number); - TF_RET_CHECK(entry->num_parameters() > param_number); + TF_RET_CHECK(0 <= alias.parameter_number); + TF_RET_CHECK(entry->num_parameters() > alias.parameter_number); const Shape& param_shape = - entry->parameter_instruction(param_number)->shape(); + entry->parameter_instruction(alias.parameter_number)->shape(); const Shape& output_shape = root->shape(); - TF_RET_CHECK(ShapeUtil::IndexIsValid(param_shape, param_index)); + TF_RET_CHECK(ShapeUtil::IndexIsValid(param_shape, alias.parameter_index)); TF_RET_CHECK(ShapeUtil::IndexIsValid(output_shape, output_index)); const Shape& param_subshape = - ShapeUtil::GetSubshape(param_shape, param_index); + ShapeUtil::GetSubshape(param_shape, alias.parameter_index); const Shape& output_subshape = ShapeUtil::GetSubshape(output_shape, output_index); TF_RET_CHECK(LayoutUtil::IsDenseArray(param_subshape)); @@ -191,19 +203,20 @@ Status HloInputOutputAliasConfig::Verify( "Expected aliased input %lld at index %s and output at index %s to " "have the same size. Input sub-shape is %s with size %lld, output " "sub-shape is %s with size %lld", - param_number, param_index.ToString(), output_index.ToString(), + alias.parameter_number, alias.parameter_index.ToString(), + output_index.ToString(), ShapeUtil::HumanStringWithLayout(param_subshape), size_func(param_subshape), ShapeUtil::HumanStringWithLayout(output_subshape), size_func(output_subshape)); } - // Check each param_number and param_index pair only show up once. No - // input can be aliased with output buffers. - TF_RET_CHECK(param_has_seen[param_number].element(param_index) == false); - - *(param_has_seen[param_number].mutable_element(param_index)) = true; - + // Check each alias.parameter_number and alias.parameter_index pair only + // show up once. No input can be aliased with output buffers. + TF_RET_CHECK(param_has_seen[alias.parameter_number].element( + alias.parameter_index) == false); + *(param_has_seen[alias.parameter_number].mutable_element( + alias.parameter_index)) = true; return Status::OK(); }); } diff --git a/tensorflow/compiler/xla/service/hlo_input_output_alias_config.h b/tensorflow/compiler/xla/service/hlo_input_output_alias_config.h index 3967743d53059c6cf1ddab8664e51fa8de7f4ac7..cd13c7a3ac7afe03fb99ed3114bdc6ac0f8ad6a7 100644 --- a/tensorflow/compiler/xla/service/hlo_input_output_alias_config.h +++ b/tensorflow/compiler/xla/service/hlo_input_output_alias_config.h @@ -32,21 +32,51 @@ class HloModule; // parameter index in the entry computation. class HloInputOutputAliasConfig { public: + // The kind of aliases which can be set. A kUserAlias is one setup at + // compilation time by the user, and has to be respected. A kSystemAlias one + // might be setup by the compiler, if it decides it is convenient to do so. + enum AliasKind { + kNoAlias, + kUserAlias, + kSystemAlias, + }; + + // Defines the alias information for a given output buffer. A given output + // buffer shape index can refer only to one parameter+index. + struct Alias { + Alias(AliasKind kind, int64 parameter_number, ShapeIndex parameter_index) + : kind(kind), + parameter_number(parameter_number), + parameter_index(std::move(parameter_index)) {} + + AliasKind kind; + int64 parameter_number; + ShapeIndex parameter_index; + }; + HloInputOutputAliasConfig() = default; - explicit HloInputOutputAliasConfig(Shape shape) : alias_(shape) {} + explicit HloInputOutputAliasConfig(Shape output_shape) + : alias_(output_shape) {} virtual ~HloInputOutputAliasConfig() = default; // Sets up alias config from `output_index` to `param_index` at // `param_number`. Status SetUpAlias(const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& param_index); + const ShapeIndex& param_index, AliasKind kind); + + // Returns the kind of alias for the given parameter number and parameter + // index. If no alias exists, AliasKind::kNoAlias is returned. + AliasKind ParameterAliasKind(int64 param_number, + const ShapeIndex& param_index) const; // Returns true if the given parameter is aliased with one of the output // buffers. bool ParameterHasAlias(int64 param_number, - const ShapeIndex& param_index) const; + const ShapeIndex& param_index) const { + return ParameterAliasKind(param_number, param_index) != AliasKind::kNoAlias; + } // Checks whether the provided output index has already been aliased. bool OutputHasAlias(const ShapeIndex& output_index) const; @@ -67,19 +97,17 @@ class HloInputOutputAliasConfig { // Returns the number of parameter and index of the parameter buffer that the // given output buffer index is aliased with. A nullopt is returned if there // is no parameter is aliased with the specific output. - absl::optional> GetAliasedParameter( + absl::optional GetAliasedParameter( const ShapeIndex& output_index) const; using AliasFn = - std::function; + std::function; // Iterates through each aliased output and input. void ForEachAlias(AliasFn fn) const; using AliasFnWithStatus = - std::function; + std::function; // Verifies that the given config is valid for the given module. // Specifically, the config's input and output should be in-bound and size of @@ -94,12 +122,10 @@ class HloInputOutputAliasConfig { private: // A ShapeTree which indicates the list of buffers that's expected to be // aliased. The key on this shape tree represents the output index. The value - // is a pair of parameter number and index into the buffer. If the value is - // nullopt, it means there is no parameter aliasing for this output. - ShapeTree>> alias_; - - // The indices of the output which have been aliased. - absl::flat_hash_set aliased_output_indices_; + // is an Alias data structure which defines the input parameter coordinates. + // If the value is nullopt, it means there is no parameter aliasing for this + // output. + ShapeTree> alias_; }; std::ostream& operator<<(std::ostream& out, diff --git a/tensorflow/compiler/xla/service/hlo_input_output_alias_config_test.cc b/tensorflow/compiler/xla/service/hlo_input_output_alias_config_test.cc index aeb9b0fdc8b6cca87731a2d4aae25120af6c3215..265bfdf7f989b0821a98c1f774cb408b78f348fe 100644 --- a/tensorflow/compiler/xla/service/hlo_input_output_alias_config_test.cc +++ b/tensorflow/compiler/xla/service/hlo_input_output_alias_config_test.cc @@ -29,7 +29,6 @@ limitations under the License. #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status_test_util.h" namespace xla { @@ -45,11 +44,12 @@ class HloInputOutputAliasConfigTest : public HloTestBase { EXPECT_TRUE(aliased_output); EXPECT_EQ(aliased_output.value(), output_index); - absl::optional> aliased_param = + absl::optional aliased_param = config.GetAliasedParameter(output_index); EXPECT_TRUE(aliased_param); - EXPECT_EQ(aliased_param.value(), std::make_pair(param_number, param_index)); + EXPECT_EQ(aliased_param->parameter_number, param_number); + EXPECT_EQ(aliased_param->parameter_index, param_index); } void expect_not_aliased(const ShapeIndex& output_index, int64 param_number, @@ -60,11 +60,12 @@ class HloInputOutputAliasConfigTest : public HloTestBase { EXPECT_FALSE(aliased_output && aliased_output == output_index); - absl::optional> aliased_param = + absl::optional aliased_param = config.GetAliasedParameter(output_index); - EXPECT_FALSE(aliased_param && aliased_param->first == param_number && - aliased_param->second == param_index); + EXPECT_FALSE(aliased_param && + aliased_param->parameter_number == param_number && + aliased_param->parameter_index == param_index); } }; @@ -84,8 +85,10 @@ ENTRY main { HloInputOutputAliasConfig config( module->entry_computation()->root_instruction()->shape()); - TF_ASSERT_OK(config.SetUpAlias(/*output_index=*/{0}, /*param_number=*/1, - /*param_index=*/{})); + TF_ASSERT_OK(config.SetUpAlias( + /*output_index=*/{0}, /*param_number=*/1, + /*param_index=*/{}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); expect_aliased(/*output_index=*/{0}, /*param_number=*/1, /*param_index=*/{}, config); @@ -114,11 +117,15 @@ ENTRY main { HloInputOutputAliasConfig config( module->entry_computation()->root_instruction()->shape()); - TF_ASSERT_OK(config.SetUpAlias(/*output_index=*/{0}, /*param_number=*/0, - /*param_index=*/{0})); + TF_ASSERT_OK(config.SetUpAlias( + /*output_index=*/{0}, /*param_number=*/0, + /*param_index=*/{0}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); - TF_ASSERT_OK(config.SetUpAlias(/*output_index=*/{1}, /*param_number=*/0, - /*param_index=*/{1})); + TF_ASSERT_OK(config.SetUpAlias( + /*output_index=*/{1}, /*param_number=*/0, + /*param_index=*/{1}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); expect_aliased(/*output_index=*/{0}, /*param_number=*/0, /*param_index=*/{0}, config); @@ -149,11 +156,15 @@ ENTRY main { HloInputOutputAliasConfig config( module->entry_computation()->root_instruction()->shape()); - TF_ASSERT_OK(config.SetUpAlias(/*output_index=*/{0}, /*param_number=*/0, - /*param_index=*/{})); + TF_ASSERT_OK(config.SetUpAlias( + /*output_index=*/{0}, /*param_number=*/0, + /*param_index=*/{}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); - TF_ASSERT_OK(config.SetUpAlias(/*output_index=*/{1}, /*param_number=*/0, - /*param_index=*/{})); + TF_ASSERT_OK(config.SetUpAlias( + /*output_index=*/{1}, /*param_number=*/0, + /*param_index=*/{}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); ASSERT_IS_NOT_OK(config.Verify(*module, [](const Shape& shape) { return ShapeUtil::ByteSizeOf(shape); @@ -176,8 +187,10 @@ ENTRY main { HloInputOutputAliasConfig config( module->entry_computation()->root_instruction()->shape()); - TF_ASSERT_OK(config.SetUpAlias(/*output_index=*/{1}, /*param_number=*/0, - /*param_index=*/{})); + TF_ASSERT_OK(config.SetUpAlias( + /*output_index=*/{1}, /*param_number=*/0, + /*param_index=*/{}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); ASSERT_IS_NOT_OK(config.Verify(*module, [](const Shape& shape) { return ShapeUtil::ByteSizeOf(shape); @@ -200,11 +213,15 @@ ENTRY main { HloInputOutputAliasConfig config( module->entry_computation()->root_instruction()->shape()); - TF_ASSERT_OK(config.SetUpAlias(/*output_index=*/{0}, /*param_number=*/0, - /*param_index=*/{})); + TF_ASSERT_OK(config.SetUpAlias( + /*output_index=*/{0}, /*param_number=*/0, + /*param_index=*/{}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); - ASSERT_IS_NOT_OK(config.SetUpAlias(/*output_index=*/{0}, /*param_number=*/1, - /*param_index=*/{})); + ASSERT_IS_NOT_OK(config.SetUpAlias( + /*output_index=*/{0}, /*param_number=*/1, + /*param_index=*/{}, + /*kind=*/HloInputOutputAliasConfig::AliasKind::kUserAlias)); } } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc index 66bc73740fed0c37bfc0f0c0bc8737a3a17fd7d4..aa1f3a2421f52c45145731a0203bd46f3ea574cf 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.cc +++ b/tensorflow/compiler/xla/service/hlo_instruction.cc @@ -83,15 +83,14 @@ StatusOr> HloInstruction::CreateFromProto( return computation_map.at(proto.called_computation_ids(index)); }; - TF_RET_CHECK(std::all_of( - proto.operand_ids().begin(), proto.operand_ids().end(), - [&instruction_map](int64 id) { return instruction_map.contains(id); })) + TF_RET_CHECK( + absl::c_all_of(proto.operand_ids(), + [&](int64 id) { return instruction_map.contains(id); })) << proto.name() << " instruction contains invalid operand id(s)"; - TF_RET_CHECK(std::all_of( - proto.called_computation_ids().begin(), - proto.called_computation_ids().end(), - [&computation_map](int64 id) { return computation_map.contains(id); })) + TF_RET_CHECK( + absl::c_all_of(proto.called_computation_ids(), + [&](int64 id) { return computation_map.contains(id); })) << proto.name() << " instruction references invalid computation id(s)"; Shape shape(proto.shape()); @@ -133,6 +132,14 @@ StatusOr> HloInstruction::CreateFromProto( absl::Span(fft_length)); 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::kSend: TF_RET_CHECK(proto.operand_ids_size() == 2) << "Send instruction should have 2 operand but sees " @@ -202,11 +209,12 @@ StatusOr> HloInstruction::CreateFromProto( << proto.operand_ids_size(); TF_RET_CHECK(proto.dimensions().size() == 1) << "Sort instruction should have 1 dimension"; + TF_RET_CHECK(proto.called_computation_ids_size() == 1) + << "Sort instruction should one called computation but sees " + << proto.called_computation_ids_size(); auto sort_operands = all_operands(); - HloInstruction* keys = sort_operands[0]; - instruction = CreateSort( - shape, proto.dimensions(0), keys, - absl::Span(sort_operands).subspan(1)); + instruction = CreateSort(shape, proto.dimensions(0), all_operands(), + computations(0)); break; } case HloOpcode::kTranspose: @@ -373,6 +381,13 @@ StatusOr> HloInstruction::CreateFromProto( CreateCollectivePermute(shape, operands(0), source_target_pairs); 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 " @@ -459,21 +474,17 @@ StatusOr> HloInstruction::CreateFromProto( << "DynamicSlice instruction should have at least 1 operands but " "sees " << proto.operand_ids_size(); - if (proto.operand_ids_size() == 2 && operands(1)->shape().rank() == 1) { - // TODO(b/118437727): Old form, remove this path. - instruction = - CreateDynamicSlice(shape, operands(0), operands(1), slice_sizes); - } else { - // New form + // TODO(b/118437727): Old form, make the check unconditional. + if (proto.operand_ids_size() != 2 || operands(1)->shape().rank() != 1) { auto expected_operands = 1 + operands(0)->shape().rank(); TF_RET_CHECK(proto.operand_ids_size() == expected_operands) << "DynamicSlice instruction should have " << expected_operands << " operands, but has " << proto.operand_ids_size(); - const auto& operand_vector = all_operands(); - instruction = CreateDynamicSlice( - shape, operands(0), absl::MakeSpan(operand_vector).subspan(1), - slice_sizes); } + const auto& operand_vector = all_operands(); + instruction = CreateDynamicSlice( + shape, operands(0), absl::MakeSpan(operand_vector).subspan(1), + slice_sizes); break; } case HloOpcode::kDynamicUpdateSlice: { @@ -481,22 +492,19 @@ StatusOr> HloInstruction::CreateFromProto( << "DynamicUpdateSlice instruction should have at least 2 operands " "but sees " << proto.operand_ids_size(); - if (proto.operand_ids_size() == 3 && operands(2)->shape().rank() == 1) { - // TODO(b/118437727): Old form, remove this path. - instruction = CreateDynamicUpdateSlice(shape, operands(0), operands(1), - operands(2)); - } else { - // New form + // TODO(b/118437727): Old form, make the check unconditional. + if (proto.operand_ids_size() != 3 || operands(2)->shape().rank() != 1) { auto expected_operands = 2 + operands(0)->shape().rank(); TF_RET_CHECK(proto.operand_ids_size() == expected_operands) << "DynamicUpdateSlice instruction should have " << expected_operands << " operands, but has " << proto.operand_ids_size(); - const auto& operand_vector = all_operands(); - instruction = - CreateDynamicUpdateSlice(shape, operands(0), operands(1), - absl::MakeSpan(operand_vector).subspan(2)); } + const auto& operand_vector = all_operands(); + instruction = + CreateDynamicUpdateSlice(shape, operands(0), operands(1), + absl::MakeSpan(operand_vector).subspan(2)); + break; } case HloOpcode::kGather: { @@ -791,6 +799,13 @@ HloInstruction::CreateGetTupleElement(const Shape& shape, fft_length); } +/* static */ std::unique_ptr +HloInstruction::CreateTriangularSolve(const Shape& shape, HloInstruction* a, + HloInstruction* b, + const TriangularSolveOptions& options) { + return absl::make_unique(shape, a, b, options); +} + /* static */ std::unique_ptr HloInstruction::CreateDot( const Shape& shape, HloInstruction* lhs, HloInstruction* rhs, const DotDimensionNumbers& dimension_numbers, @@ -833,6 +848,11 @@ HloInstruction::CreateCollectivePermute( shape, operand, source_target_pairs); } +/* static */ std::unique_ptr HloInstruction::CreateReplicaId() { + return absl::WrapUnique( + new HloInstruction(HloOpcode::kReplicaId, ShapeUtil::MakeShape(U32, {}))); +} + /* static */ std::unique_ptr HloInstruction::CreateInfeed( const Shape& infeed_shape, HloInstruction* token_operand, const string& config) { @@ -948,13 +968,6 @@ HloInstruction::CreateAddDependency(HloInstruction* data_operand, limit_indices, strides); } -/* static */ std::unique_ptr HloInstruction::CreateDynamicSlice( - const Shape& shape, HloInstruction* operand, HloInstruction* start_indices, - absl::Span slice_sizes) { - return absl::make_unique( - shape, operand, start_indices, slice_sizes); -} - /* static */ std::unique_ptr HloInstruction::CreateDynamicSlice( const Shape& shape, HloInstruction* operand, absl::Span start_indices, @@ -963,15 +976,6 @@ HloInstruction::CreateAddDependency(HloInstruction* data_operand, shape, operand, start_indices, slice_sizes); } -/* static */ std::unique_ptr -HloInstruction::CreateDynamicUpdateSlice(const Shape& shape, - HloInstruction* operand, - HloInstruction* update, - HloInstruction* start_indices) { - return absl::make_unique( - shape, operand, update, start_indices); -} - /* static */ std::unique_ptr HloInstruction::CreateDynamicUpdateSlice( const Shape& shape, HloInstruction* operand, HloInstruction* update, @@ -1165,9 +1169,10 @@ HloInstruction::CreateBroadcastSequence( } /* static */ std::unique_ptr HloInstruction::CreateSort( - const Shape& shape, int64 dimension, HloInstruction* keys, - absl::Span values) { - return absl::make_unique(shape, dimension, keys, values); + const Shape& shape, int64 dimension, + absl::Span operands, HloComputation* compare) { + return absl::make_unique(shape, dimension, operands, + compare); } /* static */ std::unique_ptr HloInstruction::CreateFusion( @@ -1359,6 +1364,7 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( case HloOpcode::kDot: case HloOpcode::kDomain: case HloOpcode::kGetDimensionSize: + case HloOpcode::kTriangularSolve: clone = CloneWithNewOperandsImpl(shape, new_operands, context); break; // Unary ops. @@ -1465,6 +1471,10 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( CHECK_EQ(new_operands.size(), 2); clone = CreateAddDependency(new_operands[0], new_operands[1]); break; + case HloOpcode::kReplicaId: + CHECK_EQ(new_operands.size(), 0); + clone = CreateReplicaId(); + break; } // SetupDerivedInstruction will setup the precision_config_ field. SetupDerivedInstruction(clone.get()); @@ -1599,12 +1609,10 @@ HloInstruction::InstructionVector HloInstruction::unique_operands() const { Status HloInstruction::AddControlDependencyTo(HloInstruction* instruction) { TF_RET_CHECK(instruction->parent() == parent()); - if (std::find(control_successors_.begin(), control_successors_.end(), - instruction) == control_successors_.end()) { + if (!absl::c_linear_search(control_successors_, instruction)) { control_successors_.push_back(instruction); - TF_RET_CHECK(std::find(instruction->control_predecessors_.begin(), - instruction->control_predecessors_.end(), - this) == instruction->control_predecessors_.end()); + TF_RET_CHECK( + !absl::c_linear_search(instruction->control_predecessors_, this)); instruction->control_predecessors_.push_back(this); } return Status::OK(); @@ -1736,6 +1744,7 @@ bool HloInstruction::IdenticalSlowPath( case HloOpcode::kReal: case HloOpcode::kRemainder: case HloOpcode::kReshape: + case HloOpcode::kReplicaId: case HloOpcode::kRoundNearestAfz: case HloOpcode::kSelect: case HloOpcode::kShiftLeft: @@ -1811,6 +1820,7 @@ bool HloInstruction::IdenticalSlowPath( case HloOpcode::kDot: case HloOpcode::kDomain: case HloOpcode::kGetDimensionSize: + case HloOpcode::kTriangularSolve: LOG(FATAL) << "Base class impl called for opcode with subclass: " << opcode(); } @@ -1853,7 +1863,7 @@ void HloInstruction::RemoveUser(HloInstruction* user) { user_set_.erase(set_it); // This is linear in the number of the users, but a vector provides a stable // iteration order and much faster traversal. - auto vec_it = std::find(users_.begin(), users_.end(), user); + auto vec_it = absl::c_find(users_, user); CHECK(vec_it != users_.end()); users_.erase(vec_it); } @@ -1871,8 +1881,7 @@ Status HloInstruction::ReplaceUseWith(HloInstruction* user, RemoveUser(user); - TF_RET_CHECK( - std::count(user->operands_.begin(), user->operands_.end(), this) >= 0); + TF_RET_CHECK(absl::c_count(user->operands_, this) >= 0); std::replace(user->operands_.begin(), user->operands_.end(), this, new_producer); new_producer->AddUser(user); @@ -1907,8 +1916,7 @@ Status HloInstruction::ReplaceOperandWithDifferentShape( VLOG(3) << "Replacing operand " << operand_num << " of " << name() << " with " << new_operand->name() << ", was " << old_operand->name(); - if (std::find(operands_.begin(), operands_.end(), old_operand) == - operands_.end()) { + if (!absl::c_linear_search(operands_, old_operand)) { old_operand->RemoveUser(this); } new_operand->AddUser(this); @@ -1963,6 +1971,7 @@ HloComputation* HloInstruction::to_apply() const { case HloOpcode::kReduce: case HloOpcode::kAllReduce: case HloOpcode::kScatter: + case HloOpcode::kSort: CHECK_EQ(called_computations_.size(), 1); return called_computations_[0]; default: @@ -1982,6 +1991,7 @@ void HloInstruction::set_to_apply(HloComputation* computation) { case HloOpcode::kReduce: case HloOpcode::kAllReduce: case HloOpcode::kScatter: + case HloOpcode::kSort: CHECK_EQ(called_computations_.size(), 1); called_computations_[0] = computation; break; @@ -2254,7 +2264,8 @@ std::vector HloInstruction::ExtraAttributesToString( opcode() == HloOpcode::kReduceWindow || opcode() == HloOpcode::kReduce || opcode() == HloOpcode::kAllReduce || - opcode() == HloOpcode::kScatter) { + opcode() == HloOpcode::kScatter || + opcode() == HloOpcode::kSort) { extra.push_back( StrCat("to_apply=", PrintName(to_apply()->name(), options))); } else if (!called_computations().empty()) { @@ -2291,6 +2302,7 @@ std::vector HloInstruction::ExtraAttributesToString( case HloOpcode::kReduce: case HloOpcode::kAllReduce: case HloOpcode::kScatter: + case HloOpcode::kSort: extra.push_back( StrCat("to_apply=\n", to_apply()->ToString(new_options))); break; @@ -2492,6 +2504,8 @@ Status HloInstruction::Visit(DfsHloVisitorBase* visitor) { return visitor->HandleAllToAll(this); case HloOpcode::kCollectivePermute: return visitor->HandleCollectivePermute(this); + case HloOpcode::kReplicaId: + return visitor->HandleReplicaId(this); case HloOpcode::kTuple: return visitor->HandleTuple(this); case HloOpcode::kMap: @@ -2594,6 +2608,8 @@ Status HloInstruction::Visit(DfsHloVisitorBase* visitor) { return visitor->HandleIota(this); case HloOpcode::kGetDimensionSize: return visitor->HandleGetDimensionSize(this); + case HloOpcode::kTriangularSolve: + return visitor->HandleTriangularSolve(this); // These opcodes are not handled here. case HloOpcode::kTrace: @@ -2945,10 +2961,10 @@ StatusOr StringToFusionKind( string PaddingConfigToString(const PaddingConfig& padding) { bool has_interior_padding = - std::any_of(padding.dimensions().begin(), padding.dimensions().end(), - [](const PaddingConfig::PaddingConfigDimension& dim) { - return dim.interior_padding() != 0; - }); + absl::c_any_of(padding.dimensions(), + [](const PaddingConfig::PaddingConfigDimension& dim) { + return dim.interior_padding() != 0; + }); return StrJoin( padding.dimensions(), "x", [&](string* out, const PaddingConfig::PaddingConfigDimension& dim) { @@ -3461,4 +3477,8 @@ const DomainMetadata& HloInstruction::operand_side_metadata() const { const DomainMetadata& HloInstruction::user_side_metadata() const { return Cast(this)->user_side_metadata(); } + +const TriangularSolveOptions& HloInstruction::triangular_solve_options() const { + return Cast(this)->triangular_solve_options(); +} } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h index 2827eed0df1fc837d391f583a9049b9e8597537d..e8ade80ef389d06b320b2386c24bb8e076fdeb48 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.h +++ b/tensorflow/compiler/xla/service/hlo_instruction.h @@ -435,6 +435,10 @@ class HloInstruction { const Shape& shape, HloInstruction* operand, FftType fft_type, absl::Span fft_length); + static std::unique_ptr CreateTriangularSolve( + const Shape& shape, HloInstruction* a, HloInstruction* b, + const TriangularSolveOptions& options); + // Creates a dot op with operands 'lhs' and 'rhs' with contracting and batch // dimensions specified in 'dimension_numbers'. static std::unique_ptr CreateDot( @@ -494,6 +498,9 @@ class HloInstruction { const Shape& shape, HloInstruction* operand, const std::vector>& source_target_pairs); + // Creates an instruction that returns a U32 replica ID. + static std::unique_ptr CreateReplicaId(); + // Creates a conversion instruction, where operand is the data to convert and // shape is the target shape for the conversion. static std::unique_ptr CreateConvert(const Shape& shape, @@ -556,10 +563,6 @@ class HloInstruction { // Creates a slice instruction, where the first operand is sliced by // start indices specified in the second operand, and by size specified in // 'slice_sizes'. - static std::unique_ptr CreateDynamicSlice( - const Shape& shape, HloInstruction* operand, - HloInstruction* start_indices, absl::Span slice_sizes); - // Same as above, but expects a span of scalar start indices. static std::unique_ptr CreateDynamicSlice( const Shape& shape, HloInstruction* operand, absl::Span start_indices, @@ -567,10 +570,6 @@ class HloInstruction { // Creates a dynamic update slice instruction, which updates a slice // of 'operand' with 'update' and 'start_indices'. - static std::unique_ptr CreateDynamicUpdateSlice( - const Shape& shape, HloInstruction* operand, HloInstruction* update, - HloInstruction* start_indices); - // Same as above, but expects a span of scalar start indices. static std::unique_ptr CreateDynamicUpdateSlice( const Shape& shape, HloInstruction* operand, HloInstruction* update, absl::Span start_indices); @@ -603,7 +602,6 @@ class HloInstruction { // f_2 = f(f_1.tuple_element(0), ..., f_1.tuple_element(N), input0.value1, // ..., inputN.value1) // ... - // TODO(b/112040122): Add support to this in HLO passes and in backends. static std::unique_ptr CreateReduce( const Shape& shape, absl::Span operands, absl::Span init_values, @@ -676,10 +674,14 @@ class HloInstruction { const Shape& shape, HloInstruction* operand, absl::Span dimensions); - // Creates a sort op, with a keys operand, and optional values operands. + // Creates a n-ary sort op with a 'compare' computation which is used for + // comparisons in the sorting algorithm. 'compare' gets 2 * n parameters, + // where parameters 2 * i and 2 * i + 1 are the values of the i-th operand at + // specific index positions which should be compared, and should return a + // PRED. static std::unique_ptr CreateSort( - const Shape& shape, int64 dimension, HloInstruction* keys, - absl::Span values = {}); + const Shape& shape, int64 dimension, + absl::Span operands, HloComputation* compare); // Creates a while instruction, given a condition computation, a body // computation, and the initial value for the input of the computations. For @@ -1247,6 +1249,10 @@ class HloInstruction { // on the instruction's existing name. void UniquifyName(NameUniquer* name_uniquer); + // Clear the unique ID of the instruction so that it can be re-assigned, such + // as for the purpose of compacting the instruction unique IDs. + void ClearUniqueIdInternal() { unique_id_ = -1; } + // Set the unique id for this instruction to "id" void SetUniqueId(int id) { CHECK_EQ(unique_id_, -1); // Should not be assigned already @@ -1559,6 +1565,9 @@ class HloInstruction { // Delegates to HloDomainInstruction::user_side_metadata(). const DomainMetadata& user_side_metadata() const; + // Delegates to HloTriangularSolveInstruction::triangular_solve_options(). + const TriangularSolveOptions& triangular_solve_options() const; + // Old methods kept for smooth subclassing transition END. protected: diff --git a/tensorflow/compiler/xla/service/hlo_instruction_test.cc b/tensorflow/compiler/xla/service/hlo_instruction_test.cc index 8048e332cb57747286758b75773b29ba154aa888..35f031f29a7aca8db7ebe2fbcfdcebb7a778d703 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction_test.cc +++ b/tensorflow/compiler/xla/service/hlo_instruction_test.cc @@ -20,6 +20,7 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/protobuf_util.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" @@ -55,13 +56,13 @@ class OpAndUserCollectingVisitor : public DfsHloVisitorWithDefault { } Status HandleParameter(HloInstruction* parameter) override { - EXPECT_EQ(0, count_.count(parameter)); + EXPECT_FALSE(count_.contains(parameter)); count_[parameter] = GetCountsForNode(parameter); return Status::OK(); } Status HandleConstant(HloInstruction* constant) override { - EXPECT_EQ(0, count_.count(constant)); + EXPECT_FALSE(count_.contains(constant)); count_[constant] = GetCountsForNode(constant); return Status::OK(); } @@ -69,25 +70,25 @@ class OpAndUserCollectingVisitor : public DfsHloVisitorWithDefault { Status HandleAdd(HloInstruction* add) override { auto lhs = add->operand(0); auto rhs = add->operand(1); - EXPECT_EQ(0, count_.count(add)); - EXPECT_GT(count_.count(lhs), 0); - EXPECT_GT(count_.count(rhs), 0); + EXPECT_FALSE(count_.contains(add)); + EXPECT_TRUE(count_.contains(lhs)); + EXPECT_TRUE(count_.contains(rhs)); count_[add] = GetCountsForNode(add); return Status::OK(); } Status HandleNegate(HloInstruction* negate) override { auto operand = negate->operand(0); - EXPECT_EQ(0, count_.count(negate)); - EXPECT_GT(count_.count(operand), 0); + EXPECT_FALSE(count_.contains(negate)); + EXPECT_TRUE(count_.contains(operand)); count_[negate] = GetCountsForNode(negate); return Status::OK(); } Status HandleMap(HloInstruction* map) override { - EXPECT_EQ(0, count_.count(map)); + EXPECT_FALSE(count_.contains(map)); for (HloInstruction* arg : map->operands()) { - EXPECT_GT(count_.count(arg), 0); + EXPECT_TRUE(count_.contains(arg)); } count_[map] = GetCountsForNode(map); return Status::OK(); @@ -96,9 +97,9 @@ class OpAndUserCollectingVisitor : public DfsHloVisitorWithDefault { Status HandleReduce(HloInstruction* reduce) override { auto arg = reduce->operand(0); auto init_value = reduce->operand(1); - EXPECT_EQ(0, count_.count(reduce)); - EXPECT_GT(count_.count(arg), 0); - EXPECT_GT(count_.count(init_value), 0); + EXPECT_FALSE(count_.contains(reduce)); + EXPECT_TRUE(count_.contains(arg)); + EXPECT_TRUE(count_.contains(init_value)); count_[reduce] = GetCountsForNode(reduce); return Status::OK(); } @@ -128,7 +129,7 @@ class OpAndUserCollectingVisitor : public DfsHloVisitorWithDefault { } // Counters for HLOs. Maps HLO to a NumOpsAndUsers. - std::unordered_map count_; + absl::flat_hash_map count_; }; TEST_F(HloInstructionTest, BasicProperties) { @@ -137,7 +138,7 @@ TEST_F(HloInstructionTest, BasicProperties) { EXPECT_EQ(HloOpcode::kParameter, parameter->opcode()); EXPECT_TRUE(ShapeUtil::IsScalarWithElementType(parameter->shape(), F32)); EXPECT_FALSE(ShapeUtil::IsScalarWithElementType(parameter->shape(), S32)); - EXPECT_EQ(0, parameter->operand_count()); + EXPECT_FALSE(parameter->operand_count()); } TEST_F(HloInstructionTest, UserWithTwoOperands) { @@ -981,9 +982,9 @@ TEST_F(HloInstructionTest, FunctionVisitor) { module->AddEntryComputation(builder.Build()); int visit_num = 0; - std::unordered_map visit_order; + absl::flat_hash_map visit_order; EXPECT_IS_OK(add->Accept([&visit_num, &visit_order](HloInstruction* inst) { - EXPECT_EQ(0, visit_order.count(inst)); + EXPECT_FALSE(visit_order.contains(inst)); visit_order[inst] = visit_num; visit_num++; return Status::OK(); diff --git a/tensorflow/compiler/xla/service/hlo_instructions.cc b/tensorflow/compiler/xla/service/hlo_instructions.cc index 7170dd7d8189bbc9ace17bb42bba63edc25b3b7b..92a74187c50db011b3c50ed6661354b5d33aef9e 100644 --- a/tensorflow/compiler/xla/service/hlo_instructions.cc +++ b/tensorflow/compiler/xla/service/hlo_instructions.cc @@ -42,11 +42,9 @@ using absl::StrJoin; bool IsInstructionElementwiseOnOperand(const HloInstruction* instruction, const HloInstruction* operand) { std::vector operand_indices = instruction->OperandIndices(operand); - return std::all_of( - operand_indices.begin(), operand_indices.end(), - [instruction](int64 operand_index) { - return instruction->IsElementwiseOnOperand(operand_index); - }); + return absl::c_all_of(operand_indices, [instruction](int64 operand_index) { + return instruction->IsElementwiseOnOperand(operand_index); + }); } string PrecisionConfigToString(const PrecisionConfig& precision_config) { @@ -203,6 +201,54 @@ std::unique_ptr HloFftInstruction::CloneWithNewOperandsImpl( fft_length_); } +HloTriangularSolveInstruction::HloTriangularSolveInstruction( + const Shape& shape, HloInstruction* a, HloInstruction* b, + const TriangularSolveOptions& options) + : HloInstruction(HloOpcode::kTriangularSolve, shape), + triangular_solve_options_(options) { + AppendOperand(a); + AppendOperand(b); +} + +HloInstructionProto HloTriangularSolveInstruction::ToProto() const { + HloInstructionProto proto = HloInstruction::ToProto(); + *proto.mutable_triangular_solve_options() = triangular_solve_options_; + return proto; +} + +std::vector HloTriangularSolveInstruction::ExtraAttributesToStringImpl( + const HloPrintOptions& options) const { + return {StrCat("left_side=", triangular_solve_options_.left_side()), + StrCat("lower=", triangular_solve_options_.lower()), + StrCat("unit_diagonal=", triangular_solve_options_.unit_diagonal()), + StrCat("transpose_a=", TriangularSolveOptions_Transpose_Name( + triangular_solve_options_.transpose_a()))}; +} + +bool HloTriangularSolveInstruction::IdenticalSlowPath( + const HloInstruction& other, + const std::function& + eq_computations) const { + const auto& casted_other = + static_cast(other); + const auto& options = triangular_solve_options(); + const auto& other_options = casted_other.triangular_solve_options(); + + return options.left_side() == other_options.left_side() && + options.lower() == other_options.lower() && + options.unit_diagonal() == other_options.unit_diagonal() && + options.transpose_a() == other_options.transpose_a(); +} + +std::unique_ptr +HloTriangularSolveInstruction::CloneWithNewOperandsImpl( + const Shape& shape, absl::Span new_operands, + HloCloneContext* context) const { + CHECK_EQ(new_operands.size(), 2); + return absl::make_unique( + shape, new_operands[0], new_operands[1], triangular_solve_options()); +} + HloSendRecvInstruction::HloSendRecvInstruction(HloOpcode opcode, const Shape& shape, int64 channel_id, @@ -385,6 +431,15 @@ HloInstructionProto HloAllReduceInstruction::ToProto() const { return proto; } +bool HloAllReduceInstruction::IsNoop() const { + for (auto replica_group : replica_groups()) { + if (replica_group.replica_ids().size() != 1) { + return false; + } + } + return !all_reduce_id(); +} + std::vector HloAllReduceInstruction::ExtraAttributesToStringImpl( const HloPrintOptions& options) const { std::vector result = @@ -602,14 +657,14 @@ std::unique_ptr HloReduceInstruction::CloneWithNewOperandsImpl( dimensions(), to_apply()); } -HloSortInstruction::HloSortInstruction(const Shape& shape, int64 dimension, - HloInstruction* keys, - absl::Span values) +HloSortInstruction::HloSortInstruction( + const Shape& shape, int64 dimension, + absl::Span operands, HloComputation* compare) : HloInstruction(HloOpcode::kSort, shape), dimensions_({dimension}) { - AppendOperand(keys); - for (auto* value : values) { + for (auto* value : operands) { AppendOperand(value); } + AppendComputation(compare); } HloInstructionProto HloSortInstruction::ToProto() const { @@ -630,15 +685,17 @@ bool HloSortInstruction::IdenticalSlowPath( const std::function& eq_computations) const { const auto& casted_other = static_cast(other); - return dimensions() == casted_other.dimensions(); + if (dimensions() != casted_other.dimensions()) { + return false; + } + return eq_computations(to_apply(), other.to_apply()); } std::unique_ptr HloSortInstruction::CloneWithNewOperandsImpl( const Shape& shape, absl::Span new_operands, HloCloneContext* context) const { - HloInstruction* keys = new_operands[0]; - return absl::make_unique(shape, dimensions(0), keys, - new_operands.subspan(1)); + return absl::make_unique(shape, dimensions(0), + new_operands, to_apply()); } HloTransposeInstruction::HloTransposeInstruction( @@ -814,8 +871,7 @@ std::vector HloSliceInstruction::ExtraAttributesToStringImpl( std::vector bounds; bounds.reserve(slice_starts_.size()); const bool omit_stride = - std::all_of(slice_strides_.begin(), slice_strides_.end(), - [](int64 stride) { return stride == 1; }); + absl::c_all_of(slice_strides_, [](int64 stride) { return stride == 1; }); for (int i = 0; i < slice_starts_.size(); ++i) { string stride_str = omit_stride ? "" : StrCat(":", slice_strides_[i]); bounds.push_back( @@ -1051,8 +1107,7 @@ HloInstruction* HloFusionInstruction::AddFusionOperand( void HloFusionInstruction::MergeFusionInstruction( HloFusionInstruction* instruction_to_merge) { - CHECK(std::find(operands().begin(), operands().end(), instruction_to_merge) != - operands().end()); + CHECK(absl::c_linear_search(operands(), instruction_to_merge)); // Clone the instruction from which to merge fused instructions. std::unique_ptr cloned = instruction_to_merge->Clone(); HloFusionInstruction* cloned_fusion = @@ -1219,8 +1274,8 @@ HloInstruction* HloFusionInstruction::CloneAndFuseInternal( // corresponding fused parameter instruction. Renumber parameters as // necessary to make parameter numbers consistent with their index in the // fused_parameter_ vector. - bool in_operand_list = std::find(operands().begin(), operands().end(), - instruction_to_fuse) != operands().end(); + bool in_operand_list = + absl::c_linear_search(operands(), instruction_to_fuse); CHECK(add_output || in_operand_list); if (instruction_to_fuse->opcode() == HloOpcode::kTuple) { // We assume all uses of a kTuple operation are GTE ops, not another @@ -1690,6 +1745,7 @@ HloInstructionProto HloConvolutionInstruction::ToProto() const { *proto.mutable_convolution_dimension_numbers() = convolution_dimension_numbers_; proto.set_feature_group_count(feature_group_count_); + proto.set_batch_group_count(batch_group_count_); *proto.mutable_precision_config() = precision_config_; return proto; } @@ -1706,6 +1762,10 @@ std::vector HloConvolutionInstruction::ExtraAttributesToStringImpl( extra.push_back(StrCat("feature_group_count=", feature_group_count_)); } + if (batch_group_count_ != 1) { + extra.push_back(StrCat("batch_group_count=", batch_group_count_)); + } + string precision_config_string = PrecisionConfigToString(precision_config_); if (!precision_config_string.empty()) { extra.push_back(precision_config_string); @@ -1723,6 +1783,9 @@ bool HloConvolutionInstruction::IdenticalSlowPath( if (feature_group_count_ != other.feature_group_count()) { return false; } + if (batch_group_count_ != other.batch_group_count()) { + return false; + } return protobuf_util::ProtobufEquals(window(), casted_other.window()) && protobuf_util::ProtobufEquals( convolution_dimension_numbers(), @@ -1841,6 +1904,7 @@ HloCustomCallInstruction::HloCustomCallInstruction( custom_call_target_(custom_call_target.begin(), custom_call_target.end()), opaque_(opaque.begin(), opaque.end()), feature_group_count_(1), + batch_group_count_(1), layout_constrained_(false) { for (auto operand : operands) { AppendOperand(operand); @@ -1855,6 +1919,7 @@ HloCustomCallInstruction::HloCustomCallInstruction( custom_call_target_(custom_call_target.begin(), custom_call_target.end()), opaque_(opaque.begin(), opaque.end()), feature_group_count_(1), + batch_group_count_(1), layout_constrained_(true), operand_shapes_with_layout_(operand_shapes_with_layout.begin(), operand_shapes_with_layout.end()) { @@ -1875,6 +1940,7 @@ HloInstructionProto HloCustomCallInstruction::ToProto() const { proto.set_custom_call_target(custom_call_target_); proto.set_custom_call_opaque(opaque_); proto.set_feature_group_count(feature_group_count_); + proto.set_batch_group_count(batch_group_count_); if (layout_constrained()) { proto.set_constrain_layout(true); for (const Shape& shape : operand_shapes_with_layout_) { @@ -1898,6 +1964,9 @@ std::vector HloCustomCallInstruction::ExtraAttributesToStringImpl( if (feature_group_count_ != 1) { extra.push_back(StrCat("feature_group_count=", feature_group_count_)); } + if (batch_group_count_ != 1) { + extra.push_back(StrCat("batch_group_count=", batch_group_count_)); + } // By contract, we print the custom call target even if // options.print_subcomputation_mode() == kOff, because the call target is not // an HloComputation. @@ -1941,6 +2010,9 @@ bool HloCustomCallInstruction::IdenticalSlowPath( if (feature_group_count_ != casted_other.feature_group_count_) { return false; } + if (batch_group_count_ != casted_other.batch_group_count_) { + return false; + } return custom_call_target_ == casted_other.custom_call_target_ && opaque_ == casted_other.opaque_; } @@ -1958,6 +2030,7 @@ HloCustomCallInstruction::CloneWithNewOperandsImpl( cloned->set_convolution_dimension_numbers(*convolution_dimension_numbers_); } cloned->set_feature_group_count(feature_group_count_); + cloned->set_batch_group_count(batch_group_count_); return std::move(cloned); } diff --git a/tensorflow/compiler/xla/service/hlo_instructions.h b/tensorflow/compiler/xla/service/hlo_instructions.h index e6111cfb57581589070b8e34556bdfe8239b4fd3..a0f2b46ba41cc6a60e28050d3cdd5e4e4583a875 100644 --- a/tensorflow/compiler/xla/service/hlo_instructions.h +++ b/tensorflow/compiler/xla/service/hlo_instructions.h @@ -131,6 +131,34 @@ class HloFftInstruction : public HloInstruction { std::vector fft_length_; }; +class HloTriangularSolveInstruction : public HloInstruction { + public: + explicit HloTriangularSolveInstruction(const Shape& shape, HloInstruction* a, + HloInstruction* b, + const TriangularSolveOptions& options); + const TriangularSolveOptions& triangular_solve_options() const { + return triangular_solve_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; + + TriangularSolveOptions triangular_solve_options_; +}; + class HloSendRecvInstruction : public HloInstruction { public: // Returns the channel id associated with the instruction. The id is @@ -253,6 +281,10 @@ class HloAllReduceInstruction : public HloCollectiveInstruction { // Returns a serialized representation of this instruction. HloInstructionProto ToProto() const override; + // Returns true if the AllReduce does no communication, so it's equivalent + // to a mem copy. + bool IsNoop() const; + private: std::vector ExtraAttributesToStringImpl( const HloPrintOptions& options) const override; @@ -414,8 +446,8 @@ class HloReduceInstruction : public HloInstruction { class HloSortInstruction : public HloInstruction { public: explicit HloSortInstruction(const Shape& shape, int64 dimension, - HloInstruction* keys, - absl::Span values = {}); + absl::Span operands, + HloComputation* compare); // Returns the dimension sizes or numbers associated with this instruction. const std::vector& dimensions() const override { return dimensions_; } int64 dimensions(int64 index) const override { return dimensions()[index]; } @@ -899,9 +931,7 @@ class HloOutfeedInstruction : public HloInstruction { HloInstruction* token_operand, absl::string_view outfeed_config); // Returns the shape for the Outfeed instruction. - const Shape& outfeed_shape() const { - return outfeed_shape_; - } + const Shape& outfeed_shape() const { return outfeed_shape_; } // Returns the config for the Outfeed instruction. const string& outfeed_config() const { return outfeed_config_; } // Returns a serialized representation of this instruction. diff --git a/tensorflow/compiler/xla/service/hlo_lexer.cc b/tensorflow/compiler/xla/service/hlo_lexer.cc index 5e81515134256a3ec4b790b38af3f42f68a79b56..2255383322873a39c7076e0f4f0dd541bc79014d 100644 --- a/tensorflow/compiler/xla/service/hlo_lexer.cc +++ b/tensorflow/compiler/xla/service/hlo_lexer.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include "absl/strings/ascii.h" #include "absl/strings/escaping.h" #include "absl/strings/numbers.h" #include "absl/strings/str_split.h" @@ -37,8 +38,8 @@ constexpr int kError = -2; // [a-zA-Z0-9_.-] bool IsIdentifierChar(char c) { - return isalnum(static_cast(c)) || c == '-' || c == '.' || - c == '_'; + return absl::ascii_isalnum(static_cast(c)) || c == '-' || + c == '.' || c == '_'; } } // namespace @@ -105,7 +106,7 @@ TokKind HloLexer::LexToken() { switch (current_char) { default: // [a-zA-Z_] - if (isalpha(static_cast(current_char)) || + if (absl::ascii_isalpha(static_cast(current_char)) || current_char == '_') { return LexIdentifier(); } @@ -152,6 +153,8 @@ TokKind HloLexer::LexToken() { return LexPercent(); case ':': return TokKind::kColon; + case '*': + return TokKind::kAsterisk; case '[': return TokKind::kLsquare; case ']': @@ -211,6 +214,15 @@ TokKind HloLexer::LexToken() { // A lone '/' is an error. return TokKind::kError; } + case '.': + if (PeekCurrentChar() == '.') { + current_ptr_++; + if (PeekCurrentChar() == '.') { + current_ptr_++; + return TokKind::kDots; + } + } + return TokKind::kError; case '"': return LexString(); } @@ -300,7 +312,7 @@ TokKind HloLexer::LexIdentifier() { // name ::= [a-zA-Z_][a-zA-Z0-9_.-]* TokKind HloLexer::LexPercent() { const char* name_start = current_ptr_; - if (isalpha(static_cast(PeekCurrentChar())) || + if (absl::ascii_isalpha(static_cast(PeekCurrentChar())) || PeekCurrentChar() == '_') { current_ptr_++; while (IsIdentifierChar(PeekCurrentChar())) { @@ -454,6 +466,8 @@ string TokKindToString(TokKind kind) { return "kComma"; case TokKind::kColon: return "kColon"; + case TokKind::kAsterisk: + return "kAsterisk"; case TokKind::kLsquare: return "kLsquare"; case TokKind::kRsquare: @@ -512,6 +526,8 @@ string TokKindToString(TokKind kind) { return "kInt"; case TokKind::kDecimal: return "kDecimal"; + case TokKind::kDots: + return "kDots"; } } diff --git a/tensorflow/compiler/xla/service/hlo_lexer.h b/tensorflow/compiler/xla/service/hlo_lexer.h index 94fac3cd8e9da7f273e7e521e21510f5188702e6..383fb4e862b8e32771879d055e663dc821a5c839 100644 --- a/tensorflow/compiler/xla/service/hlo_lexer.h +++ b/tensorflow/compiler/xla/service/hlo_lexer.h @@ -38,15 +38,17 @@ enum class TokKind { kError, // Tokens with no info. - kEqual, // = - kComma, // , - kColon, // : + kEqual, // = + kComma, // , + kColon, // : + kAsterisk, // * kLsquare, kRsquare, // [ ] kLbrace, kRbrace, // { } kLparen, kRparen, // ( ) + kDots, // ... kArrow, // -> kLeq, // <= @@ -107,7 +109,7 @@ class HloLexer { LOG(FATAL) << "This token does not have string value"; } } - tensorflow::int64 GetInt64Val() const { + int64 GetInt64Val() const { CHECK(GetKind() == TokKind::kInt); return token_state_.int64_val; } @@ -170,7 +172,7 @@ class HloLexer { const char* token_start = nullptr; TokKind current_kind; string str_val; - tensorflow::int64 int64_val; + int64 int64_val; double decimal_val; PrimitiveType primitive_type_val; }; diff --git a/tensorflow/compiler/xla/service/hlo_liveness_analysis.cc b/tensorflow/compiler/xla/service/hlo_liveness_analysis.cc index 5bf055f3c012fef687cdc275d62efdf2d4cd5e5c..e14bcfa7f67e736a4d04f5b236fb2df02cf150e0 100644 --- a/tensorflow/compiler/xla/service/hlo_liveness_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_liveness_analysis.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include "absl/container/flat_hash_set.h" #include "absl/memory/memory.h" #include "absl/strings/str_cat.h" #include "tensorflow/compiler/xla/map_util.h" @@ -36,11 +37,11 @@ namespace xla { namespace { using Worklist = std::deque; -using Workset = std::unordered_set; +using Workset = absl::flat_hash_set; void AddToWorklist(const HloInstruction* instruction, Worklist* worklist, Workset* workset) { - if (workset->count(instruction) == 0) { + if (!workset->contains(instruction)) { worklist->push_back(instruction); workset->insert(instruction); VLOG(3) << "ADD instruction: " << instruction->name(); diff --git a/tensorflow/compiler/xla/service/hlo_memory_scheduler.h b/tensorflow/compiler/xla/service/hlo_memory_scheduler.h index 7227bfb27c74758d2b79e404afc9eb97a1ca894d..76cc29cbb7848eb424d07abf11a95ffd59e9eed6 100644 --- a/tensorflow/compiler/xla/service/hlo_memory_scheduler.h +++ b/tensorflow/compiler/xla/service/hlo_memory_scheduler.h @@ -118,7 +118,7 @@ class HloTrivialScheduler : public HloModulePass { }; // A trivial pass which clears the schedule currently set on the -// HloModule. After this pass runs HloModudle::has_schedule will return false. +// HloModule. After this pass runs HloModule::has_schedule will return false. class HloDescheduler : public HloModulePass { public: HloDescheduler() = default; diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc index fe8371384c0fa3900a9022f101ff0b296439cf16..8322870cfd6a89fc6f863da8fd4a3576e8845cd7 100644 --- a/tensorflow/compiler/xla/service/hlo_module.cc +++ b/tensorflow/compiler/xla/service/hlo_module.cc @@ -107,11 +107,10 @@ HloComputation* HloModule::AddEntryComputation( } Status HloModule::RemoveEmbeddedComputation(HloComputation* to_remove) { - auto it = - std::find_if(computations_.begin(), computations_.end(), - [&to_remove](const std::unique_ptr& comp) { - return comp.get() == to_remove; - }); + auto it = absl::c_find_if( + computations_, [&to_remove](const std::unique_ptr& comp) { + return comp.get() == to_remove; + }); TF_RET_CHECK(it->get() == to_remove); computations_.erase(it); return Status::OK(); @@ -247,11 +246,39 @@ HloModuleProto HloModule::ToProto() const { return proto; } +Status HloModule::CheckUniqueNamesAndIdsForComputationsAndInstructions() const { + absl::flat_hash_set computation_names; + absl::flat_hash_set computation_ids; + absl::flat_hash_set instruction_names; + absl::flat_hash_set instruction_ids; + + for (const HloComputation* computation : computations()) { + TF_RET_CHECK(!ContainsKey(computation_names, computation->name())) + << "Computation name is not unique: " << computation->name(); + computation_names.insert(computation->name()); + + TF_RET_CHECK(!ContainsKey(computation_ids, computation->unique_id())) + << "Computation id is not unique: " << computation->unique_id(); + computation_ids.insert(computation->unique_id()); + + for (const HloInstruction* instruction : computation->instructions()) { + TF_RET_CHECK(!ContainsKey(instruction_names, instruction->name())) + << "Instruction name is not unique: " << instruction->name(); + instruction_names.insert(instruction->name()); + + TF_RET_CHECK(!ContainsKey(instruction_ids, instruction->unique_id())) + << "Instruction id is not unique: " << instruction->unique_id(); + instruction_ids.insert(instruction->unique_id()); + } + } + return Status::OK(); +} + /* static */ StatusOr> HloModule::CreateFromProto( const HloModuleProto& proto, const HloModuleConfig& module_config) { VLOG(2) << "CreateFromProto()"; - XLA_VLOG_LINES(2, proto.DebugString()); + XLA_VLOG_LINES(3, proto.DebugString()); // The ProgramShape in the passed in module config must match the shapes of // the entry parameters and root. @@ -304,11 +331,10 @@ StatusOr> HloModule::CreateFromProto( auto module = absl::make_unique(proto.name(), module_config); // Sort the computations in the proto id's order. - std::sort(computations.begin(), computations.end(), - [&](const std::unique_ptr& a, - const std::unique_ptr& b) { - return to_proto_id[a.get()] < to_proto_id[b.get()]; - }); + absl::c_sort(computations, [&](const std::unique_ptr& a, + const std::unique_ptr& b) { + return to_proto_id[a.get()] < to_proto_id[b.get()]; + }); // Add sorted computations to the module. for (auto& computation : computations) { @@ -331,28 +357,8 @@ StatusOr> HloModule::CreateFromProto( DynamicParameterBinding::CreateFromProto( proto.dynamic_parameter_binding())); - absl::flat_hash_set computation_names; - absl::flat_hash_set instruction_names; - absl::flat_hash_set computation_ids; - absl::flat_hash_set instruction_ids; - for (HloComputation* computation : module->computations()) { - TF_RET_CHECK(!ContainsKey(computation_names, computation->name())) - << "Computation name is not unique: " << computation->name(); - computation_names.insert(computation->name()); - - TF_RET_CHECK(!ContainsKey(computation_ids, computation->unique_id())) - << "Computation id is not unique: " << computation->unique_id(); - computation_ids.insert(computation->unique_id()); - for (HloInstruction* instruction : computation->instructions()) { - TF_RET_CHECK(!ContainsKey(instruction_names, instruction->name())) - << "Instruction name is not unique: " << instruction->name(); - instruction_names.insert(instruction->name()); - - TF_RET_CHECK(!ContainsKey(instruction_ids, instruction->unique_id())) - << "Instruction id is not unique: " << instruction->unique_id(); - instruction_ids.insert(instruction->unique_id()); - } - } + TF_RETURN_IF_ERROR( + module->CheckUniqueNamesAndIdsForComputationsAndInstructions()); if (proto.has_schedule()) { TF_ASSIGN_OR_RETURN( @@ -392,15 +398,12 @@ namespace { // Returns whether `hlo` is used outside the given subcomputation. // `instructions_in_subcomputation` is the instruction set of the given // subcomputation. -bool IsUsedOutsideSubcomputation( - const HloInstruction& hlo, - const std::unordered_set& instructions_in_subcomputation) { - for (HloInstruction* user : hlo.users()) { - if (!instructions_in_subcomputation.count(user)) { - return true; - } - } - return false; +bool IsUsedOutsideSubcomputation(const HloInstruction& hlo, + const absl::flat_hash_set& + instructions_in_subcomputation) { + return absl::c_any_of(hlo.users(), [&](HloInstruction* user) { + return !instructions_in_subcomputation.contains(user); + }); } } // anonymous namespace @@ -411,9 +414,9 @@ HloInstruction* HloModule::OutlineExpressionFromComputation( // A map from original instructions to their counterparts in the new outlined // function. - std::unordered_map outlined_instructions; + absl::flat_hash_map outlined_instructions; // A set that contains all instructions to be outlined. - std::unordered_set instruction_set_to_outline( + absl::flat_hash_set instruction_set_to_outline( instructions_to_outline.begin(), instructions_to_outline.end()); std::vector arguments; std::vector outputs; @@ -502,7 +505,7 @@ std::vector HloModule::MakeComputationPostOrder() const { // First determine all root computations by building a set of nonroot // computations (computations which are called by an instruction in the // module). - std::set nonroot_computations; + absl::flat_hash_set nonroot_computations; for (auto& computation : computations_) { for (auto* instruction : computation->instructions()) { for (HloComputation* called_computation : @@ -515,19 +518,19 @@ std::vector HloModule::MakeComputationPostOrder() const { // Keep track of computations which have already been added to the post // order. This prevents duplication as an embedded computation may be called // from two different root computations. - std::set added_computations; + absl::flat_hash_set added_computations; std::vector post_order; for (auto& computation : computations_) { - if (nonroot_computations.count(computation.get()) == 0) { + if (!nonroot_computations.contains(computation.get())) { for (HloComputation* embedded_computation : computation->MakeEmbeddedComputationsList()) { - if (added_computations.count(embedded_computation) == 0) { + if (!added_computations.contains(embedded_computation)) { post_order.push_back(embedded_computation); added_computations.insert(embedded_computation); } } // Root computations should only be encountered once. - CHECK_EQ(0, added_computations.count(computation.get())); + CHECK(!added_computations.contains(computation.get())); post_order.push_back(computation.get()); added_computations.insert(computation.get()); } diff --git a/tensorflow/compiler/xla/service/hlo_module.h b/tensorflow/compiler/xla/service/hlo_module.h index f1310e4b270898a21dbb4f86123edde4ba8993d0..b6fe6a5cdbd0934014f1152acd48c7a5973bead3 100644 --- a/tensorflow/compiler/xla/service/hlo_module.h +++ b/tensorflow/compiler/xla/service/hlo_module.h @@ -187,6 +187,7 @@ class HloModule { std::vector MakeNonfusionComputations() const; const HloModuleConfig& config() const { return config_; } + void set_config(HloModuleConfig& config) { config_ = config; } // Return a string representation of the module. // @@ -264,6 +265,18 @@ class HloModule { const HloSchedule& schedule() const { return *schedule_; } HloSchedule& schedule() { return *schedule_; } + HloComputation* AddComputationAndUnifyNamesAndIds( + std::unique_ptr computation, bool is_entry) { + computation->ClearUniqueIdInternal(); + for (auto* instruction : computation->instructions()) { + instruction->ClearUniqueIdInternal(); + } + return AddComputationInternal(std::move(computation), is_entry, + /*uniquify_identifiers=*/true); + } + + Status CheckUniqueNamesAndIdsForComputationsAndInstructions() const; + private: HloComputation* AddComputationInternal( std::unique_ptr computation, bool is_entry, diff --git a/tensorflow/compiler/xla/service/hlo_module_dce_test.cc b/tensorflow/compiler/xla/service/hlo_module_dce_test.cc index e535b7d74943943069b4d795cf999a3b1e963360..f6e2866204955ac024c2b6f972de449cc3df4c15 100644 --- a/tensorflow/compiler/xla/service/hlo_module_dce_test.cc +++ b/tensorflow/compiler/xla/service/hlo_module_dce_test.cc @@ -38,9 +38,7 @@ class HloModuleDceTest : public HloTestBase { // Returns whether the given instruction exists in the given computation. bool HasInstruction(const HloComputation& computation, const HloInstruction* instruction) { - return std::find(computation.instructions().begin(), - computation.instructions().end(), - instruction) != computation.instructions().end(); + return absl::c_linear_search(computation.instructions(), instruction); } // Returns whether the while instruction with name 'while_name' in diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc index 80f8ca2226b601d80af70395e485eb73246ddacb..b877081be5775bf6c75a69ffeba28d0f2cc17f90 100644 --- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc +++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc @@ -199,7 +199,7 @@ bool HloModuleGroupMetadata::IsChannelInstruction( } bool HloModuleGroupMetadata::IsCompanionInstruction(HloInstruction* hlo) const { - return companion_set_index_.count(hlo) > 0; + return companion_set_index_.contains(hlo); } bool HloModuleGroupMetadata::InstructionCommunicates( @@ -389,9 +389,10 @@ Status HloModuleGroupMetadata::AddCompanion(HloInstruction* instruction1, instruction1->opcode() == HloOpcode::kCall); VLOG(2) << "adding as companions:" << instruction1->ToString() << " and " << instruction2->ToString(); - - if (!ContainsKey(companion_set_index_, instruction1) && - !ContainsKey(companion_set_index_, instruction2)) { + if (instruction1 == instruction2) { + return Status::OK(); + } else if (!ContainsKey(companion_set_index_, instruction1) && + !ContainsKey(companion_set_index_, instruction2)) { companion_sets_.push_back( absl::make_unique>()); auto companion_set = companion_sets_.back().get(); @@ -419,7 +420,10 @@ Status HloModuleGroupMetadata::AddCompanion(HloInstruction* instruction1, for (HloInstruction* hlo : Companions(instruction2)) { companion_set_index_[hlo] = companion_set_index_[instruction1]; } - companion_sets_.erase(companion_sets_.begin() + index_to_remove); + // We can't remove the set from the vector because companion_set_index_ + // references sets by their index in this vector, so we reset to nullptr + // instead. + companion_sets_[index_to_remove].reset(nullptr); } return Status::OK(); } @@ -510,7 +514,7 @@ Status HloModuleGroupMetadata::CheckCommunicatingInstruction( HloComputation* computation = instruction->parent(); const HloModule* module = computation->parent(); if (module->entry_computation() == computation || - tracked_instructions_.count(computation) > 0) { + tracked_instructions_.contains(computation)) { return Status::OK(); } return FailedPrecondition("channel is used in disallowed computation"); diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h index 5cd0e38f3883a325dffe586f6ec46be9f6b799ab..84f7f2f31339ae9e98ea2301b6e6d94fcf4dedbb 100644 --- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h +++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h @@ -173,12 +173,13 @@ class HloModuleGroupMetadata { // Returns the number of modules for devices (excluding the host module). int64 GetDeviceModulesCount() const; - // Returns the companion instructions for the given instruction. + // Returns the companion set for the given instruction, including the + // instruction itself. // // Precondition: IsCompanionWhile(instruction) is true. const std::vector& Companions( const HloInstruction* instruction) const { - CHECK_EQ(companion_set_index_.count(instruction), 1); + CHECK(companion_set_index_.contains(instruction)); return companion_set(companion_set_index_.at(instruction)); } diff --git a/tensorflow/compiler/xla/service/hlo_opcode.h b/tensorflow/compiler/xla/service/hlo_opcode.h index 94122ac38ff2a3f7053b19e55f9a400c80ae2134..35626ba37541fc7a984ad05d12ebc22e9a08a550 100644 --- a/tensorflow/compiler/xla/service/hlo_opcode.h +++ b/tensorflow/compiler/xla/service/hlo_opcode.h @@ -137,9 +137,11 @@ namespace xla { 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(kWhile, "while") \ + V(kReplicaId, "replica-id") enum class HloOpcode { #define DECLARE_ENUM(enum_name, opcode_name, ...) enum_name, diff --git a/tensorflow/compiler/xla/service/hlo_ordering.cc b/tensorflow/compiler/xla/service/hlo_ordering.cc index ca6a154809be46d6a0305c29e2b89219de408019..0cec61c257bb84e467290fb52ec9063a32ed558d 100644 --- a/tensorflow/compiler/xla/service/hlo_ordering.cc +++ b/tensorflow/compiler/xla/service/hlo_ordering.cc @@ -367,7 +367,7 @@ bool SequentialHloOrdering::ExecutesBeforeInSameComputation( const HloInstruction* a, const HloInstruction* b) const { CHECK_EQ(a->parent(), b->parent()); // If either instruction is not in the order, then 'a' and 'b' are unordered. - if (order_position_.count(a) == 0 || order_position_.count(b) == 0) { + if (!order_position_.contains(a) || !order_position_.contains(b)) { return false; } return order_position_.at(a) < order_position_.at(b); diff --git a/tensorflow/compiler/xla/service/hlo_parser.cc b/tensorflow/compiler/xla/service/hlo_parser.cc index 730d39500b53484d5456e859c8436b3819a0a34b..20dbed07c546b3ec465e3b57c73a43c6c8f98efc 100644 --- a/tensorflow/compiler/xla/service/hlo_parser.cc +++ b/tensorflow/compiler/xla/service/hlo_parser.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/service/hlo_parser.h" +#include #include "absl/algorithm/container.h" #include "absl/memory/memory.h" @@ -21,11 +22,13 @@ 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/variant.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/service/hlo_domain_metadata.h" #include "tensorflow/compiler/xla/service/hlo_instructions.h" +#include "tensorflow/compiler/xla/service/hlo_lexer.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/hlo_schedule.h" #include "tensorflow/compiler/xla/service/hlo_sharding_metadata.h" @@ -44,8 +47,6 @@ using absl::StrCat; using absl::StrFormat; using absl::StrJoin; -const double kF16max = 65504; - // Creates and returns a schedule created using the order of the instructions in // the HloComputation::instructions() vectors in the module. HloSchedule ScheduleFromInstructionOrder(HloModule* module) { @@ -60,6 +61,10 @@ HloSchedule ScheduleFromInstructionOrder(HloModule* module) { return schedule; } +// Some functions accept either a linear index or a multi-dimensional index +// (used for indexing into sparse literals). +using LinearOrMultiIndex = absl::variant>; + // Parser for the HloModule::ToString() format text. class HloParser { public: @@ -102,7 +107,7 @@ class HloParser { // Parse a single instruction worth of text. bool ParseSingleInstruction(HloModule* module); - // ParseXXX returns false if an error occurred. + // Parses a module, returning false if an error occurred. bool ParseHloModule(HloModule* module); bool ParseComputations(HloModule* module); @@ -118,21 +123,30 @@ class HloParser { bool ParseNonTupleLiteral(Literal* literal, const Shape& shape); bool ParseDenseLiteral(Literal* literal, const Shape& shape); bool ParseSparseLiteral(Literal* literal, const Shape& shape); - template - bool ParseSparseLiteralHelper(Literal* literal, const Shape& shape); - // Sets the sub-value of literal at the given index to the given value. The - // literal's shape must have the default layout. - bool SetValueInLiteral(tensorflow::int64 value, - tensorflow::int64 linear_index, Literal* literal); - bool SetValueInLiteral(double value, tensorflow::int64 linear_index, + // Sets the sub-value of literal at the given linear or sparse index to the + // given value. If the literal is dense, it myst have the default layout. + // + // `loc` should be the source location of the value. + bool SetValueInLiteral(LocTy loc, int64 value, LinearOrMultiIndex index, + Literal* literal); + bool SetValueInLiteral(LocTy loc, double value, LinearOrMultiIndex index, Literal* literal); - bool SetValueInLiteral(bool value, tensorflow::int64 linear_index, + bool SetValueInLiteral(LocTy loc, bool value, LinearOrMultiIndex index, Literal* literal); + bool SetValueInLiteral(LocTy loc, std::complex value, + LinearOrMultiIndex index, Literal* literal); + // `loc` should be the source location of the value. template - bool SetValueInLiteralHelper(ParsedElemT value, - tensorflow::int64 linear_index, - Literal* literal); + bool SetValueInLiteralHelper(LocTy loc, ParsedElemT value, + LinearOrMultiIndex index, Literal* literal); + + // Checks whether the given value is within the range of LiteralNativeT. + // `loc` should be the source location of the value. + template + bool CheckParsedValueIsInRange(LocTy loc, ParsedElemT value); + template + bool CheckParsedValueIsInRange(LocTy loc, std::complex value); bool ParseOperands(std::vector* operands); // Fills parsed operands into 'operands' and expects a certain number of @@ -143,9 +157,9 @@ class HloParser { // Describes the start, limit, and stride on every dimension of the operand // being sliced. struct SliceRanges { - std::vector starts; - std::vector limits; - std::vector strides; + std::vector starts; + std::vector limits; + std::vector strides; }; // The data parsed for the kDomain instruction. @@ -165,6 +179,7 @@ class HloParser { kBracedInt64ListList, kHloComputation, kFftType, + kTriangularSolveTranspose, kWindow, kConvolutionDimensionNumbers, kSharding, @@ -237,16 +252,15 @@ class HloParser { bool ParseDomain(DomainData* domain); // Parses a sub-attribute of the window attribute, e.g.,size=1x2x3. - bool ParseDxD(const string& name, std::vector* result); + bool ParseDxD(const string& name, std::vector* result); // Parses window's pad sub-attriute, e.g., pad=0_0x3x3. - bool ParseWindowPad(std::vector>* pad); + bool ParseWindowPad(std::vector>* pad); bool ParseSliceRanges(SliceRanges* result); bool ParsePrecisionList(std::vector* result); bool ParseShapeList(std::vector* result); bool ParseInt64List(const TokKind start, const TokKind end, - const TokKind delim, - std::vector* result); + const TokKind delim, std::vector* result); // 'parse_and_add_item' is an lambda to parse an element in the list and add // the parsed element to the result. It's supposed to capture the result. bool ParseList(const TokKind start, const TokKind end, const TokKind delim, @@ -261,13 +275,16 @@ class HloParser { std::vector* dynamic_dimensions); bool ParseShape(Shape* result); bool ParseLayout(Layout* layout); + 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); - bool ParseInt64(tensorflow::int64* result); + bool ParseInt64(int64* result); bool ParseDouble(double* result); + bool ParseComplex(std::complex* result); bool ParseBool(bool* result); bool ParseToken(TokKind kind, const string& msg); @@ -640,11 +657,17 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, HloInstruction* instruction; switch (opcode) { case HloOpcode::kParameter: { - tensorflow::int64 parameter_number; + int64 parameter_number; if (!ParseToken(TokKind::kLparen, "expects '(' before parameter number") || - !ParseInt64(¶meter_number) || - !ParseToken(TokKind::kRparen, "expects ')' after parameter number") || + !ParseInt64(¶meter_number)) { + return false; + } + if (parameter_number < 0) { + Error(lexer_.GetLoc(), "parameter number must be >= 0"); + return false; + } + if (!ParseToken(TokKind::kRparen, "expects ')' after parameter number") || !ParseAttributes(attrs)) { return false; } @@ -666,7 +689,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kIota: { - optional iota_dimension; + optional iota_dimension; attrs["iota_dimension"] = {/*required=*/true, AttrTy::kInt64, &iota_dimension}; if (!ParseOperands(&operands, /*expected_size=*/0) || @@ -830,6 +853,14 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, HloInstruction::CreateCollectivePermute(shape, operands[0], pairs)); break; } + case HloOpcode::kReplicaId: { + if (!ParseOperands(&operands, /*expected_size=*/0) || + !ParseAttributes(attrs)) { + return false; + } + instruction = builder->AddInstruction(HloInstruction::CreateReplicaId()); + break; + } case HloOpcode::kReshape: { if (!ParseOperands(&operands, /*expected_size=*/1) || !ParseAttributes(attrs)) { @@ -861,17 +892,18 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kSort: { - optional> dimensions; + optional> dimensions; attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, &dimensions}; + optional to_apply; + attrs["to_apply"] = {/*required=*/true, AttrTy::kHloComputation, + &to_apply}; if (!ParseOperands(&operands) || !ParseAttributes(attrs) || dimensions->size() != 1) { return false; } instruction = builder->AddInstruction(HloInstruction::CreateSort( - shape, dimensions->at(0), - /*keys=*/operands[0], - /*values=*/absl::Span(operands).subspan(1))); + shape, dimensions->at(0), operands, to_apply.value())); break; } case HloOpcode::kTuple: { @@ -897,7 +929,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kRecv: { - optional channel_id; + optional channel_id; // If the is_host_transfer attribute is not present then default to false. optional is_host_transfer = false; attrs["channel_id"] = {/*required=*/true, AttrTy::kInt64, &channel_id}; @@ -913,7 +945,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kRecvDone: { - optional channel_id; + optional channel_id; // If the is_host_transfer attribute is not present then default to false. optional is_host_transfer = false; attrs["channel_id"] = {/*required=*/true, AttrTy::kInt64, &channel_id}; @@ -931,7 +963,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kSend: { - optional channel_id; + optional channel_id; // If the is_host_transfer attribute is not present then default to false. optional is_host_transfer = false; attrs["channel_id"] = {/*required=*/true, AttrTy::kInt64, &channel_id}; @@ -946,7 +978,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kSendDone: { - optional channel_id; + optional channel_id; // If the is_host_transfer attribute is not present then default to false. optional is_host_transfer = false; attrs["channel_id"] = {/*required=*/true, AttrTy::kInt64, &channel_id}; @@ -964,7 +996,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kGetTupleElement: { - optional index; + optional index; attrs["index"] = {/*required=*/true, AttrTy::kInt64, &index}; if (!ParseOperands(&operands, /*expected_size=*/1) || !ParseAttributes(attrs)) { @@ -1047,7 +1079,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, } case HloOpcode::kFft: { optional fft_type; - optional> fft_length; + optional> fft_length; attrs["fft_type"] = {/*required=*/true, AttrTy::kFftType, &fft_type}; attrs["fft_length"] = {/*required=*/true, AttrTy::kBracedInt64List, &fft_length}; @@ -1059,8 +1091,40 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, shape, operands[0], *fft_type, *fft_length)); 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}; + if (!ParseOperands(&operands, /*expected_size=*/2) || + !ParseAttributes(attrs)) { + 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::kBroadcast: { - optional> broadcast_dimensions; + optional> broadcast_dimensions; attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, &broadcast_dimensions}; if (!ParseOperands(&operands, /*expected_size=*/1) || @@ -1072,7 +1136,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kConcatenate: { - optional> dimensions; + optional> dimensions; attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, &dimensions}; if (!ParseOperands(&operands) || !ParseAttributes(attrs) || @@ -1087,7 +1151,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, optional to_apply; attrs["to_apply"] = {/*required=*/true, AttrTy::kHloComputation, &to_apply}; - optional> dimensions; + optional> dimensions; attrs["dimensions"] = {/*required=*/false, AttrTy::kBracedInt64List, &dimensions}; if (!ParseOperands(&operands) || !ParseAttributes(attrs)) { @@ -1103,7 +1167,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, optional reduce_computation; attrs["to_apply"] = {/*required=*/true, AttrTy::kHloComputation, &reduce_computation}; - optional> dimensions_to_reduce; + optional> dimensions_to_reduce; attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, &dimensions_to_reduce}; if (!ParseOperands(&operands) || !ParseAttributes(attrs)) { @@ -1124,7 +1188,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kReverse: { - optional> dimensions; + optional> dimensions; attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, &dimensions}; if (!ParseOperands(&operands, /*expected_size=*/1) || @@ -1168,7 +1232,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kDynamicSlice: { - optional> dynamic_slice_sizes; + optional> dynamic_slice_sizes; attrs["dynamic_slice_sizes"] = { /*required=*/true, AttrTy::kBracedInt64List, &dynamic_slice_sizes}; LocTy loc = lexer_.GetLoc(); @@ -1207,7 +1271,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kTranspose: { - optional> dimensions; + optional> dimensions; attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, &dimensions}; if (!ParseOperands(&operands, /*expected_size=*/1) || @@ -1221,7 +1285,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, case HloOpcode::kBatchNormTraining: { optional epsilon; attrs["epsilon"] = {/*required=*/true, AttrTy::kFloat, &epsilon}; - optional feature_index; + optional feature_index; attrs["feature_index"] = {/*required=*/true, AttrTy::kInt64, &feature_index}; if (!ParseOperands(&operands, /*expected_size=*/3) || @@ -1237,7 +1301,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, case HloOpcode::kBatchNormInference: { optional epsilon; attrs["epsilon"] = {/*required=*/true, AttrTy::kFloat, &epsilon}; - optional feature_index; + optional feature_index; attrs["feature_index"] = {/*required=*/true, AttrTy::kInt64, &feature_index}; if (!ParseOperands(&operands, /*expected_size=*/5) || @@ -1254,7 +1318,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, case HloOpcode::kBatchNormGrad: { optional epsilon; attrs["epsilon"] = {/*required=*/true, AttrTy::kFloat, &epsilon}; - optional feature_index; + optional feature_index; attrs["feature_index"] = {/*required=*/true, AttrTy::kInt64, &feature_index}; if (!ParseOperands(&operands, /*expected_size=*/5) || @@ -1336,8 +1400,8 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kReducePrecision: { - optional exponent_bits; - optional mantissa_bits; + optional exponent_bits; + optional mantissa_bits; attrs["exponent_bits"] = {/*required=*/true, AttrTy::kInt64, &exponent_bits}; attrs["mantissa_bits"] = {/*required=*/true, AttrTy::kInt64, @@ -1375,6 +1439,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, optional window; optional dnums; optional feature_group_count; + optional batch_group_count; optional> operand_layout_constraints; attrs["custom_call_target"] = {/*required=*/true, AttrTy::kString, &custom_call_target}; @@ -1384,6 +1449,8 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, AttrTy::kConvolutionDimensionNumbers, &dnums}; attrs["feature_group_count"] = {/*required=*/false, AttrTy::kInt64, &feature_group_count}; + attrs["batch_group_count"] = {/*required=*/false, AttrTy::kInt64, + &batch_group_count}; attrs["operand_layout_constraints"] = { /*required=*/false, AttrTy::kShapeList, &operand_layout_constraints}; if (!ParseOperands(&operands) || !ParseAttributes(attrs)) { @@ -1439,19 +1506,22 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, if (feature_group_count.has_value()) { instruction->set_feature_group_count(*feature_group_count); } + if (batch_group_count.has_value()) { + instruction->set_batch_group_count(*batch_group_count); + } break; } case HloOpcode::kDot: { - optional> lhs_contracting_dims; + optional> lhs_contracting_dims; attrs["lhs_contracting_dims"] = { /*required=*/false, AttrTy::kBracedInt64List, &lhs_contracting_dims}; - optional> rhs_contracting_dims; + optional> rhs_contracting_dims; attrs["rhs_contracting_dims"] = { /*required=*/false, AttrTy::kBracedInt64List, &rhs_contracting_dims}; - optional> lhs_batch_dims; + optional> lhs_batch_dims; attrs["lhs_batch_dims"] = {/*required=*/false, AttrTy::kBracedInt64List, &lhs_batch_dims}; - optional> rhs_batch_dims; + optional> rhs_batch_dims; attrs["rhs_batch_dims"] = {/*required=*/false, AttrTy::kBracedInt64List, &rhs_batch_dims}; optional> operand_precision; @@ -1495,19 +1565,19 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kGather: { - optional> offset_dims; + optional> offset_dims; attrs["offset_dims"] = {/*required=*/true, AttrTy::kBracedInt64List, &offset_dims}; - optional> collapsed_slice_dims; + optional> collapsed_slice_dims; attrs["collapsed_slice_dims"] = { /*required=*/true, AttrTy::kBracedInt64List, &collapsed_slice_dims}; - optional> start_index_map; + optional> start_index_map; attrs["start_index_map"] = {/*required=*/true, AttrTy::kBracedInt64List, &start_index_map}; - optional index_vector_dim; + optional index_vector_dim; attrs["index_vector_dim"] = {/*required=*/true, AttrTy::kInt64, &index_vector_dim}; - optional> slice_sizes; + optional> slice_sizes; attrs["slice_sizes"] = {/*required=*/true, AttrTy::kBracedInt64List, &slice_sizes}; @@ -1529,17 +1599,17 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, break; } case HloOpcode::kScatter: { - optional> update_window_dims; + optional> update_window_dims; attrs["update_window_dims"] = { /*required=*/true, AttrTy::kBracedInt64List, &update_window_dims}; - optional> inserted_window_dims; + optional> inserted_window_dims; attrs["inserted_window_dims"] = { /*required=*/true, AttrTy::kBracedInt64List, &inserted_window_dims}; - optional> scatter_dims_to_operand_dims; + optional> scatter_dims_to_operand_dims; attrs["scatter_dims_to_operand_dims"] = {/*required=*/true, AttrTy::kBracedInt64List, &scatter_dims_to_operand_dims}; - optional index_vector_dim; + optional index_vector_dim; attrs["index_vector_dim"] = {/*required=*/true, AttrTy::kInt64, &index_vector_dim}; @@ -1580,7 +1650,7 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, return TokenError(StrCat("parsing not yet implemented for op: ", HloOpcodeString(opcode))); case HloOpcode::kGetDimensionSize: - optional> dimensions; + optional> dimensions; attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, &dimensions}; if (!ParseOperands(&operands, /*expected_size=*/1) || @@ -1669,8 +1739,8 @@ bool HloParser::ParseSingleSharding(OpSharding* sharding, LocTy loc = lexer_.GetLoc(); bool maximal = false; bool replicated = false; - std::vector devices; - std::vector tile_assignment_dimensions; + std::vector devices; + std::vector tile_assignment_dimensions; while (lexer_.GetKind() != TokKind::kRbrace) { switch (lexer_.GetKind()) { case TokKind::kw_maximal: @@ -1696,7 +1766,7 @@ bool HloParser::ParseSingleSharding(OpSharding* sharding, } do { - tensorflow::int64 dim; + int64 dim; if (!ParseInt64(&dim)) { return false; } @@ -1708,7 +1778,7 @@ bool HloParser::ParseSingleSharding(OpSharding* sharding, return false; } do { - tensorflow::int64 device; + int64 device; if (!ParseInt64(&device)) { return false; } @@ -1752,10 +1822,10 @@ bool HloParser::ParseSingleSharding(OpSharding* sharding, "dimensions"); } sharding->set_type(OpSharding::Type::OpSharding_Type_OTHER); - for (tensorflow::int64 dim : tile_assignment_dimensions) { + for (int64 dim : tile_assignment_dimensions) { sharding->add_tile_assignment_dimensions(dim); } - for (tensorflow::int64 device : devices) { + for (int64 device : devices) { sharding->add_tile_assignment_devices(device); } } @@ -1816,130 +1886,146 @@ bool HloParser::ParseInstructionNames( "expects '}' at the end of instruction name list"); } -bool HloParser::SetValueInLiteral(tensorflow::int64 value, - tensorflow::int64 linear_index, - Literal* literal) { +bool HloParser::SetValueInLiteral(LocTy loc, int64 value, + LinearOrMultiIndex index, Literal* literal) { const Shape& shape = literal->shape(); switch (shape.element_type()) { case S8: - return SetValueInLiteralHelper(value, linear_index, - literal); + return SetValueInLiteralHelper(loc, value, index, literal); case S16: - return SetValueInLiteralHelper(value, linear_index, - literal); + return SetValueInLiteralHelper(loc, value, index, literal); case S32: - return SetValueInLiteralHelper(value, linear_index, - literal); + return SetValueInLiteralHelper(loc, value, index, literal); case S64: - return SetValueInLiteralHelper(value, linear_index, - literal); + return SetValueInLiteralHelper(loc, value, index, literal); case U8: - return SetValueInLiteralHelper(value, linear_index, + return SetValueInLiteralHelper(loc, value, index, literal); case U16: - return SetValueInLiteralHelper(value, linear_index, + return SetValueInLiteralHelper(loc, value, index, literal); case U32: - return SetValueInLiteralHelper(value, linear_index, + return SetValueInLiteralHelper(loc, value, index, literal); case U64: - return SetValueInLiteralHelper(value, linear_index, + return SetValueInLiteralHelper(loc, value, index, literal); case PRED: // Bool type literals with rank >= 1 are printed in 0s and 1s. - return SetValueInLiteralHelper(static_cast(value), - linear_index, literal); + return SetValueInLiteralHelper(loc, static_cast(value), index, + literal); default: LOG(FATAL) << "unknown integral primitive type " << PrimitiveType_Name(shape.element_type()); } } -bool HloParser::SetValueInLiteral(double value, tensorflow::int64 linear_index, - Literal* literal) { +bool HloParser::SetValueInLiteral(LocTy loc, double value, + LinearOrMultiIndex index, Literal* literal) { const Shape& shape = literal->shape(); switch (shape.element_type()) { case F16: - return SetValueInLiteralHelper(value, linear_index, literal); + return SetValueInLiteralHelper(loc, value, index, literal); case BF16: - return SetValueInLiteralHelper(value, linear_index, + return SetValueInLiteralHelper(loc, value, index, literal); case F32: - return SetValueInLiteralHelper(value, linear_index, literal); + return SetValueInLiteralHelper(loc, value, index, literal); case F64: - return SetValueInLiteralHelper(value, linear_index, literal); + return SetValueInLiteralHelper(loc, value, index, literal); default: LOG(FATAL) << "unknown floating point primitive type " << PrimitiveType_Name(shape.element_type()); } } -bool HloParser::SetValueInLiteral(bool value, tensorflow::int64 linear_index, - Literal* literal) { +bool HloParser::SetValueInLiteral(LocTy loc, bool value, + LinearOrMultiIndex index, Literal* literal) { const Shape& shape = literal->shape(); switch (shape.element_type()) { case PRED: - return SetValueInLiteralHelper(value, linear_index, literal); + return SetValueInLiteralHelper(loc, value, index, literal); default: LOG(FATAL) << PrimitiveType_Name(shape.element_type()) << " is not PRED type"; } } +bool HloParser::SetValueInLiteral(LocTy loc, std::complex value, + LinearOrMultiIndex index, Literal* literal) { + const Shape& shape = literal->shape(); + switch (shape.element_type()) { + case C64: + return SetValueInLiteralHelper>(loc, value, index, + literal); + case C128: + return SetValueInLiteralHelper>(loc, value, index, + literal); + default: + LOG(FATAL) << PrimitiveType_Name(shape.element_type()) + << " is not a complex type type"; + } +} + +template +string StringifyValue(T val) { + return StrCat(val); +} +template <> +string StringifyValue(std::complex val) { + return StrFormat("(%f, %f)", std::real(val), std::imag(val)); +} + template -bool HloParser::SetValueInLiteralHelper(ParsedElemT value, - tensorflow::int64 linear_index, +bool HloParser::SetValueInLiteralHelper(LocTy loc, ParsedElemT value, + LinearOrMultiIndex index, Literal* literal) { - // Check that linear_index is in range. - if (linear_index >= ShapeUtil::ElementsIn(literal->shape())) { - return TokenError( - StrCat("trys to set value ", value, " to a literal in shape ", - ShapeUtil::HumanString(literal->shape()), " at linear index ", - linear_index, ", but the index is out of range")); + if (!CheckParsedValueIsInRange(loc, value)) { + return false; } - if (std::isnan(value) || - (std::numeric_limits::has_infinity && - (std::numeric_limits::infinity() == value || - -std::numeric_limits::infinity() == value))) { - // Skip range checking for non-finite value. - } else if (literal->shape().element_type() == F16 || - literal->shape().element_type() == BF16) { - if (value > kF16max || value < -kF16max) { - return TokenError(StrCat( - "value ", value, " is out of range for literal's primitive type ", - PrimitiveType_Name(literal->shape().element_type()))); + // Check that the index is in range and assign into the literal + if (auto* linear_index = absl::get_if(&index)) { + if (*linear_index >= ShapeUtil::ElementsIn(literal->shape())) { + return Error(loc, StrCat("trys to set value ", StringifyValue(value), + " to a literal in shape ", + ShapeUtil::HumanString(literal->shape()), + " at linear index ", *linear_index, + ", but the index is out of range")); } - } else if (std::is_unsigned::value) { - CHECK((std::is_same::value || - std::is_same::value)) - << "Unimplemented checking for ParsedElemT"; + literal->data().at(*linear_index) = + static_cast(value); + } else { + auto* multi_index = absl::get_if>(&index); + CHECK(multi_index != nullptr); - ParsedElemT upper_bound; - if (sizeof(LiteralNativeT) >= sizeof(ParsedElemT)) { - upper_bound = std::numeric_limits::max(); - } else { - upper_bound = - static_cast(std::numeric_limits::max()); + auto invalid_idx = [&](string msg) { + return Error(loc, StrFormat("Invalid sparse index [%s]. %s", + absl::StrJoin(*multi_index, ", "), msg)); + }; + + const auto& shape = literal->shape(); + if (shape.rank() != multi_index->size()) { + return invalid_idx( + StrFormat("Has rank %d, but constant has shape %s, which has rank %d", + multi_index->size(), shape.ToString(), shape.rank())); } - if (value > upper_bound || value < 0) { - // Value is out of range for LiteralNativeT. - return TokenError(StrCat( - "value ", value, " is out of range for literal's primitive type ", - PrimitiveType_Name(literal->shape().element_type()))); - } - } else if (value > static_cast( - std::numeric_limits::max()) || - value < static_cast( - std::numeric_limits::lowest())) { - // Value is out of range for LiteralNativeT. - return TokenError(StrCat( - "value ", value, " is out of range for literal's primitive type ", - PrimitiveType_Name(literal->shape().element_type()))); + for (int64 i = 0; i < shape.rank(); ++i) { + auto idx = (*multi_index)[i]; + if (idx < 0) { + return invalid_idx(StrFormat( + "Sub-index value at %d, namely %d, cannot be negative.", i, idx)); + } + if (idx >= shape.dimensions(i)) { + return invalid_idx( + StrFormat("Sub-index at %d, namely %d, doesn't fit within shape " + "dimension %d in %s", + i, idx, shape.dimensions(i), shape.ToString())); + } + } + literal->AppendSparseElement(*multi_index, + static_cast(value)); } - - literal->data().at(linear_index) = - static_cast(value); return true; } @@ -1996,12 +2082,16 @@ bool HloParser::ParseNonTupleLiteral(Literal* literal, const Shape& shape) { } bool HloParser::ParseDenseLiteral(Literal* literal, const Shape& shape) { - const tensorflow::int64 rank = shape.rank(); + // Cast `rank` to int because we call shape.dimensions(int rank) below, and if + // `rank` is an int64, that's an implicit narrowing conversion, which is + // implementation-defined behavior. + const int rank = static_cast(shape.rank()); + // Create a literal with the given shape in default layout. *literal = LiteralUtil::CreateFromDimensions( shape.element_type(), AsInt64Slice(shape.dimensions())); - tensorflow::int64 nest_level = 0; - tensorflow::int64 linear_index = 0; + int64 nest_level = 0; + int64 linear_index = 0; // elems_seen_per_dim[i] is how many elements or sub-arrays we have seen for // the dimension i. For example, to parse f32[2,3] {{1, 2, 3}, {4, 5, 6}}, // when we are parsing the 2nd '{' (right before '1'), we are seeing a @@ -2009,17 +2099,35 @@ bool HloParser::ParseDenseLiteral(Literal* literal, const Shape& shape) { // the first '}' (right after '3'), it means the sub-array ends, and the // sub-array is supposed to contain exactly 3 elements, so check if // elems_seen_per_dim[1] is 3. - std::vector elems_seen_per_dim(rank); + std::vector elems_seen_per_dim(rank); auto get_index_str = [&elems_seen_per_dim](int dim) -> string { - std::vector elems_seen_until_dim( - elems_seen_per_dim.begin(), elems_seen_per_dim.begin() + dim); + std::vector elems_seen_until_dim(elems_seen_per_dim.begin(), + elems_seen_per_dim.begin() + dim); return StrCat("[", StrJoin(elems_seen_until_dim, ",", - [](string* out, const tensorflow::int64& num_elems) { + [](string* out, const int64& num_elems) { StrAppend(out, num_elems - 1); }), "]"); }; + + auto add_one_elem_seen = [&] { + if (rank > 0) { + if (nest_level != rank) { + return TokenError(absl::StrFormat( + "expects nested array in rank %d, but sees %d", rank, nest_level)); + } + elems_seen_per_dim[rank - 1]++; + if (elems_seen_per_dim[rank - 1] > shape.dimensions(rank - 1)) { + return TokenError(absl::StrFormat( + "expects %d elements on the minor-most dimension, but " + "sees more", + shape.dimensions(rank - 1))); + } + } + return true; + }; + do { switch (lexer_.GetKind()) { default: @@ -2055,6 +2163,31 @@ bool HloParser::ParseDenseLiteral(Literal* literal, const Shape& shape) { lexer_.Lex(); break; } + case TokKind::kLparen: { + if (!primitive_util::IsComplexType(shape.element_type())) { + return TokenError( + absl::StrFormat("unexpected '(' in literal. Parens are only " + "valid for complex literals")); + } + + std::complex value; + LocTy loc = lexer_.GetLoc(); + if (!add_one_elem_seen() || !ParseComplex(&value) || + !SetValueInLiteral(loc, value, linear_index++, literal)) { + return false; + } + break; + } + case TokKind::kDots: { + if (nest_level != 1) { + return TokenError(absl::StrFormat( + "expects `...` at nest level 1, but sees it at nest level %d", + nest_level)); + } + elems_seen_per_dim[0] = shape.dimensions(0); + lexer_.Lex(); + break; + } case TokKind::kComma: // Skip. lexer_.Lex(); @@ -2066,23 +2199,11 @@ bool HloParser::ParseDenseLiteral(Literal* literal, const Shape& shape) { case TokKind::kw_nan: case TokKind::kw_inf: case TokKind::kNegInf: { - if (rank > 0) { - if (nest_level != rank) { - return TokenError( - absl::StrFormat("expects nested array in rank %d, but sees %d", - rank, nest_level)); - } - elems_seen_per_dim[rank - 1]++; - if (elems_seen_per_dim[rank - 1] > shape.dimensions(rank - 1)) { - return TokenError(absl::StrFormat( - "expects %d elements on the minor-most dimension, but " - "sees more", - shape.dimensions(rank - 1))); - } - } + add_one_elem_seen(); if (lexer_.GetKind() == TokKind::kw_true || lexer_.GetKind() == TokKind::kw_false) { - if (!SetValueInLiteral(lexer_.GetKind() == TokKind::kw_true, + if (!SetValueInLiteral(lexer_.GetLoc(), + lexer_.GetKind() == TokKind::kw_true, linear_index++, literal)) { return false; } @@ -2090,12 +2211,12 @@ bool HloParser::ParseDenseLiteral(Literal* literal, const Shape& shape) { } else if (primitive_util::IsIntegralType(shape.element_type()) || shape.element_type() == PRED) { LocTy loc = lexer_.GetLoc(); - tensorflow::int64 value; + int64 value; if (!ParseInt64(&value)) { return Error(loc, StrCat("expects integer for primitive type: ", PrimitiveType_Name(shape.element_type()))); } - if (!SetValueInLiteral(value, linear_index++, literal)) { + if (!SetValueInLiteral(loc, value, linear_index++, literal)) { return false; } } else if (primitive_util::IsFloatingPointType(shape.element_type())) { @@ -2106,7 +2227,7 @@ bool HloParser::ParseDenseLiteral(Literal* literal, const Shape& shape) { loc, StrCat("expect floating point value for primitive type: ", PrimitiveType_Name(shape.element_type()))); } - if (!SetValueInLiteral(value, linear_index++, literal)) { + if (!SetValueInLiteral(loc, value, linear_index++, literal)) { return false; } } else { @@ -2123,48 +2244,7 @@ bool HloParser::ParseDenseLiteral(Literal* literal, const Shape& shape) { } bool HloParser::ParseSparseLiteral(Literal* literal, const Shape& shape) { - switch (shape.element_type()) { - case PRED: - return ParseSparseLiteralHelper(literal, shape); - case S8: - return ParseSparseLiteralHelper(literal, shape); - case S16: - return ParseSparseLiteralHelper(literal, shape); - case S32: - return ParseSparseLiteralHelper(literal, shape); - case S64: - return ParseSparseLiteralHelper(literal, shape); - case U8: - return ParseSparseLiteralHelper(literal, shape); - case U16: - return ParseSparseLiteralHelper(literal, shape); - case U32: - return ParseSparseLiteralHelper(literal, shape); - case U64: - return ParseSparseLiteralHelper(literal, shape); - case F16: - return ParseSparseLiteralHelper(literal, shape); - case F32: - return ParseSparseLiteralHelper(literal, shape); - case BF16: - return ParseSparseLiteralHelper(literal, shape); - case F64: - return ParseSparseLiteralHelper(literal, shape); - default: - return Error(lexer_.GetLoc(), - StrCat("invalid primitive type for sparse literal: ", - PrimitiveType_Name(shape.element_type()))); - } -} - -template -bool HloParser::ParseSparseLiteralHelper(Literal* literal, const Shape& shape) { - std::vector index; - - tensorflow::int64 rank = shape.rank(); - *literal = Literal(shape); - if (!ParseToken(TokKind::kLbrace, "expects '{' at the beginning of a sparse literal")) { return false; @@ -2176,61 +2256,66 @@ bool HloParser::ParseSparseLiteralHelper(Literal* literal, const Shape& shape) { break; } - LocTy index_loc = lexer_.GetLoc(); - index.clear(); + std::vector index; if (lexer_.GetKind() == TokKind::kInt) { - tensorflow::int64 single_index = lexer_.GetInt64Val(); + int64 single_index = lexer_.GetInt64Val(); lexer_.Lex(); - if (rank != 1) { - return Error( - index_loc, - StrCat("invalid single-dimensional index for shape with rank ", - rank, ": ", single_index)); - } index.push_back(single_index); } else { if (!ParseInt64List(TokKind::kLsquare, TokKind::kRsquare, TokKind::kComma, &index)) { return false; } - if (index.size() != rank) { - return Error( - index_loc, - StrCat("invalid multi-dimension index for shape with rank ", rank, - ": [", StrJoin(index, ", "), "]")); - } } if (!ParseToken(TokKind::kColon, "expects ':' after after the sparse array index and before " "the sparse array value")) { return false; } + LocTy value_loc = lexer_.GetLoc(); - LiteralNativeT value; if (lexer_.GetKind() == TokKind::kw_true || lexer_.GetKind() == TokKind::kw_false) { - value = static_cast(lexer_.GetKind() == TokKind::kw_true); + bool value = lexer_.GetKind() == TokKind::kw_true; + if (!SetValueInLiteral(lexer_.GetLoc(), value, index, literal)) { + return false; + } lexer_.Lex(); } else if (primitive_util::IsIntegralType(shape.element_type())) { - tensorflow::int64 value_s64; - if (!ParseInt64(&value_s64)) { + int64 value; + if (!ParseInt64(&value)) { return Error(value_loc, StrCat("expects integer for primitive type: ", PrimitiveType_Name(shape.element_type()))); } - value = static_cast(value_s64); + if (!SetValueInLiteral(value_loc, value, index, literal)) { + return false; + } } else if (primitive_util::IsFloatingPointType(shape.element_type())) { - double value_f64; - if (!ParseDouble(&value_f64)) { + double value; + if (!ParseDouble(&value)) { return Error(value_loc, StrCat("expects floating point value for primitive type: ", PrimitiveType_Name(shape.element_type()))); } - value = static_cast(value_f64); + if (!SetValueInLiteral(value_loc, value, index, literal)) { + return false; + } + } else if (primitive_util::IsComplexType(shape.element_type())) { + std::complex value; + if (!ParseComplex(&value)) { + return Error(value_loc, + StrCat("expects complex value for primitive type: ", + PrimitiveType_Name(shape.element_type()))); + } + if (!SetValueInLiteral(value_loc, value, index, literal)) { + return false; + } } else { LOG(FATAL) << "Unexpected element type: " << PrimitiveType_Name(shape.element_type()); } + if (lexer_.GetKind() != TokKind::kRbrace && !ParseToken(TokKind::kComma, "expects ',' separator between sparse array elements")) { @@ -2244,14 +2329,114 @@ bool HloParser::ParseSparseLiteralHelper(Literal* literal, const Shape& shape) { StrCat("number of sparse elements exceeds maximum for layout: ", ShapeUtil::HumanStringWithLayout(shape))); } - - literal->AppendSparseElement(index, value); } literal->SortSparseElements(); return true; } +// MaxFiniteValue is a type-traits helper used by +// HloParser::CheckParsedValueIsInRange. +template +struct MinMaxFiniteValue { + static T max() { return std::numeric_limits::max(); } + static T min() { return std::numeric_limits::lowest(); } +}; + +template <> +struct MinMaxFiniteValue { + static double max() { + // Sadly this is not constexpr, so this forces `value` to be a method. + return static_cast(Eigen::NumTraits::highest()); + } + static double min() { return -max(); } +}; + +template <> +struct MinMaxFiniteValue { + static double max() { return static_cast(bfloat16::highest()); } + static double min() { return -max(); } +}; + +template +bool HloParser::CheckParsedValueIsInRange(LocTy loc, ParsedElemT value) { + PrimitiveType literal_ty = + primitive_util::NativeToPrimitiveType(); + if (std::isnan(value) || + (std::numeric_limits::has_infinity && + (std::numeric_limits::infinity() == value || + -std::numeric_limits::infinity() == value))) { + // Skip range checking for non-finite value. + } else if (std::is_unsigned::value) { + CHECK((std::is_same::value || + std::is_same::value)) + << "Unimplemented checking for ParsedElemT"; + + ParsedElemT upper_bound; + if (sizeof(LiteralNativeT) >= sizeof(ParsedElemT)) { + upper_bound = std::numeric_limits::max(); + } else { + upper_bound = + static_cast(std::numeric_limits::max()); + } + if (value > upper_bound || value < 0) { + // Value is out of range for LiteralNativeT. + return Error(loc, StrCat("value ", value, + " is out of range for literal's primitive type ", + PrimitiveType_Name(literal_ty), " namely [0, ", + upper_bound, "].")); + } + } else if (value > MinMaxFiniteValue::max() || + value < MinMaxFiniteValue::min()) { + // Value is out of range for LiteralNativeT. + return Error(loc, StrCat("value ", value, + " is out of range for literal's primitive type ", + PrimitiveType_Name(literal_ty), " namely [", + MinMaxFiniteValue::min(), ", ", + MinMaxFiniteValue::max(), "].")); + } + return true; +} + +template +bool HloParser::CheckParsedValueIsInRange(LocTy loc, + std::complex value) { + // e.g. `float` for std::complex + using LiteralComplexComponentT = + decltype(std::real(std::declval())); + + // We could do simply + // + // return CheckParsedValueIsInRange(std::real(value)) && + // CheckParsedValueIsInRange(std::imag(value)); + // + // but this would give bad error messages on failure. + + auto check_component = [&](absl::string_view name, double v) { + if (std::isnan(v) || v == std::numeric_limits::infinity() || + v == -std::numeric_limits::infinity()) { + // Skip range-checking for non-finite values. + return true; + } + + double min = MinMaxFiniteValue::min(); + double max = MinMaxFiniteValue::max(); + if (v < min || v > max) { + // Value is out of range for LitearlComplexComponentT. + return Error( + loc, + StrCat(name, " part ", v, + " is out of range for literal's primitive type ", + PrimitiveType_Name( + primitive_util::NativeToPrimitiveType()), + ", namely [", min, ", ", max, "].")); + } + return true; + }; + return check_component("real", std::real(value)) && + check_component("imaginary", std::imag(value)); +} + // operands ::= '(' operands1 ')' // operands1 // ::= /*empty*/ @@ -2409,24 +2594,23 @@ bool HloParser::ParseAttributeHelper( return true; } case AttrTy::kInt64: { - tensorflow::int64 result; + int64 result; if (!ParseInt64(&result)) { return false; } - static_cast*>(attr_out_ptr) - ->emplace(result); + static_cast*>(attr_out_ptr)->emplace(result); return true; } case AttrTy::kInt32: { - tensorflow::int64 result; + int64 result; if (!ParseInt64(&result)) { return false; } - if (result != static_cast(result)) { + if (result != static_cast(result)) { return Error(attr_loc, "value out of range for int32"); } - static_cast*>(attr_out_ptr) - ->emplace(static_cast(result)); + static_cast*>(attr_out_ptr) + ->emplace(static_cast(result)); return true; } case AttrTy::kFloat: { @@ -2466,6 +2650,15 @@ bool HloParser::ParseAttributeHelper( static_cast*>(attr_out_ptr)->emplace(result); return true; } + case AttrTy::kTriangularSolveTranspose: { + TriangularSolveOptions::Transpose result; + if (!ParseTriangularSolveTranspose(&result)) { + return false; + } + static_cast*>(attr_out_ptr) + ->emplace(result); + return true; + } case AttrTy::kWindow: { Window result; if (!ParseWindow(&result, /*expect_outer_curlies=*/true)) { @@ -2510,19 +2703,19 @@ bool HloParser::ParseAttributeHelper( return true; } case AttrTy::kBracedInt64List: { - std::vector result; + std::vector result; if (!ParseInt64List(TokKind::kLbrace, TokKind::kRbrace, TokKind::kComma, &result)) { return false; } - static_cast>*>(attr_out_ptr) + static_cast>*>(attr_out_ptr) ->emplace(result); return true; } case AttrTy::kBracedInt64ListList: { - std::vector> result; + std::vector> result; auto parse_and_add_item = [&]() { - std::vector item; + std::vector item; if (!ParseInt64List(TokKind::kLbrace, TokKind::kRbrace, TokKind::kComma, &item)) { return false; @@ -2534,8 +2727,7 @@ bool HloParser::ParseAttributeHelper( parse_and_add_item)) { return false; } - static_cast>>*>( - attr_out_ptr) + static_cast>>*>(attr_out_ptr) ->emplace(result); return true; } @@ -2736,7 +2928,7 @@ bool HloParser::ParseConvolutionDimensionNumbers( absl::string_view rhs = split2[0]; absl::string_view out = split2[1]; - const tensorflow::int64 rank = lhs.length(); + const int64 rank = lhs.length(); if (rank != rhs.length() || rank != out.length()) { return TokenError( "convolution lhs, rhs, and output must have the same rank"); @@ -2746,7 +2938,7 @@ bool HloParser::ParseConvolutionDimensionNumbers( } auto is_unique = [](string str) -> bool { - std::sort(str.begin(), str.end()); + absl::c_sort(str); return std::unique(str.begin(), str.end()) == str.end(); }; @@ -2847,7 +3039,7 @@ bool HloParser::ParseSliceRanges(SliceRanges* result) { if (!ParseToken(TokKind::kLbrace, "expects '{' to start ranges")) { return false; } - std::vector> ranges; + std::vector> ranges; if (lexer_.GetKind() == TokKind::kRbrace) { // empty return ParseToken(TokKind::kRbrace, "expects '}' to end ranges"); @@ -2917,9 +3109,9 @@ bool HloParser::ParseShapeList(std::vector* result) { // ::= int64_val (delim int64_val)* bool HloParser::ParseInt64List(const TokKind start, const TokKind end, const TokKind delim, - std::vector* result) { + std::vector* result) { auto parse_and_add_item = [&]() { - tensorflow::int64 i; + int64 i; if (!ParseInt64(&i)) { return false; } @@ -2995,7 +3187,7 @@ bool HloParser::ParseParamList() { bool HloParser::ParseDimensionSizes(std::vector* dimension_sizes, std::vector* dynamic_dimensions) { auto parse_and_add_item = [&]() { - tensorflow::int64 i; + int64 i; bool is_dynamic = false; if (lexer_.GetKind() == TokKind::kLeq) { is_dynamic = true; @@ -3012,22 +3204,108 @@ bool HloParser::ParseDimensionSizes(std::vector* dimension_sizes, parse_and_add_item); } -// layout ::= '{' int64_list '}' +// tiles +// ::= /*empty*/ +// ::= 'T' '(' dim_list ')' +// dim_list +// ::= /*empty*/ +// ::= (int64 | '*') (',' (int64 | '*'))* +bool HloParser::ParseTiles(std::vector* tiles) { + auto parse_and_add_tile_dimension = [&]() { + tensorflow::int64 i; + if (ParseInt64(&i)) { + tiles->back().add_dimensions(i); + return true; + } + if (lexer_.GetKind() == TokKind::kAsterisk) { + tiles->back().add_dimensions(Tile::kCombineDimension); + lexer_.Lex(); + return true; + } + return false; + }; + + do { + tiles->push_back(Tile()); + if (!ParseList(TokKind::kLparen, TokKind::kRparen, TokKind::kComma, + parse_and_add_tile_dimension)) { + return false; + } + } while (lexer_.GetKind() == TokKind::kLparen); + return true; +} + +// layout ::= '{' int64_list (':' tiles element_size_in_bits)? '}' +// element_size_in_bits +// ::= /*empty*/ +// ::= 'E' '(' int64 ')' bool HloParser::ParseLayout(Layout* layout) { std::vector minor_to_major; + std::vector tiles; + tensorflow::int64 element_size_in_bits = 0; + auto parse_and_add_item = [&]() { - tensorflow::int64 i; + int64 i; if (!ParseInt64(&i)) { return false; } minor_to_major.push_back(i); return true; }; - if (!ParseList(TokKind::kLbrace, TokKind::kRbrace, TokKind::kComma, - parse_and_add_item)) { + + if (!ParseToken(TokKind::kLbrace, + StrCat("expects layout to start with ", + TokKindToString(TokKind::kLbrace)))) { + return false; + } + if (lexer_.GetKind() != TokKind::kRbrace) { + if (lexer_.GetKind() == TokKind::kInt) { + // Parse minor to major. + do { + if (!parse_and_add_item()) { + return false; + } + } while (EatIfPresent(TokKind::kComma)); + } + + if (lexer_.GetKind() == TokKind::kColon) { + lexer_.Lex(); + if (lexer_.GetKind() == TokKind::kIdent && lexer_.GetStrVal() == "T") { + lexer_.Lex(); + ParseTiles(&tiles); + } + + if (lexer_.GetKind() == TokKind::kIdent && lexer_.GetStrVal() == "E") { + // Parse element size in bits. + lexer_.Lex(); + if (!ParseToken(TokKind::kLparen, + StrCat("expects element size in bits to start with ", + TokKindToString(TokKind::kLparen)))) { + return false; + } + if (!ParseInt64(&element_size_in_bits)) { + return false; + } + if (!ParseToken(TokKind::kRparen, + StrCat("expects element size in bits to end with ", + TokKindToString(TokKind::kRparen)))) { + return false; + } + } + } + } + if (!ParseToken(TokKind::kRbrace, + StrCat("expects layout to end with ", + TokKindToString(TokKind::kRbrace)))) { return false; } - *layout = LayoutUtil::MakeLayout(minor_to_major); + + std::vector vec_tiles(tiles.size()); + for (int i = 0; i < tiles.size(); i++) { + vec_tiles[i] = Tile(tiles[i]); + } + *layout = + LayoutUtil::MakeLayout(minor_to_major, vec_tiles, element_size_in_bits); return true; } @@ -3079,7 +3357,7 @@ bool HloParser::ParseShape(Shape* result) { lexer_.Lex(); const string message = "expects a brace-bracketed integer for sparse layout"; - tensorflow::int64 max_sparse_elements; + int64 max_sparse_elements; if (!ParseToken(TokKind::kLbrace, message) || !ParseInt64(&max_sparse_elements) || !ParseToken(TokKind::kRbrace, message)) { @@ -3099,13 +3377,20 @@ bool HloParser::ParseShape(Shape* result) { // // The open brace could either be the start of a computation or the start of a // layout for the f32[123] shape. We consider it the start of a layout if the - // next token after the open brace is a integer + // next token after the open brace is an integer or a colon. if (lexer_.GetKind() == TokKind::kLbrace && - lexer_.LookAhead() == TokKind::kInt) { + (lexer_.LookAhead() == TokKind::kInt || + lexer_.LookAhead() == TokKind::kColon)) { Layout layout; if (!ParseLayout(&layout)) { return false; } + if (layout.minor_to_major_size() != result->rank()) { + return Error( + lexer_.GetLoc(), + StrFormat("Dimensions size is %ld, but minor to major size is %ld.", + result->rank(), layout.minor_to_major_size())); + } *result->mutable_layout() = layout; } return true; @@ -3148,15 +3433,14 @@ bool HloParser::ParseString(string* result) { return true; } -bool HloParser::ParseDxD(const string& name, - std::vector* result) { +bool HloParser::ParseDxD(const string& name, std::vector* result) { LocTy loc = lexer_.GetLoc(); if (!result->empty()) { return Error(loc, StrFormat("sub-attribute '%s=' already exists", name)); } // 1D if (lexer_.GetKind() == TokKind::kInt) { - tensorflow::int64 number; + int64 number; if (!ParseInt64(&number)) { return Error(loc, StrFormat("expects sub-attribute '%s=i'", name)); } @@ -3175,8 +3459,7 @@ bool HloParser::ParseDxD(const string& name, return TokenError("expects token type kInt or kDxD"); } -bool HloParser::ParseWindowPad( - std::vector>* pad) { +bool HloParser::ParseWindowPad(std::vector>* pad) { LocTy loc = lexer_.GetLoc(); if (!pad->empty()) { return Error(loc, "sub-attribute 'pad=' already exists"); @@ -3186,7 +3469,7 @@ bool HloParser::ParseWindowPad( } string str = lexer_.GetStrVal(); for (const auto& padding_dim_str : absl::StrSplit(str, 'x')) { - std::vector low_high; + std::vector low_high; if (!SplitToInt64s(padding_dim_str, '_', &low_high) || low_high.size() != 2) { return Error(loc, @@ -3209,7 +3492,7 @@ bool HloParser::ParsePaddingConfig(PaddingConfig* padding) { LocTy loc = lexer_.GetLoc(); string str = lexer_.GetStrVal(); for (const auto& padding_dim_str : absl::StrSplit(str, 'x')) { - std::vector padding_dim; + std::vector padding_dim; if (!SplitToInt64s(padding_dim_str, '_', &padding_dim) || (padding_dim.size() != 2 && padding_dim.size() != 3)) { return Error(loc, @@ -3231,7 +3514,7 @@ bool HloParser::ParseMetadata(OpMetadata* metadata) { optional op_type; optional op_name; optional source_file; - optional source_line; + optional source_line; attrs["op_type"] = {/*required=*/false, AttrTy::kString, &op_type}; attrs["op_name"] = {/*required=*/false, AttrTy::kString, &op_name}; attrs["source_file"] = {/*required=*/false, AttrTy::kString, &source_file}; @@ -3283,6 +3566,22 @@ 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) { @@ -3334,7 +3633,7 @@ bool HloParser::ParsePrecision(PrecisionConfig::Precision* result) { return true; } -bool HloParser::ParseInt64(tensorflow::int64* result) { +bool HloParser::ParseInt64(int64* result) { VLOG(1) << "ParseInt64"; if (lexer_.GetKind() != TokKind::kInt) { return TokenError("expects integer"); @@ -3346,9 +3645,18 @@ bool HloParser::ParseInt64(tensorflow::int64* result) { bool HloParser::ParseDouble(double* result) { switch (lexer_.GetKind()) { - case TokKind::kDecimal: - *result = lexer_.GetDecimalVal(); + case TokKind::kDecimal: { + double val = lexer_.GetDecimalVal(); + // If GetDecimalVal returns +/-inf, that means that we overflowed + // `double`. + if (std::isinf(val)) { + return TokenError(StrCat("Constant is out of range for double (+/-", + std::numeric_limits::max(), + ") and so is unparsable.")); + } + *result = val; break; + } case TokKind::kInt: *result = static_cast(lexer_.GetInt64Val()); break; @@ -3368,6 +3676,42 @@ bool HloParser::ParseDouble(double* result) { return true; } +bool HloParser::ParseComplex(std::complex* result) { + if (lexer_.GetKind() != TokKind::kLparen) { + return TokenError("expects '(' before complex number"); + } + lexer_.Lex(); + + double real; + LocTy loc = lexer_.GetLoc(); + if (!ParseDouble(&real)) { + return Error(loc, + "expect floating-point value for real part of complex number"); + } + + if (lexer_.GetKind() != TokKind::kComma) { + return TokenError( + absl::StrFormat("expect comma after real part of complex literal")); + } + lexer_.Lex(); + + double imag; + loc = lexer_.GetLoc(); + if (!ParseDouble(&imag)) { + return Error( + loc, + "expect floating-point value for imaginary part of complex number"); + } + + if (lexer_.GetKind() != TokKind::kRparen) { + return TokenError(absl::StrFormat("expect ')' after complex number")); + } + + *result = std::complex(real, imag); + lexer_.Lex(); + return true; +} + bool HloParser::ParseBool(bool* result) { if (lexer_.GetKind() != TokKind::kw_true && lexer_.GetKind() != TokKind::kw_false) { diff --git a/tensorflow/compiler/xla/service/hlo_parser_test.cc b/tensorflow/compiler/xla/service/hlo_parser_test.cc index 76b8a5bc117e653b3b2c09c7c95f26f1d4f27c7b..203a7dba22110063b54467bd8e550fa8f23c68d1 100644 --- a/tensorflow/compiler/xla/service/hlo_parser_test.cc +++ b/tensorflow/compiler/xla/service/hlo_parser_test.cc @@ -551,6 +551,17 @@ ENTRY %Transpose.v2 () -> s32[1,2,3] { ROOT %transpose = s32[1,2,3]{2,1,0} transpose(s32[1,2,3]{2,1,0} %constant), dimensions={0,1,2} } +)" +}, +{ +"TransposeC128", +R"(HloModule TransposeC128_module + +ENTRY %Transpose.v3 (input: c128[1,2,3]) -> c128[1,2,3] { + %input = c128[1,2,3]{2,1,0} parameter(0) + ROOT %transpose = c128[1,2,3]{2,1,0} transpose(c128[1,2,3]{2,1,0} %input), dimensions={0,1,2} +} + )" }, // Dynamic slice @@ -771,7 +782,17 @@ ENTRY %fusion.v3 () -> f32[3,2,1,1] { R"(HloModule sparse_f32 ENTRY %sparse () -> f32[2,3,4] { - ROOT %foo = f32[2,3,4]sparse{10} constant({[0, 1, 2]: 1, [1, 2, 3]: 2, [2, 3, 4]: 3}) + ROOT %foo = f32[2,3,4]sparse{10} constant({[0, 1, 2]: 1, [1, 2, 2]: 2, [1, 2, 3]: 3}) +} + +)" +}, +{ +"SparseC128", +R"(HloModule sparse_c128 + +ENTRY %sparse () -> c128[2,3,4] { + ROOT %foo = c128[2,3,4]sparse{10} constant({[0, 1, 2]: (1, 0), [1, 2, 2]: (2, 5), [1, 2, 3]: (3, 10)}) } )" @@ -883,6 +904,28 @@ ENTRY %CustomCallWithLayoutConstraints (p0: (f32[2,2], f32[42,2,3]), p1: f32[123 ROOT %custom-call = (f32[1,2,3]{0,2,1}, f32[1,2,3]{1,2,0}) custom-call((f32[2,2]{0,1}, f32[42,2,3]{0,1,2}) %p0, f32[123,4]{0,1} %p1), custom_call_target="baz", operand_layout_constraints={(f32[2,2]{1,0}, f32[42,2,3]{2,0,1}), f32[123,4]{1,0}} } +)" +}, +// Parse c64 literal +{ +"ParseC64Literal", +R"(HloModule ParseC64Literal + +ENTRY %ParseC64Literal () -> c64[2] { + ROOT %c = c64[2]{0} constant({(1, 2), (-inf, nan)}) +} + +)" +}, +// Parse c128 literal +{ +"ParseC128Literal", +R"(HloModule ParseC128Literal + +ENTRY %ParseC128Literal () -> c128[2] { + ROOT %c = c128[2]{0} constant({(1, 2), (-inf, nan)}) +} + )" }, }); @@ -1004,9 +1047,15 @@ ENTRY ReducePrecision { "SortKey", R"(HloModule sort +compare { + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + ROOT lt = pred[] less-than(p.0.lhs, p.0.rhs) +} + ENTRY Sort { x = f32[1024]{0} parameter(0) - ROOT sorted = f32[1024]{0} sort(x), dimensions={0} + ROOT sorted = f32[1024]{0} sort(x), dimensions={0}, to_apply=compare } )" @@ -1016,10 +1065,18 @@ ENTRY Sort { "SortKeyValue", R"(HloModule sort +compare { + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + ROOT lt = pred[] less-than(p.0.lhs, p.0.rhs) +} + ENTRY Sort { keys = f32[1024]{0} parameter(0) values = s32[1024]{0} parameter(1) - ROOT sorted = (f32[1024]{0}, s32[1024]{0}) sort(keys, values), dimensions={0} + ROOT sorted = (f32[1024]{0}, s32[1024]{0}) sort(keys, values), dimensions={0}, to_apply=compare } )" @@ -1029,9 +1086,15 @@ ENTRY Sort { "SortKeyR2", R"(HloModule sort +compare { + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + ROOT lt = pred[] less-than(p.0.lhs, p.0.rhs) +} + ENTRY Sort { x = f32[1024,16]{0,1} parameter(0) - ROOT sorted = f32[1024,16]{0,1} sort(x), dimensions={0} + ROOT sorted = f32[1024,16]{0,1} sort(x), dimensions={0}, to_apply=compare } )" @@ -1041,10 +1104,18 @@ ENTRY Sort { "SortKeyValueR2", R"(HloModule sort +compare { + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + ROOT lt = pred[] less-than(p.0.lhs, p.0.rhs) +} + ENTRY Sort { keys = f32[1024,16]{0,1} parameter(0) values = s32[1024,16]{0,1} parameter(1) - ROOT sorted = (f32[1024,16]{0,1}, s32[1024,16]{0,1}) sort(keys, values), dimensions={0} + ROOT sorted = (f32[1024,16]{0,1}, s32[1024,16]{0,1}) sort(keys, values), dimensions={0}, to_apply=compare } )" @@ -1054,12 +1125,24 @@ ENTRY Sort { "SortManyValues", R"(HloModule sort +compare { + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + p.2.lhs = u32[] parameter(4) + p.2.rhs = u32[] parameter(5) + p.3.lhs = f32[] parameter(6) + p.3.rhs = f32[] parameter(7) + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + ROOT lt = pred[] less-than(p.0.lhs, p.0.rhs) +} + ENTRY Sort { keys = f32[1024,16]{0,1} parameter(0) values.0 = s32[1024,16]{0,1} parameter(1) values.1 = u32[1024,16]{0,1} parameter(2) values.2 = f32[1024,16]{0,1} parameter(3) - ROOT sorted = (f32[1024,16]{0,1}, s32[1024,16]{0,1}, u32[1024,16]{0,1}, f32[1024,16]{0,1}) sort(keys, values.0, values.1, values.2), dimensions={0} + ROOT sorted = (f32[1024,16]{0,1}, s32[1024,16]{0,1}, u32[1024,16]{0,1}, f32[1024,16]{0,1}) sort(keys, values.0, values.1, values.2), dimensions={0}, to_apply=compare } )" @@ -1237,6 +1320,17 @@ ENTRY CollectivePermute { ROOT root = f32[128,32]{0,1} collective-permute(input), source_target_pairs={{0,1},{1,2},{2,3}} } +)" +}, +// replica-id +{ +"ReplicaId", +R"(HloModule replica-id + +ENTRY Replica-id { + ROOT replica-id = u32[] replica-id() +} + )" }, // Iota @@ -1266,10 +1360,18 @@ ENTRY Computation { "ScheduledModule", R"(HloModule scheduled_module, is_scheduled=true +compare { + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + ROOT lhs = pred[] less-than(p.0.lhs, p.0.rhs) +} + ENTRY Sort { keys = f32[1024]{0} parameter(0) values = s32[1024]{0} parameter(1) - ROOT sorted = (f32[1024]{0}, s32[1024]{0}) sort(keys, values), dimensions={0} + ROOT sorted = (f32[1024]{0}, s32[1024]{0}) sort(keys, values), dimensions={0}, to_apply=compare } )" @@ -1303,6 +1405,30 @@ ENTRY AddDependency { ROOT sum = f32[] add(neg, exp) } +)" +}, + +// A module containing constants equal to the min/max values of various data +// types. +{ +"MinMaxValues", +R"(HloModule MinMaxValues + +ENTRY MinMaxValues { + x.s8 = s8[2]{0} constant({-128, 127}) + x.s16 = s16[2]{0} constant({-32768, 32767}) + x.s32 = s32[2]{0} constant({-2147483648, 2147483647}) + x.u8 = u8[2]{0} constant({0, 255}) + x.u16 = u16[2]{0} constant({0, 65535}) + x.u32 = u32[2]{0} constant({0, 4294967295}) + x.f16 = f16[2]{0} constant({-65504, 65504}) + x.bf16 = bf16[2]{0} constant({-3.38953e+38, 3.38953e+38}) + x.f32 = f32[2]{0} constant({-3.40282e+38, 3.40282e+38}) + x.f64 = f64[2]{0} constant({-1.79769e+308, 1.79769e+308}) + x.c64 = c64[2]{0} constant({(-3.40282e+38, 3.40282e+38), (3.40282e+38, -3.40282e+38)}) + ROOT c.c128 = c128[2]{0} constant({(-1.79769e+308, 1.79769e+308), (1.79769e+308, -1.79769e+308)}) +} + )" }, }); @@ -1329,7 +1455,7 @@ class HloParameterizedParserTest protected: // Expects "ToString(ParseHloString(string)) == string", that is, parses the // string, asserts that it succeeded, stringifies the parsed module, and - // checks that the it equals the original string. + // checks that it equals the original string. void ExpectEqual() { const string& original = GetParam().module_string; TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, @@ -1360,20 +1486,20 @@ TEST_P(HloParserTestLongProto, Run) { ExpectEqual(); } TEST_P(HloParserTestShort, Run) { ExpectEqual(); } TEST_P(HloParserTestShortProto, Run) { ExpectEqual(); } -INSTANTIATE_TEST_CASE_P(HloParserTestSuccessInstantiation, HloParserTestLong, - ::testing::ValuesIn(CreateTestCases()), - TestDataToString); -INSTANTIATE_TEST_CASE_P(HloParserTestSuccessInstantiation, - HloParserTestLongProto, - ::testing::ValuesIn(CreateTestCases()), - TestDataToString); -INSTANTIATE_TEST_CASE_P(HloParserTestSuccessInstantiation, HloParserTestShort, - ::testing::ValuesIn(CreateShortTestCases()), - TestDataToString); -INSTANTIATE_TEST_CASE_P(HloParserTestSuccessInstantiation, - HloParserTestShortProto, - ::testing::ValuesIn(CreateShortTestCases()), - TestDataToString); +INSTANTIATE_TEST_SUITE_P(HloParserTestSuccessInstantiation, HloParserTestLong, + ::testing::ValuesIn(CreateTestCases()), + TestDataToString); +INSTANTIATE_TEST_SUITE_P(HloParserTestSuccessInstantiation, + HloParserTestLongProto, + ::testing::ValuesIn(CreateTestCases()), + TestDataToString); +INSTANTIATE_TEST_SUITE_P(HloParserTestSuccessInstantiation, HloParserTestShort, + ::testing::ValuesIn(CreateShortTestCases()), + TestDataToString); +INSTANTIATE_TEST_SUITE_P(HloParserTestSuccessInstantiation, + HloParserTestShortProto, + ::testing::ValuesIn(CreateShortTestCases()), + TestDataToString); class HloParserTest : public ::testing::Test { protected: @@ -1532,6 +1658,37 @@ ENTRY %ConstantF16Overflow.v4 () -> f16[] { "is out of range for literal's primitive type F16"); } +TEST_F(HloParserTest, ConstantBf16NoOverflow) { + // 65505 is in range for bf16. + const string original = R"( + HloModule test_module + ENTRY test { + ROOT c = bf16[] constant(-65505) + })"; + EXPECT_EQ(Status::OK(), ParseHloString(original).status()); +} + +TEST_F(HloParserTest, ConstantBf16Overflow) { + // 1e100 is out of range for bf16. + const string original = R"( + HloModule test_module + ENTRY test { + ROOT c = bf16[] constant(1e100) + })"; + ExpectHasSubstr(ParseHloString(original).status().error_message(), + "out of range"); +} + +TEST_F(HloParserTest, ConstantF16OverflowInSparseArray) { + const string original = R"( + HloModule test_module + ENTRY test { + ROOT c = f16[5]sparse{10} constant({[0]: 0, [1]: -65505}) + })"; + ExpectHasSubstr(ParseHloString(original).status().error_message(), + "is out of range for literal's primitive type F16"); +} + TEST_F(HloParserTest, ConstantUnsignedUnderflow) { const string original = R"( HloModule ConstantUnsignedUnderflow_module @@ -1566,6 +1723,46 @@ TEST_F(HloParserTest, ConstantUnsignedInt64Overflow) { EXPECT_NE(Status::OK(), result.status()); } +TEST_F(HloParserTest, ConstantC64Overflow) { + const string original = R"( + HloModule test_module + ENTRY test () -> c64[] { + ROOT c = c64[] constant((1e100, 0)) + })"; + auto result = ParseHloString(original); + EXPECT_NE(Status::OK(), result.status()); +} + +TEST_F(HloParserTest, ConstantC64Underflow) { + const string original = R"( + HloModule test_module + ENTRY test () -> c64[] { + ROOT c = c64[] constant((0, -1e100)) + })"; + auto result = ParseHloString(original); + EXPECT_NE(Status::OK(), result.status()); +} + +TEST_F(HloParserTest, ConstantF64Overflow) { + const string original = R"( + HloModule test_module + ENTRY test { + ROOT c = f64[] constant(1.8e308) + })"; + auto result = ParseHloString(original); + EXPECT_NE(Status::OK(), result.status()); +} + +TEST_F(HloParserTest, ConstantF64Underflow) { + const string original = R"( + HloModule test_module + ENTRY test { + ROOT c = f64[] constant(-1.8e308) + })"; + auto result = ParseHloString(original); + EXPECT_NE(Status::OK(), result.status()); +} + TEST_F(HloParserTest, ConstantWithExp) { const string original = R"(HloModule ConstantWithExp_module @@ -1581,6 +1778,19 @@ ENTRY %ConstantWithExp.v4 () -> f32[] { // printed as "300". } +TEST_F(HloParserTest, ShortConstant) { + const string original = R"(HloModule ShortCOnstant_module + +ENTRY %ShortConstant.v4 () -> f32[67,89] { + ROOT %constant.1 = f32[67,89]{1,0} constant({...}) +} + +)"; + auto result = ParseHloString(original); + TF_EXPECT_OK(result.status()); + EXPECT_EQ(result.ValueOrDie()->ToString(HloPrintOptions()), original); +} + TEST_F(HloParserTest, AttibutesAnyOrder) { const string original = R"(HloModule any_order_module @@ -2280,6 +2490,46 @@ ENTRY %entrycomp (p: f32[2,2]) -> f32[2,2] { " with the shape of the operand instruction f32[2,2]{1,0}."); } +TEST_F(HloParserTest, OutOfRangeSparseIndex) { + const string original = R"( + HloModule test_module + ENTRY test { + ROOT c = f16[5]sparse{10} constant({[100]: 0}) + })"; + ExpectHasSubstr(ParseHloString(original).status().error_message(), + "Invalid sparse index"); +} + +TEST_F(HloParserTest, NegativeSparseIndex) { + const string original = R"( + HloModule test_module + ENTRY test { + ROOT c = f16[5]sparse{10} constant({-1: 0}) + })"; + ExpectHasSubstr(ParseHloString(original).status().error_message(), + "Invalid sparse index"); +} + +TEST_F(HloParserTest, SparseIndexWithRankTooLarge) { + const string original = R"( + HloModule test_module + ENTRY test { + ROOT c = f16[5]sparse{10} constant({[0, 0]: 0}) + })"; + ExpectHasSubstr(ParseHloString(original).status().error_message(), + "Invalid sparse index"); +} + +TEST_F(HloParserTest, SparseIndexWithRankTooSmall) { + const string original = R"( + HloModule test_module + ENTRY test { + ROOT c = f16[5, 5]sparse{10} constant({[0]: 0}) + })"; + ExpectHasSubstr(ParseHloString(original).status().error_message(), + "Invalid sparse index"); +} + TEST_F(HloParserTest, ParseShapeStringR2F32) { string shape_string = "f32[123,456]"; TF_ASSERT_OK_AND_ASSIGN(Shape actual, ParseShape(shape_string)); @@ -2324,6 +2574,60 @@ TEST_F(HloParserTest, ParseShapeStringWithLayout) { << "actual: " << ShapeUtil::HumanString(actual); } +TEST_F(HloParserTest, ParseShapeStringWithTilingLayout) { + // One tile. + string shape_string = "f32[123,456]{0,1:T(2,128)}"; + TF_ASSERT_OK_AND_ASSIGN(Shape actual, ParseShape(shape_string)); + Shape expected = + ShapeUtil::MakeShapeWithLayout(F32, {123, 456}, {0, 1}, {Tile({2, 128})}); + EXPECT_EQ(expected, actual) + << "expected: " << ShapeUtil::HumanStringWithLayout(expected) + << "actual: " << ShapeUtil::HumanStringWithLayout(actual); + + // Tile with negative dimension size for combining dimensions. + shape_string = "f32[123,456,789]{0,1,2:T(2, * , 128)}"; + TF_ASSERT_OK_AND_ASSIGN(actual, ParseShape(shape_string)); + expected = + ShapeUtil::MakeShapeWithLayout(F32, {123, 456, 789}, {0, 1, 2}, + {Tile({2, Tile::kCombineDimension, 128})}); + EXPECT_EQ(expected, actual) + << "expected: " << ShapeUtil::HumanStringWithLayout(expected) + << "actual: " << ShapeUtil::HumanStringWithLayout(actual); + + // Two tiles. + shape_string = "bf16[123,456,789]{2,1,0:T(2,*,128)(2,1)}"; + TF_ASSERT_OK_AND_ASSIGN(actual, ParseShape(shape_string)); + expected = ShapeUtil::MakeShapeWithLayout( + BF16, {123, 456, 789}, {2, 1, 0}, + {Tile({2, Tile::kCombineDimension, 128}), Tile({2, 1})}); + EXPECT_EQ(expected, actual) + << "expected: " << ShapeUtil::HumanStringWithLayout(expected) + << "actual: " << ShapeUtil::HumanStringWithLayout(actual); + + // Tile with element size in bits. + shape_string = "pred[123,456]{1,0:T(2,128)E(1)}"; + TF_ASSERT_OK_AND_ASSIGN(actual, ParseShape(shape_string)); + expected = ShapeUtil::MakeShapeWithLayout(PRED, {123, 456}, {1, 0}, + {Tile({2, 128})}, 1); + EXPECT_EQ(expected, actual) + << "expected: " << ShapeUtil::HumanStringWithLayout(expected) + << "actual: " << ShapeUtil::HumanStringWithLayout(actual); + + // Element size in bits without tile. + shape_string = "pred[123,456]{1,0:E(1)}"; + TF_ASSERT_OK_AND_ASSIGN(actual, ParseShape(shape_string)); + expected = ShapeUtil::MakeShapeWithLayout(PRED, {123, 456}, {1, 0}, {}, 1); + EXPECT_EQ(expected, actual) + << "expected: " << ShapeUtil::HumanStringWithLayout(expected) + << "actual: " << ShapeUtil::HumanStringWithLayout(actual); + + // Wrong minor_to_major. + shape_string = "f32[123,456,789]{1:T(2, * , 128)}"; + auto result = ParseShape(shape_string); + ExpectHasSubstr(result.status().error_message(), + "Dimensions size is 3, but minor to major size is 1."); +} + TEST_F(HloParserTest, ParseShapeStringWithSparseLayout) { string shape_string = "f32[123,456]sparse{10}"; TF_ASSERT_OK_AND_ASSIGN(Shape actual, ParseShape(shape_string)); @@ -2380,5 +2684,13 @@ TEST_F(HloParserTest, ParseDynamicTuple) { << "actual: " << ShapeUtil::HumanString(actual); } +TEST_F(HloParserTest, NegativeParameterNumber) { + const string hlo_string = "par0 = f32[3,5] parameter(-1)"; + auto result = ParseHloString(hlo_string); + ASSERT_FALSE(result.status().ok()); + EXPECT_THAT(result.status().error_message(), + ::testing::HasSubstr("parameter number must be >= 0")); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_pass_fix.h b/tensorflow/compiler/xla/service/hlo_pass_fix.h index 791b1a97b0b82edf19ff1588fd8d5d996ac0fef4..35dc9c0029f9871334cb500c6b71f0c86ab136d7 100644 --- a/tensorflow/compiler/xla/service/hlo_pass_fix.h +++ b/tensorflow/compiler/xla/service/hlo_pass_fix.h @@ -19,6 +19,7 @@ limitations under the License. #include #include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_module_group.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" @@ -39,9 +40,36 @@ class HloPassFix : public Pass { int64 iteration_count = 0; int64 limit = std::max(static_cast(1000), module->instruction_count()); + VLOG(3) << "Running HloPassFix."; while (changed_this_iteration) { TF_ASSIGN_OR_RETURN(changed_this_iteration, Pass::Run(module)); changed |= changed_this_iteration; + VLOG(3) << "changed_this_iteration: " << changed_this_iteration; + ++iteration_count; + if (iteration_count == limit) { + LOG(ERROR) + << "Unexpectedly high number of iterations in HLO passes (" + << iteration_count + << ")\nIf compilation hangs here, please file a bug with XLA."; + } + } + return changed; + } + + StatusOr RunOnModuleGroup(HloModuleGroup* module_group) override { + bool changed = false; + bool changed_this_iteration = true; + int64 iteration_count = 0; + int64 limit = 1000; + for (const HloModule* module : module_group->modules()) { + limit = std::max(limit, module->instruction_count()); + } + VLOG(3) << "Running HloPassFix."; + while (changed_this_iteration) { + TF_ASSIGN_OR_RETURN(changed_this_iteration, + Pass::RunOnModuleGroup(module_group)); + changed |= changed_this_iteration; + VLOG(3) << "changed_this_iteration: " << changed_this_iteration; ++iteration_count; if (iteration_count == limit) { LOG(ERROR) diff --git a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc index 33ce7e23a82d840676bba5f1ca9c0ffc4433465d..ae8c08cf1d16ad6738962f3be7c1b5512110b1d1 100644 --- a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc +++ b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc @@ -89,7 +89,7 @@ std::vector HloPassPipeline::GetEnabledPasses( std::vector enabled_passes; for (auto& pass : passes_) { - if (disabled_pass_names.count(string(pass->name())) == 0) { + if (!disabled_pass_names.contains(pass->name())) { enabled_passes.push_back(pass.get()); } } diff --git a/tensorflow/compiler/xla/service/hlo_profile_printer.cc b/tensorflow/compiler/xla/service/hlo_profile_printer.cc index 5eb707a957e49d86cdb2f72b72ce750bf29b8fd2..9cc202aa9f5fe5a20a9da05251ea811137ccaadb 100644 --- a/tensorflow/compiler/xla/service/hlo_profile_printer.cc +++ b/tensorflow/compiler/xla/service/hlo_profile_printer.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_profile_printer.h" +#include "absl/algorithm/container.h" #include "absl/strings/str_cat.h" #include "tensorflow/compiler/xla/service/human_readable_profile_builder.h" @@ -34,11 +35,10 @@ string PrintHloProfile(const HloProfilePrinterData& hlo_profile_printer_data, for (const HloComputationInfo& computation_info : hlo_profile_printer_data.computation_infos()) { const auto& instruction_infos = computation_info.instruction_infos(); - bool any_instruction_profiled = - std::any_of(instruction_infos.begin(), instruction_infos.end(), - [&](const HloInstructionInfo& instruction_info) { - return counters[instruction_info.profile_index()] != 0; - }); + bool any_instruction_profiled = absl::c_any_of( + instruction_infos, [&](const HloInstructionInfo& instruction_info) { + return counters[instruction_info.profile_index()] != 0; + }); if (!any_instruction_profiled) { continue; diff --git a/tensorflow/compiler/xla/service/hlo_reachability.cc b/tensorflow/compiler/xla/service/hlo_reachability.cc index edaa4c59e2674e5f165c468059747d3dd2d54218..0fced7f15bdaf1dbe349e3b0fc6ada68393c6512 100644 --- a/tensorflow/compiler/xla/service/hlo_reachability.cc +++ b/tensorflow/compiler/xla/service/hlo_reachability.cc @@ -49,7 +49,7 @@ void HloReachabilityMap::SetReachabilityToUnionHelper( absl::Span inputs, const HloInstruction* instruction, BitVector* bit_vector) { // If instruction is part of inputs, don't reset the bit_vector. - if (std::find(inputs.begin(), inputs.end(), instruction) == inputs.end()) { + if (!absl::c_linear_search(inputs, instruction)) { bit_vector->SetToZero(); } bit_vector->Set(GetIndex(instruction)); diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization.cc b/tensorflow/compiler/xla/service/hlo_rematerialization.cc index ac74e2432f2176e13eaf7d4a1934a50ee89d1042..a175e4643de2ac6ce07ac00da914d7ab7acca541 100644 --- a/tensorflow/compiler/xla/service/hlo_rematerialization.cc +++ b/tensorflow/compiler/xla/service/hlo_rematerialization.cc @@ -57,6 +57,15 @@ using ::tensorflow::strings::HumanReadableNumBytes; // Returns true if the given instruction is rematerializable. bool IsRematerializable(const HloInstruction* instruction) { + if (instruction->opcode() == HloOpcode::kCopy) { + if (LayoutUtil::Equal(instruction->shape().layout(), + instruction->operand(0)->shape().layout())) { + // Don't rematerialize copies added by copy insertion (layout doesn't + // change). + return false; + } + } + // Don't rematerialize instructions with side effects or instructions which // cannot be cloned safely. switch (instruction->opcode()) { @@ -179,7 +188,8 @@ class InstructionList { Item* CreateItem(HloInstruction* inst) { Item* item = new Item; item->instruction = inst; - CHECK(item_map_.insert({inst, item}).second) << "inserting inst twice"; + CHECK(item_map_.insert({inst, item}).second) + << "inserting inst twice " << inst->name(); return item; } @@ -235,8 +245,7 @@ class InstructionList { } // Now scan forwards until we find one of the before_instructions. - while (std::find(before_instructions.begin(), before_instructions.end(), - min_position_item) == before_instructions.end()) { + while (!absl::c_linear_search(before_instructions, min_position_item)) { min_position_item = min_position_item->next; } return InsertBefore(to_insert, min_position_item); @@ -302,7 +311,7 @@ ItemList GetUsers(const InstructionList& instruction_list, // A buffer may be used by the instruction via more than one alias. For // example, a buffer which appears in more than one element of a tuple. Item* user_item = instruction_list.GetItem(user); - if (std::find(users.begin(), users.end(), user_item) == users.end()) { + if (!absl::c_linear_search(users, user_item)) { users.push_back(user_item); } } @@ -418,11 +427,12 @@ class MemoryUsageTracker { // the given uses. Buffer& RematerializeBuffer(const Buffer& original_buffer, Item* remat_item, ItemList&& rematerialized_uses) { - CHECK(original_buffer.defining_instruction->placed); - CHECK(!original_buffer.has_indirect_uses); - CHECK(!original_buffer.live_out); + CHECK(original_buffer.defining_instruction->placed) + << original_buffer.defining_instruction->instruction->name(); + CHECK(!original_buffer.has_indirect_uses) << original_buffer.ToString(); + CHECK(!original_buffer.live_out) << original_buffer.ToString(); for (Item* use : rematerialized_uses) { - CHECK(!use->placed); + CHECK(!use->placed) << use->instruction->name(); } return NewBuffer(remat_item, original_buffer.size, std::move(rematerialized_uses), /*live_out=*/false, @@ -456,8 +466,7 @@ class MemoryUsageTracker { return false; } const BufferIdList& in_progress_uses = in_progress_item_->buffers_used; - return std::find(in_progress_uses.begin(), in_progress_uses.end(), - buffer_id) != in_progress_uses.end(); + return absl::c_linear_search(in_progress_uses, buffer_id); } // Returns whether the given instruction is live at the current program @@ -535,8 +544,7 @@ MemoryUsageTracker::MemoryUsageTracker( bool unused; for (Item* user_item : GetUsers(instruction_list_, logical_buffer, points_to_analysis, &unused)) { - if (std::find(buffer->users.begin(), buffer->users.end(), - user_item) == buffer->users.end()) { + if (!absl::c_linear_search(buffer->users, user_item)) { buffer->users.push_back(user_item); buffer->unfinished_user_count++; user_item->buffers_used.push_back(buffer->id); @@ -677,8 +685,8 @@ Status MemoryUsageTracker::AddRematerializedInstruction(Item* original_item, << ", remat_instruction = " << remat_item->instruction->name(); TF_RET_CHECK(in_progress_item_ != nullptr); - TF_RET_CHECK(original_item->placed); - TF_RET_CHECK(!remat_item->placed); + TF_RET_CHECK(original_item->placed) << original_item->instruction->name(); + TF_RET_CHECK(!remat_item->placed) << remat_item->instruction->name(); // Construct the list of buffers used and defined by the rematerialization. remat_item->buffers_used = original_item->buffers_used; @@ -707,7 +715,7 @@ Status MemoryUsageTracker::AddRematerializedInstruction(Item* original_item, ItemList unplaced_users; for (Item* user : old_buffer.users) { if (user->placed) { - CHECK(IsFinished(user)); + CHECK(IsFinished(user)) << user->instruction->name(); placed_users.push_back(user); } else { unplaced_users.push_back(user); @@ -784,8 +792,7 @@ bool MemoryUsageTracker::Check() const { for (const Buffer& buffer : buffers_) { if (buffer.defining_instruction->instruction == instruction) { - CHECK(std::find(defined_buffers.begin(), defined_buffers.end(), - buffer.id) != defined_buffers.end()) + CHECK(absl::c_linear_search(defined_buffers, buffer.id)) << "Instruction " << instruction->name() << " defined buffers is missing: " << buffer.ToString(); } @@ -808,8 +815,7 @@ bool MemoryUsageTracker::Check() const { int64 unfinished_uses = 0; for (Item* user : buffer.users) { const BufferIdList& used_buffers = user->buffers_used; - CHECK(std::find(used_buffers.begin(), used_buffers.end(), buffer.id) != - used_buffers.end()) + CHECK(absl::c_linear_search(used_buffers, buffer.id)) << "Instruction " << user->instruction->name() << " used buffers is missing " << buffer.ToString(); if (!IsFinished(user)) { @@ -836,10 +842,10 @@ int64 RematerializationCost(const HloInstruction* instruction, // If none of the users of 'instruction' have been placed in the sequence (as // tracked by memory_tracker), then rematerialization of 'instruction' is a // zero-cost move of 'instruction' in the sequence. - if (!std::any_of(instruction->users().begin(), instruction->users().end(), - [&memory_tracker](const HloInstruction* inst) { - return memory_tracker.IsPlaced(inst); - })) { + if (!absl::c_any_of(instruction->users(), + [&memory_tracker](const HloInstruction* inst) { + return memory_tracker.IsPlaced(inst); + })) { return 0; } @@ -1094,7 +1100,7 @@ StatusOr HloRematerialization::RematerializeComputation( Item* successor_item = instruction_list.GetItem(successor); // Assert to make sure we never remat an operation with control // successor already placed. - CHECK(!successor_item->placed); + CHECK(!successor_item->placed) << successor_item->instruction->name(); place_before.push_back(successor_item); } instruction_list.InsertBeforeInstructions(remat_item, place_before); @@ -1164,7 +1170,7 @@ StatusOr HloRematerialization::RematerializeComputation( // Verify some invariants on the memory tracker. CHECK_EQ(memory_tracker.memory_usage(), 0); for (auto* instruction : computation->instructions()) { - CHECK(memory_tracker.IsPlaced(instruction)); + CHECK(memory_tracker.IsPlaced(instruction)) << instruction->name(); } VLOG(1) << "In computation " << computation->name() << " rematerialized " diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc b/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc index 22c3c40a93a1ddcd36659483fcc79fede32dd2c3..102a360ad8116d8781baf9cb7627a920f4a687c4 100644 --- a/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc +++ b/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc @@ -499,6 +499,52 @@ TEST_F(HloRematerializationTest, InstructionRematerializedMultipleTimes) { EXPECT_THAT(add_4->operand(0), op::Broadcast(param)); } +TEST_F(HloRematerializationTest, CopyNotRematerialized) { + // Test that copies are not rematerialized. + auto module = CreateNewVerifiedModule(); + + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, vec1024_shape_, "param")); + + auto copy = builder.AddInstruction( + HloInstruction::CreateUnary(vec1024_shape_, HloOpcode::kCopy, param)); + + auto negate_a_1 = builder.AddInstruction( + HloInstruction::CreateUnary(vec1024_shape_, HloOpcode::kNegate, copy)); + + auto negate_a_2 = builder.AddInstruction(HloInstruction::CreateUnary( + vec1024_shape_, HloOpcode::kNegate, negate_a_1)); + + auto negate_b_1 = builder.AddInstruction( + HloInstruction::CreateUnary(vec1024_shape_, HloOpcode::kNegate, copy)); + + auto negate_b_2 = builder.AddInstruction(HloInstruction::CreateUnary( + vec1024_shape_, HloOpcode::kNegate, negate_b_1)); + + builder.AddInstruction(HloInstruction::CreateTuple({negate_a_2, negate_b_2})); + + HloComputation* entry_computation = + module->AddEntryComputation(builder.Build()); + + TF_ASSERT_OK_AND_ASSIGN(bool changed, + RunHloRematerialization( + /*memory_limit_bytes=*/1 * 1024, module.get())); + + auto count_copies = [](const HloComputation* computation) { + int64 copy_count = 0; + for (auto* instruction : computation->instructions()) { + if (instruction->opcode() == HloOpcode::kCopy) { + copy_count++; + } + } + return copy_count; + }; + EXPECT_TRUE(changed); + + EXPECT_EQ(count_copies(entry_computation), 1); +} + class IndirectUseTest : public HloRematerializationTest, public ::testing::WithParamInterface {}; @@ -588,8 +634,8 @@ TEST_P(IndirectUseTest, IndirectUseNotRematerialized) { } } -INSTANTIATE_TEST_CASE_P(IndirectUseTestInstantiation, IndirectUseTest, - ::testing::Values(true, false)); +INSTANTIATE_TEST_SUITE_P(IndirectUseTestInstantiation, IndirectUseTest, + ::testing::Values(true, false)); } // namespace diff --git a/tensorflow/compiler/xla/service/hlo_runner.cc b/tensorflow/compiler/xla/service/hlo_runner.cc index 5a9b820a9d7f58695383b21c9e2126cf98970c83..84399f17e5e0b0b1f29cded17b605571bcfa8843 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.cc +++ b/tensorflow/compiler/xla/service/hlo_runner.cc @@ -168,6 +168,35 @@ StatusOr HloRunner::Execute(std::unique_ptr module, /*profile=*/profile); } +StatusOr HloRunner::Execute( + std::unique_ptr executable, + const absl::Span arguments, + ExecutionProfile* profile) { + TF_ASSIGN_OR_RETURN(std::vector argument_buffers, + TransferLiteralsToDevice(arguments)); + TF_ASSIGN_OR_RETURN(ScopedShapedBuffer result, + ExecuteWithDeviceBuffers( + /*module=*/std::move(executable), + /*arguments=*/argument_buffers, + /*profile=*/profile)); + return TransferLiteralFromDevice(result); +} + +StatusOr HloRunner::Execute(std::unique_ptr executable, + const absl::Span arguments, + ExecutionProfile* profile) { + // Construct a vector of plain pointers for the arguments. + std::vector argument_pointers; + argument_pointers.reserve(arguments.size()); + for (const auto& argument : arguments) { + argument_pointers.push_back(&argument); + } + return Execute( + /*module=*/std::move(executable), + /*arguments=*/argument_pointers, + /*profile=*/profile); +} + StatusOr HloRunner::ExecuteWithDeviceBuffers( std::unique_ptr module, const absl::Span arguments, bool run_hlo_passes, @@ -383,9 +412,7 @@ ServiceExecutableRunOptions HloRunner::GetServiceRunOptionsForDevice( if (device_assignment != nullptr) { run_options.set_device_assignment(device_assignment); } - return ServiceExecutableRunOptions( - run_options, backend().StreamBorrower(), - /*xla_intra_op_thread_pool=*/backend().eigen_intra_op_thread_pool()); + return ServiceExecutableRunOptions(run_options, backend().StreamBorrower()); } Backend& HloRunner::backend() { diff --git a/tensorflow/compiler/xla/service/hlo_runner.h b/tensorflow/compiler/xla/service/hlo_runner.h index bb792cf8c9825ff67ca33bbcf2c3c32b1a0ecb85..a6e6015d6a5e2ad6e85cf2411f1a740c0987d8b4 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.h +++ b/tensorflow/compiler/xla/service/hlo_runner.h @@ -124,6 +124,14 @@ class HloRunner { bool run_hlo_passes = true, ExecutionProfile* profile = nullptr); + StatusOr Execute(std::unique_ptr executable, + const absl::Span arguments, + ExecutionProfile* profile = nullptr); + + StatusOr Execute(std::unique_ptr executable, + const absl::Span arguments, + ExecutionProfile* profile = nullptr); + // As Execute(), but accepts and returns device buffers instead of host // buffers. StatusOr ExecuteWithDeviceBuffers( diff --git a/tensorflow/compiler/xla/service/hlo_schedule.cc b/tensorflow/compiler/xla/service/hlo_schedule.cc index 8f6eb974c5179b420c8f961393ca923e0a3b3530..e75373501cffac6a736be89e9f6139b6ff2cdbc1 100644 --- a/tensorflow/compiler/xla/service/hlo_schedule.cc +++ b/tensorflow/compiler/xla/service/hlo_schedule.cc @@ -140,7 +140,7 @@ Status HloSchedule::UpdateComputationSchedule( std::queue worklist; for (HloInstruction* instruction : computation->instructions()) { - if (ids_in_schedule.count(instruction->unique_id()) == 0) { + if (!ids_in_schedule.contains(instruction->unique_id())) { // This is a newly added instruction which is not in the schedule. if (instruction->operands().empty()) { worklist.push(instruction); @@ -204,7 +204,7 @@ Status HloSchedule::Update() { std::vector nonfusion_computations = module_->MakeNonfusionComputations(); for (const HloComputation* computation : nonfusion_computations) { - TF_RET_CHECK(sequences_.count(computation->unique_id()) == 1) + TF_RET_CHECK(sequences_.contains(computation->unique_id())) << "Computation " << computation->name() << " not in HloSchedule."; } if (sequences_.size() > nonfusion_computations.size()) { @@ -215,7 +215,7 @@ Status HloSchedule::Update() { nonfusion_computations_ids.insert(computation->unique_id()); } for (auto it = sequences_.begin(); it != sequences_.end();) { - if (nonfusion_computations_ids.count(it->first) == 0) { + if (!nonfusion_computations_ids.contains(it->first)) { sequences_.erase(it++); } else { ++it; @@ -244,7 +244,7 @@ Status HloSchedule::Verify() const { << "Schedule has " << sequences_.size() << " sequences, but module has " << nonfusion_computations.size() << " non-fusion computations"; for (const HloComputation* computation : nonfusion_computations) { - TF_RET_CHECK(sequences_.count(computation->unique_id()) == 1) + TF_RET_CHECK(sequences_.contains(computation->unique_id())) << "Computation " << computation->name() << " missing from HLO schedule."; } @@ -268,7 +268,7 @@ Status HloSchedule::Verify() const { << instruction_position.size() << " instructions, expected " << computation->instruction_count(); for (const HloInstruction* instruction : computation->instructions()) { - TF_RET_CHECK(instruction_position.count(instruction) == 1) + TF_RET_CHECK(instruction_position.contains(instruction)) << "Instruction " << instruction->name() << " is not in schedule"; } diff --git a/tensorflow/compiler/xla/service/hlo_schedule.h b/tensorflow/compiler/xla/service/hlo_schedule.h index 486ddbf499de80c634bc497158cd79ca066cc866..a5f54ae2c33259d080631061dff9ae40b41495dc 100644 --- a/tensorflow/compiler/xla/service/hlo_schedule.h +++ b/tensorflow/compiler/xla/service/hlo_schedule.h @@ -110,7 +110,7 @@ class HloSchedule { // Returns true if the schedule has a sequence for the given computation. bool is_computation_scheduled(const HloComputation* computation) const { - return sequences_.count(computation->unique_id()) == 1; + return sequences_.contains(computation->unique_id()); } // Updates the schedule such that it is (again) a valid schedule for the diff --git a/tensorflow/compiler/xla/service/hlo_sharding.cc b/tensorflow/compiler/xla/service/hlo_sharding.cc index bdd6ec1169a7d89cb4978a624d5606a5fb89df0f..37cc146bd7a6f2aef9373bd4afd8572ffac6473c 100644 --- a/tensorflow/compiler/xla/service/hlo_sharding.cc +++ b/tensorflow/compiler/xla/service/hlo_sharding.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_sharding.h" +#include "absl/container/flat_hash_set.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "tensorflow/compiler/xla/overflow_util.h" @@ -106,13 +107,12 @@ string HloSharding::ToString() const { bool HloSharding::UsesDevice(int64 device) const { if (IsTuple()) { - return std::any_of( - tuple_elements_.begin(), tuple_elements_.end(), - [&](const HloSharding& s) { return s.UsesDevice(device); }); + return absl::c_any_of(tuple_elements_, [&](const HloSharding& s) { + return s.UsesDevice(device); + }); } const auto& devices = tile_assignment_; - return replicated_ || - std::find(devices.begin(), devices.end(), device) != devices.end(); + return replicated_ || absl::c_linear_search(devices, device); } std::map HloSharding::UsedDevices(int64* count) const { @@ -316,7 +316,7 @@ Status HloSharding::ValidateNonTuple(const Shape& shape, // All tile assignments must be less than the number of available cores and // unique. Status status = Status::OK(); - std::set seen_cores; + absl::flat_hash_set seen_cores; tile_assignment_.Each( [&](absl::Span indices, int32 core) { // Don't overwrite a bad status, so we report the first error. @@ -324,7 +324,7 @@ Status HloSharding::ValidateNonTuple(const Shape& shape, if (core >= num_devices) { status = tensorflow::errors::InvalidArgument(StrCat( "core ", core, " > ", num_devices, " in tile assignment")); - } else if (seen_cores.count(core) != 0) { + } else if (seen_cores.contains(core)) { status = tensorflow::errors::InvalidArgument( StrCat("core ", core, " is not unique in tile assignment")); } diff --git a/tensorflow/compiler/xla/service/hlo_sharding.h b/tensorflow/compiler/xla/service/hlo_sharding.h index 9775505f8608ced3e33abe376f4922cc6a972726..5789ae09988d2a85247c5b8c037a172b3699f3b7 100644 --- a/tensorflow/compiler/xla/service/hlo_sharding.h +++ b/tensorflow/compiler/xla/service/hlo_sharding.h @@ -101,8 +101,8 @@ class HloSharding { if (!IsTuple()) { return replicated_; } - return std::all_of(tuple_elements_.begin(), tuple_elements_.end(), - [](const HloSharding& s) { return s.IsReplicated(); }); + return absl::c_all_of( + tuple_elements_, [](const HloSharding& s) { return s.IsReplicated(); }); } // Returns true if the tile size is the same as the input size. @@ -110,8 +110,9 @@ class HloSharding { if (!IsTuple()) { return maximal_; } - return std::all_of(tuple_elements_.begin(), tuple_elements_.end(), - [](const HloSharding& s) { return s.IsTileMaximal(); }); + return absl::c_all_of(tuple_elements_, [](const HloSharding& s) { + return s.IsTileMaximal(); + }); } // Returns true if the sharding defines an operation on the given device. diff --git a/tensorflow/compiler/xla/service/hlo_sharding_metadata.cc b/tensorflow/compiler/xla/service/hlo_sharding_metadata.cc index b414d2a66328284a4d8be8d35206bb837a2b3a58..094d98bc6e54028557f6d38cd165bf34e1fb8c46 100644 --- a/tensorflow/compiler/xla/service/hlo_sharding_metadata.cc +++ b/tensorflow/compiler/xla/service/hlo_sharding_metadata.cc @@ -99,7 +99,7 @@ std::vector LocatePassThroughDomainLinks( << "Instruction is not a kDomain: " << instruction->ToString(); for (HloInstruction* user : instruction->users()) { if (user->opcode() == HloOpcode::kDomain && - domain.exit_domains.count(user) != 0) { + domain.exit_domains.contains(user)) { pass_through.emplace_back(user, instruction); VLOG(2) << "Found passthrough domain link:"; VLOG(2) << " " << user->ToString(); @@ -253,7 +253,7 @@ StatusOr ApplyShardingFromUsers(HloInstruction* instruction, instruction->shape(), HloSharding::AssignDevice(kUnassignedDevice)); for (HloInstruction* user : instruction->users()) { if (user->opcode() == HloOpcode::kDomain && - domain.exit_domains.count(user) > 0) { + domain.exit_domains.contains(user)) { // If a user is a domain and it is registered in the domain exits, then // the instruction sharding is taken directly from the domain, and no // further users need to be visited. diff --git a/tensorflow/compiler/xla/service/hlo_tfgraph_builder.cc b/tensorflow/compiler/xla/service/hlo_tfgraph_builder.cc index 4ea39a7628317867eb054bb14de48f721354cd9d..6925dc37dbe9dc90e79d315cf41a3416e2084c81 100644 --- a/tensorflow/compiler/xla/service/hlo_tfgraph_builder.cc +++ b/tensorflow/compiler/xla/service/hlo_tfgraph_builder.cc @@ -61,8 +61,7 @@ void CleanNodeName(string* name) { name->erase(std::remove(name->begin(), name->end(), '%'), name->end()); const string chars_to_replace = "<>[]"; auto pred = [&](char c) { - return std::find(chars_to_replace.begin(), chars_to_replace.end(), c) != - chars_to_replace.end(); + return absl::c_linear_search(chars_to_replace, c); }; std::replace_if(name->begin(), name->end(), pred, '_'); } @@ -163,11 +162,11 @@ void HloTfGraphBuilder::SetNodeAttrs(const HloInstruction* instruction, // For tuples, emit the full shape because the layout of a tuple is not // represented in a single Layout field. layout_string = ShapeUtil::HumanStringWithLayout(instruction->shape()); - } else { - layout_string = StrCat( - "{", - absl::StrJoin(LayoutUtil::MinorToMajor(instruction->shape()), ","), - "}"); + } else if (instruction->shape().has_layout()) { + // For non-tuples, only emit the layout when the shape has a Layout. + // This extra check is required because LayoutUtil::HasLayout ignores + // token, opaque types etc. + layout_string = instruction->shape().layout().ToString(); } attrs["layout"].set_s(layout_string); } diff --git a/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc b/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc index 1e2b31a1f2bb4865faafc3d14e2b194e3aa171a1..498abcfe04d963575fb9200443efb7d911a6293e 100644 --- a/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc +++ b/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/tensor_shape.pb.h" +#include "tensorflow/core/lib/core/status_test_util.h" namespace xla { namespace hlo_graph_dumper { @@ -178,6 +179,23 @@ TEST_F(HloTfGraphBuilderTest, EmbeddedComputationsDiamond) { EXPECT_GT(generator_.GetGraphDef().node_size(), 0); } +TEST_F(HloTfGraphBuilderTest, TokenHasNoLayout) { + auto builder = HloComputation::Builder("Token"); + auto token = builder.AddInstruction(HloInstruction::CreateToken()); + OpMetadata metadata; + metadata.set_op_name("x"); + metadata.set_op_type("y"); + token->set_metadata(metadata); + TF_ASSERT_OK(generator_.AddComputation(*builder.Build())); + GraphDef graph_def = generator_.GetGraphDef(); + ASSERT_EQ(graph_def.node_size(), 1); + const auto &node = graph_def.node(0); + ASSERT_EQ(GetNodeAttr(node, "type").s(), "TOKEN"); + ASSERT_EQ(GetNodeAttr(node, "layout").s(), ""); + ASSERT_EQ(GetNodeAttr(node, "tf_op_name").s(), "x"); + ASSERT_EQ(GetNodeAttr(node, "tf_op_type").s(), "y"); +} + } // namespace } // namespace hlo_graph_dumper } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_value.cc b/tensorflow/compiler/xla/service/hlo_value.cc index d409df06be44a79d1bf7c90ce3ff34fca975fda4..218b33b2ac2b86edc30b2f014ba206c71da37682 100644 --- a/tensorflow/compiler/xla/service/hlo_value.cc +++ b/tensorflow/compiler/xla/service/hlo_value.cc @@ -209,7 +209,7 @@ std::ostream& operator<<(std::ostream& out, const HloValue& value) { } void HloValueSet::SortAndUniquifyValues() { - std::sort(values_.begin(), values_.end(), HloValue::IdLessThan); + absl::c_sort(values_, HloValue::IdLessThan); values_.erase(std::unique(values_.begin(), values_.end(), HloValue::IdEqual), values_.end()); } diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc index 14e29533c2687f8fd6d258d4bcce2d0c0acd514d..39153cc06c40765909ab4d4cc29ee123f144bb55 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.cc +++ b/tensorflow/compiler/xla/service/hlo_verifier.cc @@ -50,6 +50,7 @@ bool IsCallerInstruction(HloInstruction* hlo) { case HloOpcode::kReduceWindow: case HloOpcode::kScatter: case HloOpcode::kSelectAndScatter: + case HloOpcode::kSort: case HloOpcode::kFusion: return true; default: @@ -167,6 +168,15 @@ Status ShapeVerifier::HandleFft(HloInstruction* fft) { return CheckShape(fft, expected); } +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(), + hlo->triangular_solve_options())); + return CheckShape(hlo, expected); +} + Status ShapeVerifier::HandleAllReduce(HloInstruction* crs) { std::vector operand_shapes; for (const HloInstruction* operand : crs->operands()) { @@ -184,6 +194,10 @@ Status ShapeVerifier::HandleAllToAll(HloInstruction* hlo) { ShapeInference::InferAllToAllTupleShape(operand_shapes)); } +Status ShapeVerifier::HandleReplicaId(HloInstruction* hlo) { + return CheckShape(hlo, ShapeUtil::MakeShape(U32, {})); +} + Status ShapeVerifier::HandleCollectivePermute(HloInstruction* hlo) { TF_RETURN_IF_ERROR(CheckOperandCount(hlo, 1)); return CheckShape(hlo, ShapeInference::InferCollectivePermuteShape( @@ -323,13 +337,48 @@ Status ShapeVerifier::HandleSort(HloInstruction* sort) { return InternalError("Expected at least 1 operand for %s instruction: %s", HloOpcodeString(sort->opcode()), sort->ToString()); } + HloComputation* compare = sort->to_apply(); + + // Check that the 'compare' computation returns a PRED. + Shape compare_shape = compare->root_instruction()->shape(); + if (!ShapesSame(compare_shape, ShapeUtil::MakeShape(PRED, {}))) { + return InternalError( + "The Sort compare computation shape does not lead to a scalar " + "predicate shape: %s", + StringifyShape(compare_shape)); + } + + // Check that the number of parameters of the 'compare' computation is + // correct. + TF_RETURN_IF_ERROR( + CheckParameterCount(sort, compare, sort->operand_count() * 2)); + + // Verify that the operands of the compare computation have the correct scalar + // shapes. + for (int64 parameter_idx = 0; parameter_idx < compare->num_parameters(); + ++parameter_idx) { + int64 operand_idx = parameter_idx / 2; + Shape expected_scalar_shape = ShapeUtil::MakeShape( + sort->operand(operand_idx)->shape().element_type(), {}); + Shape actual_parameter_shape = + compare->parameter_instruction(parameter_idx)->shape(); + if (!ShapeUtil::CompatibleIgnoringFpPrecision(expected_scalar_shape, + actual_parameter_shape)) { + return InternalError( + "Expected the %lld-th parameter of the compare computation of sort " + "to have shape %s, but got %s", + parameter_idx, StringifyShape(expected_scalar_shape), + StringifyShape(actual_parameter_shape)); + } + } + + // Verify that all operand shapes have the same dimensions. for (int64 operand = 1; operand < sort->operand_count(); ++operand) { if (!ShapeUtil::SameDimensions(sort->operand(0)->shape(), sort->operand(operand)->shape())) { return InternalError( - "Expected sort to have to have the same dimensions for the keys " - "and the values. Keys shape is: %s\n, Values shape (operand index " - "%lld) is: %s", + "Expected sort to have to have the same dimensions for all operands. " + "First operand shape is: %s\n, shape (operand index %lld) is: %s", StringifyShape(sort->operand(0)->shape()), operand, StringifyShape(sort->operand(operand)->shape())); } @@ -349,6 +398,9 @@ 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."); + } const int64 rank = iota->shape().rank(); if (rank == 0) { return InternalError("Iota does not support scalars."); @@ -369,6 +421,24 @@ Status ShapeVerifier::HandleGetTupleElement(HloInstruction* get_tuple_element) { get_tuple_element->tuple_index())); } +namespace { +Status SameElementTypesForOperandsAndToApplyParameters( + const HloInstruction& instruction, int64 num_operands_to_check) { + const ProgramShape& to_apply = instruction.to_apply()->ComputeProgramShape(); + for (int i = 0; i < num_operands_to_check; ++i) { + const Shape& parameter_shape = to_apply.parameters(i); + const Shape& operand_shape = instruction.operands()[i]->shape(); + if (!ShapeUtil::SameElementType(parameter_shape, operand_shape)) { + return InvalidArgument( + "Shape mismatch between to_apply computation" + " parameter and operand %d in %s.", + i, instruction.ToString().c_str()); + } + } + return Status::OK(); +} +} // namespace + Status ShapeVerifier::HandleReduce(HloInstruction* reduce) { if (reduce->operand_count() % 2 != 0) { return InternalError( @@ -380,13 +450,27 @@ Status ShapeVerifier::HandleReduce(HloInstruction* reduce) { for (const HloInstruction* operand : reduce->operands()) { operand_shapes.push_back(&operand->shape()); } - return CheckShape(reduce, ShapeInference::InferReduceShape( - operand_shapes, reduce->dimensions(), - reduce->to_apply()->ComputeProgramShape())); + TF_RETURN_IF_ERROR( + CheckShape(reduce, ShapeInference::InferReduceShape( + operand_shapes, reduce->dimensions(), + reduce->to_apply()->ComputeProgramShape()))); + + return allow_mixed_precision_ + ? Status::OK() + : SameElementTypesForOperandsAndToApplyParameters( + *reduce, reduce->operands().size() - 1); } 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()) { + return InternalError( + "Bitcast can not change the element type from %s to %s", + PrimitiveType_Name(bitcast->operand(0)->shape().element_type()), + PrimitiveType_Name(bitcast->shape().element_type())); + } return Status::OK(); } @@ -496,38 +580,23 @@ Status ShapeVerifier::HandleSlice(HloInstruction* slice) { } Status ShapeVerifier::HandleDynamicSlice(HloInstruction* dynamic_slice) { - const DebugOptions& debug_options = - dynamic_slice->GetModule()->config().debug_options(); - const bool allow_scalar_indices = - debug_options.xla_allow_scalar_index_dynamic_ops(); - if (!allow_scalar_indices) { - TF_RETURN_IF_ERROR(CheckOperandCount(dynamic_slice, 2)); - } return CheckShape( dynamic_slice, ShapeInference::InferDynamicSliceShape( dynamic_slice->operand(0)->shape(), Cast(dynamic_slice)->index_shapes(), - dynamic_slice->dynamic_slice_sizes(), allow_scalar_indices)); + dynamic_slice->dynamic_slice_sizes())); } Status ShapeVerifier::HandleDynamicUpdateSlice( HloInstruction* dynamic_update_slice) { - const DebugOptions& debug_options = - dynamic_update_slice->GetModule()->config().debug_options(); - const bool allow_scalar_indices = - debug_options.xla_allow_scalar_index_dynamic_ops(); - if (!allow_scalar_indices) { - TF_RETURN_IF_ERROR(CheckOperandCount(dynamic_update_slice, 3)); - } return CheckShape( dynamic_update_slice, ShapeInference::InferDynamicUpdateSliceShape( dynamic_update_slice->operand(0)->shape(), dynamic_update_slice->operand(1)->shape(), Cast(dynamic_update_slice) - ->index_shapes(), - allow_scalar_indices)); + ->index_shapes())); } Status ShapeVerifier::HandleTuple(HloInstruction* tuple) { @@ -545,19 +614,31 @@ Status ShapeVerifier::HandleMap(HloInstruction* map) { // arbitrary map dimensions. std::vector map_dims(max_operand_rank); std::iota(map_dims.begin(), map_dims.end(), 0); - return CheckShape(map, ShapeInference::InferMapShape( - operand_shapes, - map->to_apply()->ComputeProgramShape(), map_dims)); + + TF_RETURN_IF_ERROR(CheckShape( + map, + ShapeInference::InferMapShape( + operand_shapes, map->to_apply()->ComputeProgramShape(), map_dims))); + + return allow_mixed_precision_ + ? Status::OK() + : SameElementTypesForOperandsAndToApplyParameters( + *map, map->operands().size()); } Status ShapeVerifier::HandleReduceWindow(HloInstruction* reduce_window) { TF_RETURN_IF_ERROR(CheckOperandCount(reduce_window, 2)); - return CheckShape( + TF_RETURN_IF_ERROR(CheckShape( reduce_window, ShapeInference::InferReduceWindowShape( reduce_window->operand(0)->shape(), reduce_window->operand(1)->shape(), reduce_window->window(), - reduce_window->to_apply()->ComputeProgramShape())); + reduce_window->to_apply()->ComputeProgramShape()))); + + return allow_mixed_precision_ + ? Status::OK() + : SameElementTypesForOperandsAndToApplyParameters(*reduce_window, + 1); } Status ShapeVerifier::HandleSelectAndScatter(HloInstruction* instruction) { @@ -709,7 +790,6 @@ Status CheckMixedPrecisionOperands(const HloInstruction* instruction) { case HloOpcode::kRecv: case HloOpcode::kRecvDone: case HloOpcode::kReducePrecision: - case HloOpcode::kSelect: case HloOpcode::kTupleSelect: case HloOpcode::kSend: case HloOpcode::kSendDone: diff --git a/tensorflow/compiler/xla/service/hlo_verifier.h b/tensorflow/compiler/xla/service/hlo_verifier.h index a1a6aba9728c137d17487b5914f67cb3966fc12b..a9b5e9a3e6eec19e125188a192694fcaadfe2322 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.h +++ b/tensorflow/compiler/xla/service/hlo_verifier.h @@ -52,9 +52,11 @@ class ShapeVerifier : public DfsHloVisitor { Status HandleDot(HloInstruction* dot) override; Status HandleConvolution(HloInstruction* convolution) override; Status HandleFft(HloInstruction* fft) override; + Status HandleTriangularSolve(HloInstruction* hlo) override; Status HandleAllReduce(HloInstruction* crs) override; Status HandleAllToAll(HloInstruction* hlo) override; Status HandleCollectivePermute(HloInstruction* hlo) override; + Status HandleReplicaId(HloInstruction* hlo) override; Status HandleReducePrecision(HloInstruction* reduce_precision) override; Status HandleInfeed(HloInstruction*) override; Status HandleOutfeed(HloInstruction*) override; @@ -168,8 +170,13 @@ class ShapeVerifier : public DfsHloVisitor { // An interface used to encapsulate target-specific verification quirks. class TargetVerifierMetadata { public: + TargetVerifierMetadata(std::function shape_size_function) + : shape_size_function_(shape_size_function) {} + // Returns a target-specific shape size. - virtual int64 ShapeSize(const Shape& shape) const = 0; + int64 ShapeSize(const Shape& shape) const { + return shape_size_function_(shape); + } virtual std::unique_ptr GetVerifier() const = 0; @@ -178,20 +185,23 @@ class TargetVerifierMetadata { TargetVerifierMetadata(const TargetVerifierMetadata&) = delete; TargetVerifierMetadata& operator=(const TargetVerifierMetadata&) = delete; + + private: + // Returns a target-specific shape size. + std::function shape_size_function_; }; // The default implementation of TargetVerifierMetadata, used unless the target // needs to override it. class DefaultVerifierMetadata : public TargetVerifierMetadata { public: - DefaultVerifierMetadata(bool layout_sensitive, bool allow_mixed_precision) - : layout_sensitive_(layout_sensitive), + DefaultVerifierMetadata( + bool layout_sensitive, bool allow_mixed_precision, + std::function shape_size_function) + : TargetVerifierMetadata(shape_size_function), + layout_sensitive_(layout_sensitive), allow_mixed_precision_(allow_mixed_precision) {} - int64 ShapeSize(const Shape& shape) const override { - return ShapeUtil::ByteSizeOf(shape); - } - // Creates a ShapeVerifier that checks that shapes match inferred // expectations. This creates a new verifier every time because ShapeVerifier, // being a DfsHloVisitor, is stateful. We want a clean object for each run of @@ -210,11 +220,14 @@ class DefaultVerifierMetadata : public TargetVerifierMetadata { // the module. class HloVerifier : public HloModulePass { public: - explicit HloVerifier(bool layout_sensitive, bool allow_mixed_precision, - std::function - instruction_can_change_layout_func = {}) + explicit HloVerifier( + bool layout_sensitive, bool allow_mixed_precision, + std::function + instruction_can_change_layout_func = {}, + std::function shape_size_func = + [](const Shape& shape) { return ShapeUtil::ByteSizeOf(shape); }) : target_metadata_(absl::make_unique( - layout_sensitive, allow_mixed_precision)), + layout_sensitive, allow_mixed_precision, shape_size_func)), instruction_can_change_layout_func_( std::move(instruction_can_change_layout_func)) { CHECK(instruction_can_change_layout_func_ == nullptr || layout_sensitive); diff --git a/tensorflow/compiler/xla/service/hlo_verifier_test.cc b/tensorflow/compiler/xla/service/hlo_verifier_test.cc index 91f247a9bb3af79e27bc028cfc4b4963fb2b4b26..523890b3c7268c06cdb6aaa67749f26a1cb62855 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier_test.cc +++ b/tensorflow/compiler/xla/service/hlo_verifier_test.cc @@ -450,8 +450,9 @@ TEST_F(HloVerifierTestLayoutSensitive, SliceWithLayoutChangeNotAllowed) { HloModule SliceWithLayoutChange ENTRY SliceWithLayoutChange { par0 = f32[4,5]{0,1} parameter(0) - par1 = s32[2] parameter(1) - ROOT dslice0 = f32[3,4]{1,0} dynamic-slice(par0, par1), + par1 = s32[] parameter(1) + par2 = s32[] parameter(2) + ROOT dslice0 = f32[3,4]{1,0} dynamic-slice(par0, par1, par2), dynamic_slice_sizes={3,4} } )"; @@ -480,5 +481,138 @@ TEST_F(HloVerifierTestLayoutSensitive, ConcatWithLayoutChangeNotAllowed) { EXPECT_THAT(status.error_message(), HasSubstr("Instruction shouldn't change layouts")); } + +TEST_F(HloVerifierTest, BitcastCanNotChangeElementType) { + const char* const hlo_string = R"( + HloModule Module + + ENTRY BitcastCanNotChangeElementType { + constant.0 = f32[2] constant({0.0, 0.0}) + ROOT bitcast = s32[2] bitcast(constant.0) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto module, ParseHloString(hlo_string)); + + auto status = verifier().Run(module.get()).status(); + ASSERT_FALSE(status.ok()); + EXPECT_THAT(status.error_message(), + HasSubstr("Bitcast can not change the element type")); +} + +TEST_F(HloVerifierTest, SelectMixedPrecisionNotAllowed) { + const char* const hlo_string = R"( + HloModule Module + + ENTRY SelectMixedPrecisionNotAllowed { + p0 = pred[] parameter(0) + p1 = f32[32] parameter(1) + p2 = bf16[32] parameter(2) + ROOT select = f32[32] select(p0, p1, p2) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto module, ParseHloString(hlo_string)); + + auto status = verifier().Run(module.get()).status(); + ASSERT_FALSE(status.ok()); + EXPECT_THAT(status.error_message(), + HasSubstr("Seen floating point types of different precisions")); +} + +TEST_F(HloVerifierTestAllowMixedPrecision, SelectMixedPrecisionAllowed) { + const char* const hlo_string = R"( + HloModule Module + + ENTRY SelectMixedPrecisionAllowed { + p0 = pred[] parameter(0) + p1 = f32[32] parameter(1) + p2 = bf16[32] parameter(2) + ROOT select = f32[32] select(p0, p1, p2) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto module, ParseHloString(hlo_string)); + + auto status = verifier().Run(module.get()).status(); + ASSERT_TRUE(status.ok()); +} + +TEST_F(HloVerifierTest, IotaNonArrayResult) { + const char* const hlo_string = R"( + HloModule IotaTupleResult + + ENTRY kernelEntry { + ROOT iota = () iota(), iota_dimension=24 + } + )"; + + TF_ASSERT_OK_AND_ASSIGN(auto module, ParseHloString(hlo_string)); + + auto status = verifier().Run(module.get()).status(); + ASSERT_FALSE(status.ok()); + EXPECT_THAT(status.error_message(), + HasSubstr("does not support non-array result")); +} + +static const char* const kMapOperandComputationMismatchHlo = R"( + HloModule MapOperandComputationMismatch + + Computation { + param0 = f32[] parameter(0) + constant = f32[] constant(1) + ROOT add = f32[] add(param0, constant) + } + + ENTRY kernelEntry { + param = f64[] parameter(0) + ROOT map = f32[] map(param), dimensions={}, to_apply=Computation +})"; + +TEST_F(HloVerifierTest, MapOperandComputationMismatch) { + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseHloString(kMapOperandComputationMismatchHlo)); + auto status = verifier().Run(module.get()).status(); + ASSERT_FALSE(status.ok()); + EXPECT_THAT( + status.error_message(), + HasSubstr( + "Shape mismatch between to_apply computation parameter and operand")); +} + +TEST_F(HloVerifierTestAllowMixedPrecision, MapOperandComputationMismatch) { + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseHloString(kMapOperandComputationMismatchHlo)); + auto status = verifier().Run(module.get()).status(); + ASSERT_TRUE(status.ok()); +} + +static const char* const kReduceOperandComputationMismatchHlo = R"( + HloModule ReduceOperandComputationMismatch + computation { + x = f32[] parameter(0) + y = f32[] parameter(1) + ROOT add = f32[] add(x, y) + } + + ENTRY kernelEntry { + arg0 = f16[64,64,224,224]{3,2,1,0} parameter(0) + constant = f16[] constant(0) + reduce = f16[64]{0} reduce(arg0, constant), dimensions={0,2,3}, to_apply=computation + })"; + +TEST_F(HloVerifierTest, ReduceOperandComputationMismatch) { + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseHloString(kReduceOperandComputationMismatchHlo)); + auto status = verifier().Run(module.get()).status(); + ASSERT_FALSE(status.ok()); + EXPECT_THAT(status.error_message(), + HasSubstr("Expected instruction to have shape equal to f32[64]")); +} + +TEST_F(HloVerifierTestAllowMixedPrecision, ReduceOperandComputationMismatch) { + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseHloString(kReduceOperandComputationMismatchHlo)); + auto status = verifier().Run(module.get()).status(); + ASSERT_TRUE(status.ok()); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/human_readable_profile_builder.cc b/tensorflow/compiler/xla/service/human_readable_profile_builder.cc index 90904ac00110457bcc3b8974816a7080c4ab89fc..88fc62bd1e2a7830b3f61738a8642308ef4225a7 100644 --- a/tensorflow/compiler/xla/service/human_readable_profile_builder.cc +++ b/tensorflow/compiler/xla/service/human_readable_profile_builder.cc @@ -128,9 +128,9 @@ string HumanReadableProfileBuilder::ToString() const { // Sort ops in decreasing order of cycles, and print them. std::vector sorted_ops(op_infos_); - std::sort( - sorted_ops.begin(), sorted_ops.end(), - [](const OpInfo& a, const OpInfo& b) { return a.cycles > b.cycles; }); + absl::c_sort(sorted_ops, [](const OpInfo& a, const OpInfo& b) { + return a.cycles > b.cycles; + }); for (const auto& op : sorted_ops) { print_op(op); } diff --git a/tensorflow/compiler/xla/service/indexed_array_analysis.cc b/tensorflow/compiler/xla/service/indexed_array_analysis.cc index a41cf714c5ed6adcf7aa1f0e54cf052594834b5d..c5d32a4b9ad8c708ec0870173fa72320238e8464 100644 --- a/tensorflow/compiler/xla/service/indexed_array_analysis.cc +++ b/tensorflow/compiler/xla/service/indexed_array_analysis.cc @@ -27,7 +27,6 @@ limitations under the License. #include "tensorflow/compiler/xla/util.h" namespace xla { -namespace gtl = ::tensorflow::gtl; namespace { using Analysis = IndexedArrayAnalysis; @@ -103,7 +102,7 @@ Status IndexedArrayAnalysis::TraverseAndPopulateCache( do { const HloInstruction* instr = stack.back(); - if (cache_.count(instr)) { + if (cache_.contains(instr)) { stack.pop_back(); continue; } @@ -111,9 +110,9 @@ Status IndexedArrayAnalysis::TraverseAndPopulateCache( switch (FindOrDie(dfs_state_map, instr)) { case kDiscovered: { for (const HloInstruction* operand : instr->operands()) { - if (!cache_.count(operand)) { + if (!cache_.contains(operand)) { stack.push_back(operand); - CHECK(!dfs_state_map.count(operand) || + CHECK(!dfs_state_map.contains(operand) || dfs_state_map[operand] == kDiscovered); dfs_state_map[operand] = kDiscovered; } diff --git a/tensorflow/compiler/xla/service/indexed_array_analysis_test.cc b/tensorflow/compiler/xla/service/indexed_array_analysis_test.cc index 295465c8481bcb7d1385192febe0d09614e393b3..62107b5a88d4e37552fa5a6384700a9291a9c655 100644 --- a/tensorflow/compiler/xla/service/indexed_array_analysis_test.cc +++ b/tensorflow/compiler/xla/service/indexed_array_analysis_test.cc @@ -13,9 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include - #include "tensorflow/compiler/xla/service/indexed_array_analysis.h" +#include "absl/strings/ascii.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/tests/test_utils.h" @@ -43,7 +42,7 @@ class IndexedArrayAnalysisTest : public HloTestBase { string result; for (char c : text) { - if (!isspace(c)) { + if (!absl::ascii_isspace(c)) { result.push_back(c); } else if (!result.empty() && result.back() != ' ') { result.push_back(' '); diff --git a/tensorflow/compiler/xla/service/instruction_fusion.cc b/tensorflow/compiler/xla/service/instruction_fusion.cc index f6f25c44131226d35f8e927d62defee291bf46dd..a5767774f2dcebf8fd309823d48dcc3269b3f594 100644 --- a/tensorflow/compiler/xla/service/instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/instruction_fusion.cc @@ -95,6 +95,7 @@ bool IsAlwaysDuplicable(const HloInstruction& instruction) { case HloOpcode::kPad: case HloOpcode::kReal: case HloOpcode::kReducePrecision: + case HloOpcode::kReplicaId: case HloOpcode::kReshape: case HloOpcode::kReverse: case HloOpcode::kRoundNearestAfz: @@ -157,6 +158,7 @@ bool IsAlwaysDuplicable(const HloInstruction& instruction) { case HloOpcode::kSort: case HloOpcode::kTanh: case HloOpcode::kTrace: + case HloOpcode::kTriangularSolve: case HloOpcode::kWhile: case HloOpcode::kGetDimensionSize: return true; @@ -178,19 +180,18 @@ bool InstructionFusion::EffectivelyAtMostUnary(HloInstruction* hlo) { output_rank = std::max(output_rank, ShapeUtil::TrueRank(subshape)); } }); - return std::count_if(hlo->operands().begin(), hlo->operands().end(), - [output_rank](HloInstruction* operand) { - if (operand->opcode() == HloOpcode::kBroadcast || - operand->opcode() == HloOpcode::kIota) { - return false; - } - if (operand->opcode() == HloOpcode::kConstant && - ShapeUtil::IsEffectiveScalar(operand->shape())) { - return false; - } - return ShapeUtil::TrueRank(operand->shape()) >= - output_rank; - }) <= 1; + return absl::c_count_if( + hlo->operands(), [output_rank](HloInstruction* operand) { + if (operand->opcode() == HloOpcode::kBroadcast || + operand->opcode() == HloOpcode::kIota) { + return false; + } + if (operand->opcode() == HloOpcode::kConstant && + ShapeUtil::IsEffectiveScalar(operand->shape())) { + return false; + } + return ShapeUtil::TrueRank(operand->shape()) >= output_rank; + }) <= 1; } bool InstructionFusion::CanFuseOnAllPaths( @@ -409,9 +410,8 @@ class ReversePostOrderFusionQueue : public FusionQueue { } sorted_operand_numbers.push_back(i); } - std::sort( - sorted_operand_numbers.begin(), sorted_operand_numbers.end(), - [&](int64 i, int64 j) { + absl::c_sort( + sorted_operand_numbers, [&](int64 i, int64 j) { // Instructions with higher priority in the queue come first. return ( FindOrDie(post_order_index_, instruction->mutable_operand(i)) > diff --git a/tensorflow/compiler/xla/service/interpreter/BUILD b/tensorflow/compiler/xla/service/interpreter/BUILD index a981d94a999e3d322986bc2bfd56a5b0b5d175fc..8cd936268994c2a25c2c0debe0a003d1d05cbd0b 100644 --- a/tensorflow/compiler/xla/service/interpreter/BUILD +++ b/tensorflow/compiler/xla/service/interpreter/BUILD @@ -1,12 +1,12 @@ -licenses(["notice"]) # Apache 2.0 - -package(default_visibility = ["//visibility:public"]) - load( "//tensorflow/core:platform/default/build_config_root.bzl", "if_static", ) +licenses(["notice"]) # Apache 2.0 + +package(default_visibility = ["//visibility:public"]) + cc_library( name = "interpreter_transfer_manager", srcs = ["interpreter_transfer_manager.cc"], @@ -34,6 +34,7 @@ cc_library( "//tensorflow/compiler/xla/service:algebraic_simplifier", "//tensorflow/compiler/xla/service:compiler", "//tensorflow/compiler/xla/service:computation_placer", + "//tensorflow/compiler/xla/service:dynamic_index_splitter", "//tensorflow/compiler/xla/service:executable", "//tensorflow/compiler/xla/service:flatten_call_graph", "//tensorflow/compiler/xla/service:hlo", @@ -47,8 +48,11 @@ cc_library( "//tensorflow/compiler/xla/service:hlo_subcomputation_unification", "//tensorflow/compiler/xla/service:layout_assignment", "//tensorflow/compiler/xla/service:map_inliner", + "//tensorflow/compiler/xla/service:reduce_precision_insertion", "//tensorflow/compiler/xla/service:reshape_mover", + "//tensorflow/compiler/xla/service:triangular_solve_expander", "//tensorflow/compiler/xla/service:while_loop_simplifier", + "//tensorflow/compiler/xla/service/cpu:custom_call_target_registry", "//tensorflow/core:lib", "//tensorflow/stream_executor", "@com_google_absl//absl/memory", @@ -115,6 +119,8 @@ cc_library( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", "//tensorflow/core:stream_executor_headers_lib", + "//tensorflow/stream_executor/host:host_stream", + "//tensorflow/stream_executor/host:host_timer", "@com_google_absl//absl/types:span", ], ) diff --git a/tensorflow/compiler/xla/service/interpreter/compiler.cc b/tensorflow/compiler/xla/service/interpreter/compiler.cc index d37ae94bf6c4c697bbf30390c02a5099271e00a4..792773c676984aa280c1b20cb7fd0fc7c9425f6c 100644 --- a/tensorflow/compiler/xla/service/interpreter/compiler.cc +++ b/tensorflow/compiler/xla/service/interpreter/compiler.cc @@ -21,6 +21,8 @@ limitations under the License. #include "absl/memory/memory.h" #include "tensorflow/compiler/xla/service/algebraic_simplifier.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" #include "tensorflow/compiler/xla/service/flatten_call_graph.h" #include "tensorflow/compiler/xla/service/hlo_constant_folding.h" #include "tensorflow/compiler/xla/service/hlo_cse.h" @@ -31,7 +33,9 @@ limitations under the License. #include "tensorflow/compiler/xla/service/interpreter/executable.h" #include "tensorflow/compiler/xla/service/layout_assignment.h" #include "tensorflow/compiler/xla/service/map_inliner.h" +#include "tensorflow/compiler/xla/service/reduce_precision_insertion.h" #include "tensorflow/compiler/xla/service/reshape_mover.h" +#include "tensorflow/compiler/xla/service/triangular_solve_expander.h" #include "tensorflow/compiler/xla/service/while_loop_simplifier.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/core/lib/core/errors.h" @@ -40,12 +44,51 @@ limitations under the License. namespace xla { namespace interpreter { +namespace { + +// Handles custom_call ops during evaluation by routing them through the global +// CPU registry used by other CPU-based backends. +StatusOr HandleEvaluatorCustomCall( + HloInstruction* custom_call, absl::Span operands) { + // Find the target C function in the global registry. + auto* registry = xla::cpu::CustomCallTargetRegistry::Global(); + void* target_fn = registry->Lookup(custom_call->custom_call_target()); + if (!target_fn) { + return NotFound("Custom call target '%s' was not registered", + custom_call->custom_call_target()); + } + + // Populate pointers to operand and output literal data. + std::vector operand_data; + operand_data.reserve(operands.size()); + for (const auto* literal : operands) { + operand_data.push_back(literal->untyped_data()); + } + auto output = Literal::CreateFromShape(custom_call->shape()); + void* output_data = output.untyped_data(); + + // Call the target function matching the C ABI used by the CPU backends. + auto* typed_fn = reinterpret_cast(target_fn); + (*typed_fn)(output_data, operand_data.data()); + + return std::move(output); +} + +} // namespace + Status InterpreterCompiler::RunHloOptimization(HloModule* hlo_module) { HloPassPipeline pipeline("Interpreter"); + pipeline.AddPass(); + pipeline.AddPass(); pipeline.AddPass( hlo_module->mutable_entry_computation_layout(), LayoutAssignment::InstructionCanChangeLayout); + + ReducePrecisionInsertion::AddPasses( + &pipeline, hlo_module->config().debug_options(), + ReducePrecisionInsertion::PassTiming::BEFORE_OPTIMIZATION); + return pipeline.Run(hlo_module).status(); } @@ -75,10 +118,12 @@ StatusOr> InterpreterCompiler::RunBackend( // In this case we are using an HloEvaluator at execution time, so we don't // need to compile anything - // Create executable from only the Hlo module. auto evaluator = absl::make_unique(); evaluator->set_use_fast_path( hlo_module->config().debug_options().xla_hlo_evaluator_use_fast_path()); + evaluator->set_custom_call_handler(HandleEvaluatorCustomCall); + + // Create executable from only the Hlo module. std::unique_ptr executable = absl::make_unique(std::move(hlo_module), std::move(evaluator)); diff --git a/tensorflow/compiler/xla/service/layout_assignment.cc b/tensorflow/compiler/xla/service/layout_assignment.cc index 1c4ac178f0b9a2789d4e6d4554e3a56e0e7cf77f..2dd9e055503e01f5c1ace5e6cd3dc64012828b71 100644 --- a/tensorflow/compiler/xla/service/layout_assignment.cc +++ b/tensorflow/compiler/xla/service/layout_assignment.cc @@ -147,12 +147,9 @@ bool LayoutConstraints::OperandBufferForwarded( PointsToSet::BufferSet* output_buffers = GetBufferSet(instruction); PointsToSet::BufferSet* operand_buffers = GetBufferSet(instruction->operand(operand_no)); - for (const LogicalBuffer* output_buffer : *output_buffers) { - if (operand_buffers->count(output_buffer) > 0) { - return true; - } - } - return false; + return absl::c_any_of(*output_buffers, [&](const LogicalBuffer* b) { + return operand_buffers->count(b) > 0; + }); } Status LayoutConstraints::SetBufferLayout(const Layout& layout, @@ -1236,6 +1233,23 @@ Status LayoutAssignment::PropagateUseConstraintToDefs( }); } +namespace { +// A transpose or a reshape that only changes trivial dimensions have meaningful +// layouts that are valuable to propagate in a depthfirst manner to avoid +// unassigned layouts in the graph. +bool InstructionShouldPropagateDepthFirst(const HloInstruction& hlo) { + switch (hlo.opcode()) { + case HloOpcode::kReshape: + return std::get<0>(hlo.ReshapeMerelyInsertsOrDeletes1SizedDimensions()); + case HloOpcode::kTranspose: + return true; + default: + return false; + } +} + +} // namespace + Status LayoutAssignment::PropagateOperandConstraint( const OperandLayoutConstraint& operand_constraint, LayoutConstraints* constraints) { @@ -1370,7 +1384,7 @@ Status LayoutAssignment::PropagateOperandConstraint( TF_RETURN_IF_ERROR(constraints->SetBufferLayout( *layout, *buffer, /*mandatory=*/user->opcode() == HloOpcode::kReduce, - /*dfs=*/false)); + /*dfs=*/InstructionShouldPropagateDepthFirst(*user))); } } return Status::OK(); @@ -1420,11 +1434,9 @@ Status LayoutAssignment::PropagateBufferConstraintToOperands( ChooseOperandLayoutFromOutputLayout(buffer_constraint.layout(), instruction, operand_no); if (operand_layout != nullptr) { - // Do not propagate operand constraints of transposes and reshapes, it - // tends to create really bad layouts. TF_RETURN_IF_ERROR(constraints->SetArrayOperandLayout( *operand_layout, instruction, operand_no, /*mandatory=*/false, - /*dfs=*/false)); + /*dfs=*/InstructionShouldPropagateDepthFirst(*instruction))); } } else { VLOG(6) << "Operand already has a constraint " @@ -1573,8 +1585,9 @@ Status SetFusionLayouts(HloInstruction* fusion) { fused_instruction->mutable_shape())); } else if (fused_instruction->opcode() == HloOpcode::kInfeed) { // Nop; leave the infeed layout alone. - } else { + } else if (fusion->fusion_kind() != HloInstruction::FusionKind::kCustom) { // Other instructions don't have layouts inside of fusion nodes. + // But do not clear layouts for other instructions in custom fusion nodes. LayoutUtil::ClearLayout(fused_instruction->mutable_shape()); } } @@ -2057,6 +2070,7 @@ bool LayoutAssignment::InstructionCanChangeLayout( case HloOpcode::kSort: case HloOpcode::kSubtract: case HloOpcode::kTanh: + case HloOpcode::kTriangularSolve: case HloOpcode::kTupleSelect: case HloOpcode::kWhile: return false; @@ -2082,6 +2096,7 @@ bool LayoutAssignment::InstructionCanChangeLayout( case HloOpcode::kRecv: case HloOpcode::kRecvDone: case HloOpcode::kReduce: + case HloOpcode::kReplicaId: case HloOpcode::kReshape: case HloOpcode::kRng: case HloOpcode::kSend: @@ -2120,7 +2135,7 @@ Status LayoutAssignment::ClearPreviousPassSideEffects(HloModule* module) { for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) { if (instruction->opcode() == HloOpcode::kCopy && - added_copies_.count(instruction) > 0) { + added_copies_.contains(instruction)) { VLOG(5) << "Removing added copy: " << instruction->ToString(); TF_RETURN_IF_ERROR( instruction->ReplaceAllUsesWith(instruction->mutable_operand(0))); diff --git a/tensorflow/compiler/xla/service/layout_assignment.h b/tensorflow/compiler/xla/service/layout_assignment.h index 3b081de3c7826c3c11a7d87d542835d0ecce1b7e..5701cb5b025e563247d46d0d24f81a5f886fc23b 100644 --- a/tensorflow/compiler/xla/service/layout_assignment.h +++ b/tensorflow/compiler/xla/service/layout_assignment.h @@ -243,7 +243,7 @@ class ChannelLayoutConstraints { // Returns true if channel_id has a layout constraint. bool IsChannelConstrained(int64 channel_id) const { - return constraints_.count(channel_id) > 0; + return constraints_.contains(channel_id); } // Given `shape`, apply the layout for `channel_id`. `channel_id` must already @@ -276,7 +276,7 @@ class ChannelLayoutConstraints { } private: - std::unordered_map constraints_; + absl::flat_hash_map constraints_; }; // HLO pass which assigns layouts to all instructions in the HLO module while diff --git a/tensorflow/compiler/xla/service/layout_assignment_test.cc b/tensorflow/compiler/xla/service/layout_assignment_test.cc index 387b385157ab6ece65c692de65b4da33038f1f30..c8cf3c47d380012fdb0206c0d20d67e6a13017ae 100644 --- a/tensorflow/compiler/xla/service/layout_assignment_test.cc +++ b/tensorflow/compiler/xla/service/layout_assignment_test.cc @@ -960,8 +960,9 @@ TEST_F(LayoutAssignmentTest, CopyDSliceOperandToAvoidImplicitLayoutChange) { ENTRY CopyDSliceOperandToAvoidImplicitLayoutChange { par0 = f32[3,4]{1,0} parameter(0) par1 = f32[4,5]{0,1} parameter(1) - par2 = s32[2] parameter(2) - dslice0 = f32[3,4] dynamic-slice(par1, par2), dynamic_slice_sizes={3,4} + par2 = s32[] parameter(2) + par3 = s32[] parameter(3) + dslice0 = f32[3,4] dynamic-slice(par1, par2, par3), dynamic_slice_sizes={3,4} ROOT add0 = f32[3,4]{1,0} add(par0,dslice0) } )"; @@ -982,7 +983,7 @@ TEST_F(LayoutAssignmentTest, CopyDSliceOperandToAvoidImplicitLayoutChange) { m::Parameter(), m::DynamicSlice( m::Copy(m::Parameter(1)).WithShapeEqualTo(&shape_copy), - m::Parameter(2))))); + m::Parameter(2), m::Parameter(3))))); } TEST_F(LayoutAssignmentTest, CopyConcatOperandToAvoidImplicitLayoutChange) { diff --git a/tensorflow/compiler/xla/service/llvm_ir/BUILD b/tensorflow/compiler/xla/service/llvm_ir/BUILD index 728a66b388f0f9af480ff88b5e96990a26e36af5..c5d59fb28e02ce229967fb3856012d608fb83c5d 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/BUILD +++ b/tensorflow/compiler/xla/service/llvm_ir/BUILD @@ -39,7 +39,6 @@ cc_library( "//tensorflow/compiler/xla/service:logical_buffer", "//tensorflow/core:lib", "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/strings", "@llvm//:core", ], @@ -169,6 +168,7 @@ cc_library( "//tensorflow/compiler/xla/service:elemental_ir_emitter", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", "@llvm//:core", diff --git a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc index 643ecd0fbaa546c551097b29e74ccd49418e1466..ce3d922ca7a9bdea3a520959a8b8d284bc3e0d64 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc @@ -81,9 +81,7 @@ void AliasAnalysis::AddAliasingInformationToIrArray(const HloInstruction& hlo, if (hlo.opcode() == HloOpcode::kParameter) { const std::vector& parameter_instructions = module_.entry_computation()->parameter_instructions(); - if (std::find(parameter_instructions.begin(), - parameter_instructions.end(), - &hlo) != parameter_instructions.end()) { + if (absl::c_linear_search(parameter_instructions, &hlo)) { array->MarkInvariantOverWholeProgram(context_); } } diff --git a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h index 2b46b3c3964b15548dbacc8b0ada0047a0fa85b6..12e2f449e23ac2511aac576fed893f5a9ef510c0 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h +++ b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h @@ -76,15 +76,12 @@ class AliasAnalysis { // A map from a buffer slice to metadata corresponding to its alias.scope // metadata. The index kParameterAliasSet is used to hold aliasing // information for parameters. - absl::flat_hash_map + absl::flat_hash_map alias_scope_metadata_; // A map from a buffer slice to metadata corresponding to its noalias // metadata. - absl::flat_hash_map - noalias_metadata_; + absl::flat_hash_map noalias_metadata_; }; } // namespace llvm_ir diff --git a/tensorflow/compiler/xla/service/llvm_ir/buffer_assignment_util.cc b/tensorflow/compiler/xla/service/llvm_ir/buffer_assignment_util.cc index bdce4a171b8a58f617f1d56e6cf6db5354846703..1ea5a42b0b398818b0946eaa9e214100007bada4 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/buffer_assignment_util.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/buffer_assignment_util.cc @@ -1,4 +1,4 @@ -/* 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. @@ -41,14 +41,26 @@ static const HloInstruction& InstrForConstantBufferAllocation( return *const_instr; } -string ConstantBufferAllocationToGlobalName( - const BufferAllocation& allocation) { - string instr_name = InstrForConstantBufferAllocation(allocation).name(); +string SanitizeConstantName(const HloInstruction& instr) { + CHECK_EQ(instr.opcode(), HloOpcode::kConstant); + string instr_name = instr.name(); for (char& c : instr_name) { - if (c == '.') { + // Having a hyphen or a dot in a global variable name can crash the LLVM PTX + // backend. + if (c == '.' || c == '-') { c = '_'; } } + return instr_name; +} + +string ConstantBufferAllocationToGlobalName( + const BufferAllocation& allocation) { + const HloInstruction& instr = InstrForConstantBufferAllocation(allocation); + string instr_name = instr.name(); + // Check that names are sanitized and stored in the HLO instructions + // before constant buffer allocation. + DCHECK_EQ(instr_name, SanitizeConstantName(instr)); return absl::StrCat("buffer_for_", instr_name); } diff --git a/tensorflow/compiler/xla/service/llvm_ir/buffer_assignment_util.h b/tensorflow/compiler/xla/service/llvm_ir/buffer_assignment_util.h index bfb6eecb87f6a1b756b3a8da3377f608dd7f0be7..03e98a66900095889292cbff9d9924a9abe83ab0 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/buffer_assignment_util.h +++ b/tensorflow/compiler/xla/service/llvm_ir/buffer_assignment_util.h @@ -20,6 +20,10 @@ limitations under the License. namespace xla { namespace llvm_ir { +// Sanitizes the HLO constant instruction name so that it can be used for the +// name of the corresponding constant buffer. In particular, it replaces . and +// - with _. +string SanitizeConstantName(const HloInstruction& instr); // In XLA:GPU we map constant buffer allocations to globals in the generated // LLVM IR. This function gives us the name of the global variable a constant // buffer is mapped to. Not used on XLA:CPU. diff --git a/tensorflow/compiler/xla/service/llvm_ir/dynamic_update_slice_util.cc b/tensorflow/compiler/xla/service/llvm_ir/dynamic_update_slice_util.cc index c66eaec8fb0e4c03f6967fec0cf0ae9661cdf470..3acceccfa556103c15fe229c41e96e618ac59c80 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/dynamic_update_slice_util.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/dynamic_update_slice_util.cc @@ -113,20 +113,10 @@ Status EmitDynamicUpdateSliceInPlace(absl::Span operand_arrays, Shape output_shape = output_array.GetShape(); Shape update_shape = update_array.GetShape(); - IndexGenerator start_indices_generator; - // TODO(b/118437727): Remove the R1 path, and rename the variables. - if (start_indices_array.GetShape().rank() == 1) { - start_indices_generator = [&](int64 index) { - return start_indices_array.EmitReadArrayElement( - IrArray::Index({b->getInt64(index)}), b); - }; - } else { - start_indices_generator = [&](int64 index) { - return operand_arrays[2 + index].EmitReadArrayElement( - IrArray::Index(b->getInt64Ty()), b); - }; - } - + IndexGenerator start_indices_generator = [&](int64 index) { + return operand_arrays[2 + index].EmitReadArrayElement( + IrArray::Index(b->getInt64Ty()), b); + }; ElementGenerator update_array_generator = [&](const IrArray::Index& index) { return update_array.EmitReadArrayElement(index, b); }; @@ -178,21 +168,11 @@ static Status EmitFusedDynamicUpdateSliceInPlaceImpl( TF_RETURN_IF_ERROR(dynamic_update_slice->Accept(&fused_emitter)); ElementGenerator update_array_generator = fused_emitter.GetGenerator(update); - // TODO(b/118437727): Remove the R1 path, and rename the variables. - IndexGenerator start_indices_generator; - if (start_indices->shape().rank() == 1) { - start_indices_generator = [&](int64 index) { - return fused_emitter.GetGenerator(start_indices)( - IrArray::Index({b->getInt64(index)})); - }; - } else { - start_indices_generator = [&](int64 index) { - ElementGenerator element_generator = - fused_emitter.GetGenerator(dynamic_update_slice->operand(2 + index)); - return element_generator(IrArray::Index(b->getInt64Ty())); - }; - } - + IndexGenerator start_indices_generator = [&](int64 index) { + ElementGenerator element_generator = + fused_emitter.GetGenerator(dynamic_update_slice->operand(2 + index)); + return element_generator(IrArray::Index(b->getInt64Ty())); + }; bool is_signed = ShapeUtil::ElementIsSigned(start_indices->shape()); return EmitDynamicUpdateSliceInPlaceImpl( update_shape, start_indices_generator, is_signed, update_array_generator, diff --git a/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.cc b/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.cc index 03a475b40b5a1f5100091fc11d744792b641cb04..e440f05e2b2f0d4a2a4c7b326b4881183de4d235 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.cc @@ -35,7 +35,7 @@ using llvm_ir::IrArray; Status FusedIrEmitter::DefaultAction(HloInstruction* hlo) { indexed_generators_[hlo] = [=](const IrArray::Index& index) -> StatusOr { - if (generated_value_cache_[hlo].count(index.multidim()) > 0) { + if (generated_value_cache_[hlo].contains(index.multidim())) { llvm::Value* generated_value = generated_value_cache_[hlo][index.multidim()]; llvm::BasicBlock* generated_value_bb = nullptr; diff --git a/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.h b/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.h index 1b9c61f6700e2a1309b21e499f4a9e2439ed3702..e6d52a580c04a920d3f0e8ed6f39c1cae587cf1b 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.h +++ b/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.h @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_map.h" #include "absl/types/optional.h" #include "absl/types/span.h" #include "llvm/IR/IRBuilder.h" @@ -134,8 +135,9 @@ class FusedIrEmitter : public DfsHloVisitorWithDefault { // Cache of generated values, lest we regenerate an element of a node with // multiple outgoing edges - std::unordered_map, llvm::Value*>> + absl::flat_hash_map< + const HloInstruction*, + absl::flat_hash_map, llvm::Value*>> generated_value_cache_; }; diff --git a/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc b/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc index f658761dff69f0441a46989fa86b96e71b7701b2..8ee07ae8331e986f9d271be5e39065f0d87853b1 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc @@ -104,8 +104,8 @@ IrArray::Index::Index(absl::Span multidim, CHECK(LayoutUtil::HasLayout(shape)); } -IrArray::IrArray(llvm::Value* base_ptr, const Shape& shape) - : base_ptr_(base_ptr), shape_(&shape) { +IrArray::IrArray(llvm::Value* base_ptr, Shape shape) + : base_ptr_(base_ptr), shape_(std::move(shape)) { TF_CHECK_OK(ShapeUtil::ValidateShape(shape)); CHECK(base_ptr_->getType()->isPointerTy()); int depth = 0; diff --git a/tensorflow/compiler/xla/service/llvm_ir/ir_array.h b/tensorflow/compiler/xla/service/llvm_ir/ir_array.h index d6d84994ee147f4b8c1a333b0eaccdf6e0a2219b..b706ebd311cbb706e7e4698b93319e37e664d10a 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/ir_array.h +++ b/tensorflow/compiler/xla/service/llvm_ir/ir_array.h @@ -130,6 +130,11 @@ class IrArray { CHECK_LE(index, size()); mutable_multidim().insert(mutable_multidim().begin() + index, value); } + void InsertAt(int64 index, int64 count, llvm::Value* value) { + CHECK_LE(index, size()); + mutable_multidim().insert(mutable_multidim().begin() + index, count, + value); + } using iterator = std::vector::iterator; using const_iterator = std::vector::const_iterator; @@ -189,6 +194,8 @@ class IrArray { return llvm::ConstantInt::get(index_type_, c); } + void ClearLinearIndex() { linear_ = nullptr; } + private: // Changing the multi-dimensional index invalidates the linear index. std::vector& mutable_multidim() { @@ -220,11 +227,11 @@ class IrArray { }; // Default constructor. Constructs an IrArray in a null status. - IrArray() : base_ptr_(nullptr), shape_(nullptr) {} + IrArray() : base_ptr_(nullptr) {} // Construct an IrArray with the given base pointer and shape. base_ptr is a // pointer type pointing to the first element(lowest address) of the array. - IrArray(llvm::Value* base_ptr, const Shape& shape); + IrArray(llvm::Value* base_ptr, Shape shape); // Default implementations of copying and moving. IrArray(IrArray&& other) = default; @@ -236,7 +243,6 @@ class IrArray { llvm::Type* GetElementLlvmType() const { return element_type_; } const Shape& GetShape() const { - CHECK(shape_ != nullptr); return *shape_; } @@ -331,7 +337,7 @@ class IrArray { llvm::Type* element_type_; // Shape of the XLA array. - const Shape* shape_; + absl::optional shape_; // The list of key/value pairs used when attaching metadata to emitted // loads/stores for this array. They keys are the metadata kinds and the diff --git a/tensorflow/compiler/xla/service/llvm_ir/ir_builder_mixin.h b/tensorflow/compiler/xla/service/llvm_ir/ir_builder_mixin.h index abc06fb7b4245294df2dc20d25a22ac4fdaeb4cf..02c719502ee7b0a732ae74acec364f89d51ae0c1 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/ir_builder_mixin.h +++ b/tensorflow/compiler/xla/service/llvm_ir/ir_builder_mixin.h @@ -254,6 +254,11 @@ class IrBuilderMixin { return mixin_builder()->CreateFCmpOLT(std::forward(args)...); } + template + llvm::Value* FCmpOLE(Args&&... args) { + return mixin_builder()->CreateFCmpOLE(std::forward(args)...); + } + template llvm::Value* FCmpONE(Args&&... args) { return mixin_builder()->CreateFCmpONE(std::forward(args)...); @@ -264,6 +269,11 @@ class IrBuilderMixin { return mixin_builder()->CreateFCmpUNE(std::forward(args)...); } + template + llvm::Value* FCmpUNO(Args&&... args) { + return mixin_builder()->CreateFCmpUNO(std::forward(args)...); + } + template llvm::Value* FDiv(Args&&... args) { return mixin_builder()->CreateFDiv(std::forward(args)...); diff --git a/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.cc b/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.cc index cebbc4290163d4e98003cd7b5df6ec906509a446..cd8dd72cd775d5e0b52f96a2326367da0775e7eb 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.cc @@ -123,7 +123,8 @@ KernelMappingScheme::KernelMappingScheme( dims_in_elems_(dims_in_elems.begin(), dims_in_elems.end()), tile_sizes_{1, tile_size_y, tile_size_x}, num_threads_x_(num_threads_x), - num_threads_y_(num_threads_y) { + num_threads_y_(num_threads_y), + dilated_x_(true) { DCHECK_EQ(dims_in_elems_.size(), 3); DCHECK_EQ(req_block_sizes.size(), 3); diff --git a/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.h b/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.h index fb633b12e60d1a9f3103fb2919ad2c3f3f14de20..f802cc27d519e621262f328903697373aa8c284c 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.h +++ b/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.h @@ -117,7 +117,10 @@ class KernelMappingScheme { int64 GetNumberOfTilesInOneBlock() const { return absl::c_accumulate(block_sizes_, 1, std::multiplies()); } - + int64 GetNumberOfTilesInOneBlockForDimension(int d) const { + DCHECK(d >= DimZ && d <= DimX); + return block_sizes_[d]; + } int64 GetNumberOfBlocks() const { return absl::c_accumulate(dims_in_blocks_, 1, std::multiplies()); } @@ -147,6 +150,16 @@ class KernelMappingScheme { GetNumberOfThreadsForDimensionY(); } + bool DilatedX() const { return dilated_x_; } + void SetDilatedX(bool v) { + dilated_x_ = v; + if (!dilated_x_) { + // dilated_x_=false is for the purpose of vectorization, which requires + // GetTileSizeForDimension(DimX) to be a multiplier of num_threads_x_. + CHECK_EQ(GetTileSizeForDimension(DimX) % num_threads_x_, 0); + } + } + IrArray::Index EmitBlockIndex(llvm::Type* index_ty); // Returns the index for the first tile in the block with the given block // index. @@ -186,6 +199,13 @@ class KernelMappingScheme { int64 num_threads_x_; // Number of threads used to process elements in the Y direction of a tile. int64 num_threads_y_; + + // When num_threads_x threads process a total of tile_size_x elements in the + // X dimension of a tile, each threads process n=tile_size_x/num_threads_x + // elements. When dilated_x=false, the n elements processed by a thread are + // contiguous. On the other hand, when dilated_x=true the n elements are + // dilated by a factor of num_threads_x. + bool dilated_x_; }; // A class to represent information for tiled parameters to support IR emission diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc b/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc index fe320bbe727111fbc986cc1fbc217feed74d30f1..3a35405a2da0af386e01bb48bed56ad194048543 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc @@ -25,7 +25,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/logging.h" namespace xla { diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc index 75c4245f5ff39f700187622a3f6ae2a8edc4d5da..807296329c07b8e4ac630486a1e1f59e4fdfa009 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc @@ -188,7 +188,16 @@ llvm::Type* PrimitiveTypeToIrType(PrimitiveType element_type, } return cplx_t; } - // A Tuple contains an array of pointers. Use i8*. + case C128: { + auto cplx_t = module->getTypeByName("complex128"); + if (cplx_t == nullptr) { + return llvm::StructType::create( + {llvm::Type::getDoubleTy(module->getContext()), + llvm::Type::getDoubleTy(module->getContext())}, + "complex128", /*isPacked=*/true); + } + return cplx_t; + } // A Tuple contains an array of pointers. Use i8*. case TUPLE: // An Opaque is like a void*, use i8*. case OPAQUE: @@ -621,6 +630,10 @@ llvm::Function* CreateFunction(llvm::FunctionType* function_type, function->setCallingConv(llvm::CallingConv::C); function->addFnAttr("no-frame-pointer-elim", "false"); + // Generate unwind information so that GDB can crawl through the stack frames + // created by the JIT compiled code. + function->setHasUWTable(); + if (enable_fast_math) { function->addFnAttr("unsafe-fp-math", "true"); function->addFnAttr("no-infs-fp-math", "true"); diff --git a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc index 0dc120e0b0df47f261435f490a8459b49d989b53..a689881e65ec3a7ddf606c36bdd64b749cfe358e 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc @@ -23,7 +23,6 @@ limitations under the License. #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/protobuf.h" diff --git a/tensorflow/compiler/xla/service/llvm_ir/sort_util.cc b/tensorflow/compiler/xla/service/llvm_ir/sort_util.cc index 89b6a36f96beedbcb7322e6164ac59221650d3d8..d71addec9b7317dfe16e9d7e5380c3cfda0b8c06 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/sort_util.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/sort_util.cc @@ -24,6 +24,7 @@ limitations under the License. #include "llvm/ADT/APInt.h" #include "llvm/IR/BasicBlock.h" #include "llvm/IR/Constants.h" +#include "llvm/IR/DerivedTypes.h" #include "llvm/IR/Instructions.h" #include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/primitive_util.h" @@ -45,13 +46,14 @@ namespace llvm_ir { namespace { // Adds the inner comparison loop body where we compare elements. -void EmitCompareLoopBody( - int64 iteration_bound, PrimitiveType key_type, int64 num_values, - int64 iota_values_parameter_index, llvm::Value* element_pair_index, +Status EmitCompareLoopBody( + int64 iteration_bound, int64 num_values, llvm::Value* element_pair_index, int64 xor_mask, llvm::Type* index_type, - std::function read_element, + std::function + element_address, std::function write_element, + const EmitCallToNestedComputationCallback& emit_compare_callback, llvm::IRBuilder<>* b, bool needs_bounds_checks = true) { auto index_typed_constant = [&](int64 value) { return llvm::ConstantInt::get(index_type, value); @@ -108,74 +110,44 @@ void EmitCompareLoopBody( // if (is_smaller_index && index_is_inbounds) KernelSupportLibrary ksl(b); - ksl.If("smaller_comparison_index", do_comparison, [&]() { - auto key1 = read_element(0, current_keys_index); - auto key2 = read_element(0, compare_keys_index); - auto compare_key1 = key1; - auto compare_key2 = key2; - bool is_signed_comparison = true; - if (primitive_util::IsFloatingPointType(key_type)) { - // We would like a total order of floating point numbers so that the - // sort has a predictable behavior in the presence of NaNs. Rather - // than using floating point comparison, we use the following trick: - // If f is a float, and - // x = bit_cast(f); - // y = x < 0 ? 0x7FFFFFFF - x : x; - // then y is ordered as an int32 such that finite values have the - // obvious order, -0 is ordered before 0, and -NaN and NaN appear at - // the beginning and end of the ordering. - auto k = b->getInt(llvm::APInt::getSignedMaxValue( - key1->getType()->getPrimitiveSizeInBits())); - auto comparison_type = k->getType(); - auto zero = llvm::ConstantInt::get(comparison_type, 0); - auto maybe_flip = [&](llvm::Value* v) { - return b->CreateSelect(b->CreateICmp(llvm::ICmpInst::ICMP_SLT, v, zero), - b->CreateSub(k, v), v); - }; - compare_key1 = b->CreateBitCast(key1, comparison_type); - compare_key2 = b->CreateBitCast(key2, comparison_type); - compare_key1 = maybe_flip(compare_key1); - compare_key2 = maybe_flip(compare_key2); - } else if (!primitive_util::IsSignedIntegralType(key_type)) { - is_signed_comparison = false; - } - // If key2 < key1 - auto is_smaller_than = - b->CreateICmp(is_signed_comparison ? llvm::ICmpInst::ICMP_SLT - : llvm::ICmpInst::ICMP_ULT, - compare_key2, compare_key1); - if (iota_values_parameter_index >= 0) { - auto keys_equal = b->CreateICmpEQ(compare_key1, compare_key2); - auto key_index1 = - read_element(iota_values_parameter_index, current_keys_index); - auto key_index2 = - read_element(iota_values_parameter_index, compare_keys_index); - auto index_is_smaller_than = - b->CreateICmp(llvm::ICmpInst::ICMP_ULT, key_index2, key_index1); - is_smaller_than = b->CreateOr( - is_smaller_than, b->CreateAnd(keys_equal, index_is_smaller_than)); + return ksl.IfWithStatus("smaller_comparison_index", do_comparison, [&]() { + std::vector values_to_compare; + for (int i = 0; i < num_values; ++i) { + values_to_compare.push_back(element_address(i, compare_keys_index)); + values_to_compare.push_back(element_address(i, current_keys_index)); } + llvm::Module* module = b->GetInsertBlock()->getParent()->getParent(); + llvm::Value* compare_return_buffer = llvm_ir::EmitAllocaAtFunctionEntry( + llvm_ir::PrimitiveTypeToIrType(PRED, module), "compare_return_buffer", + b); + TF_RETURN_IF_ERROR( + emit_compare_callback(values_to_compare, compare_return_buffer)); + llvm::Value* result = b->CreateLoad(compare_return_buffer); + + // Check if the 'compare' function returns true. + llvm::Value* is_smaller_than = + b->CreateICmpNE(result, llvm::ConstantInt::get(result->getType(), 0), + "boolean_predicate"); ksl.If("is_smaller_than", is_smaller_than, [&]() { - // Swap key1 with key2. - write_element(0, current_keys_index, key2); - write_element(0, compare_keys_index, key1); - for (int64 i = 1; i <= num_values; ++i) { - // Also swap the values. - auto value1 = read_element(i, current_keys_index); - auto value2 = read_element(i, compare_keys_index); - write_element(i, current_keys_index, value2); - write_element(i, compare_keys_index, value1); + for (int64 i = 0; i < num_values; ++i) { + // Swap the values. + auto value1 = b->CreateLoad(values_to_compare[i * 2]); + auto value2 = b->CreateLoad(values_to_compare[i * 2 + 1]); + write_element(i, current_keys_index, value1); + write_element(i, compare_keys_index, value2); } }); + return Status::OK(); }); } -void EmitTiledCompareLoop( +Status EmitTiledCompareLoop( const IrArray::Index& tiled_keys_index, int64 dimension_to_sort, - int64 dimension_to_sort_bound, PrimitiveType keys_type, - absl::Span xor_masks, const std::vector& params, - const std::vector& param_shmem_buffers, - int64 iota_values_parameter_index, int64 tile_size, llvm::IRBuilder<>* b) { + int64 dimension_to_sort_bound, absl::Span xor_masks, + const std::vector& params, + const std::vector& param_shmem_buffers, int64 tile_size, + const EmitCallToNestedComputationCallback& emit_compare_callback, + llvm::IRBuilder<>* b) { KernelSupportLibrary ksl(b); llvm::Value* thread_id = llvm_ir::EmitCallToIntrinsic( llvm::Intrinsic::nvvm_read_ptx_sreg_tid_x, {}, {}, b); @@ -200,7 +172,7 @@ void EmitTiledCompareLoop( [&]() { auto cache_index = b->CreateShl(thread_id, value_one); read_or_write(cache_index, current_keys_index); - // Increment to go the next index position. + // Increment to go to the next index position. current_keys_index = b->CreateAdd(current_keys_index, value_one); // Here we check whether the next index position is within bounds. ksl.If("inner_smaller_keys_index", @@ -230,10 +202,18 @@ void EmitTiledCompareLoop( llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::nvvm_barrier0, {}, {}, b); // Now emit the bodies of the comparison loops. - auto read_element = [&](int64 operand, llvm::Value* index) { - return b->CreateLoad( + auto element_address = [&](int64 operand, llvm::Value* index) { + auto shared_memory_address = b->CreateGEP(param_shmem_buffers[operand], - {tiled_keys_index.GetConstantWithIndexType(0), index})); + {tiled_keys_index.GetConstantWithIndexType(0), index}); + auto ptr_type = shared_memory_address->getType(); + // We need a generic pointer with address space 0 instead of a pointer to + // shared memory (address space 3) so that we can pass it to the comparison + // computation. + return b->CreateAddrSpaceCast( + shared_memory_address, + llvm::PointerType::get(ptr_type->getPointerElementType(), + /*AddressSpace=*/0)); }; auto write_element = [&](int64 operand, llvm::Value* index, llvm::Value* value) { @@ -252,7 +232,7 @@ void EmitTiledCompareLoop( if (dimension_to_sort_bound % tile_size) { // Otherwise we need a bounds check for the last tile. The last tile has // size 'dimension_to_sort_bound' % 'tile_size'. - ksl.If( + TF_RETURN_IF_ERROR(ksl.IfWithStatus( "is_last_tile", b->CreateICmpUGE( b->CreateMul(tiled_keys_index[dimension_to_sort], @@ -260,24 +240,24 @@ void EmitTiledCompareLoop( tiled_keys_index.GetConstantWithIndexType( RoundDownToNearest(dimension_to_sort_bound, tile_size))), [&]() { - EmitCompareLoopBody(dimension_to_sort_bound % tile_size, keys_type, - params.size() - 1, iota_values_parameter_index, - element_pair_index, xor_mask, - tiled_keys_index.GetType(), read_element, - write_element, b); + return EmitCompareLoopBody( + dimension_to_sort_bound % tile_size, params.size(), + element_pair_index, xor_mask, tiled_keys_index.GetType(), + element_address, write_element, emit_compare_callback, b); }, [&]() { - EmitCompareLoopBody(tile_size, keys_type, params.size() - 1, - iota_values_parameter_index, element_pair_index, - xor_mask, tiled_keys_index.GetType(), - read_element, write_element, b, - /*needs_bounds_checks=*/false); - }); + return EmitCompareLoopBody( + tile_size, params.size(), element_pair_index, xor_mask, + tiled_keys_index.GetType(), element_address, write_element, + emit_compare_callback, b, + /*needs_bounds_checks=*/false); + })); } else { - EmitCompareLoopBody(tile_size, keys_type, params.size() - 1, - iota_values_parameter_index, element_pair_index, - xor_mask, tiled_keys_index.GetType(), read_element, - write_element, b, /*needs_bounds_checks=*/false); + TF_RETURN_IF_ERROR(EmitCompareLoopBody( + tile_size, params.size(), element_pair_index, xor_mask, + tiled_keys_index.GetType(), element_address, write_element, + emit_compare_callback, b, + /*needs_bounds_checks=*/false)); } // Wait until all comparisons have happened. llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::nvvm_barrier0, {}, {}, b); @@ -301,17 +281,16 @@ void EmitTiledCompareLoop( // same location in shared memory because we have exactly tile_size / 2 many // threads, and the linear index calculated by ParallelLoopEmitter uses // linear_index = blockIdx.x * blockDim.x + threadIdx.x; + return Status::OK(); } } // namespace -Status EmitSortInPlace(int64 dimension_to_sort, const IrArray& keys_array, - const std::vector& values_arrays, - int64 iota_values_parameter_index, - absl::string_view name, - absl::Span xor_masks, llvm::IRBuilder<>* b, - const gpu::LaunchDimensions& launch_dimensions, - int64 num_iterations_in_sort_dim, - const int64 tile_size) { +Status EmitSortInPlace( + int64 dimension_to_sort, const std::vector& values_arrays, + absl::string_view name, absl::Span xor_masks, + llvm::IRBuilder<>* b, const gpu::LaunchDimensions& launch_dimensions, + int64 num_iterations_in_sort_dim, const int64 tile_size, + const EmitCallToNestedComputationCallback& emit_compare_callback) { // Iterate through the keys shape in physical order, but skip the dimension to // sort and make it the innermost loop which is the loop where the comparisons // happen. In the dimension to sort, if we use tiling, we iterate through it @@ -321,7 +300,7 @@ Status EmitSortInPlace(int64 dimension_to_sort, const IrArray& keys_array, // within those 64 elements and are therefore independent of the other // comparisons). - const Shape& keys_shape = keys_array.GetShape(); + const Shape& keys_shape = values_arrays[0].GetShape(); int64 rank = keys_shape.rank(); int64 dimension_to_sort_bound = keys_shape.dimensions(dimension_to_sort); std::vector dimensions_in_iteration_order(rank); @@ -338,18 +317,16 @@ Status EmitSortInPlace(int64 dimension_to_sort, const IrArray& keys_array, Shape iteration_shape = ShapeUtil::MakeShape(keys_shape.element_type(), dimensions_in_iteration_order); - std::vector params(1, keys_array); - params.insert(params.end(), values_arrays.begin(), values_arrays.end()); // Allocate shared memory for the tiled compare loop. - std::vector param_shmem_buffers(params.size(), nullptr); + std::vector param_shmem_buffers(values_arrays.size(), nullptr); if (xor_masks.size() > 1) { llvm::Module* module = b->GetInsertBlock()->getParent()->getParent(); - for (int64 i = 0; i < params.size(); ++i) { - llvm::Type* tile_type = - llvm::ArrayType::get(llvm_ir::PrimitiveTypeToIrType( - params[i].GetShape().element_type(), module), - tile_size); + for (int64 i = 0; i < values_arrays.size(); ++i) { + llvm::Type* tile_type = llvm::ArrayType::get( + llvm_ir::PrimitiveTypeToIrType( + values_arrays[i].GetShape().element_type(), module), + tile_size); param_shmem_buffers[i] = llvm_ir::AllocateSharedMemoryTile( module, tile_type, absl::StrCat(name, "_tile_param_", i)); } @@ -376,25 +353,24 @@ Status EmitSortInPlace(int64 dimension_to_sort, const IrArray& keys_array, keys_index[iteration_order_to_logical_order[i]] = tiles_index[i]; } if (xor_masks.size() > 1) { - EmitTiledCompareLoop(keys_index, dimension_to_sort, - dimension_to_sort_bound, keys_shape.element_type(), - xor_masks, params, param_shmem_buffers, - iota_values_parameter_index, tile_size, b); + TF_RETURN_IF_ERROR(EmitTiledCompareLoop( + keys_index, dimension_to_sort, dimension_to_sort_bound, xor_masks, + values_arrays, param_shmem_buffers, tile_size, emit_compare_callback, + b)); } else { - auto read_element = [&](int64 operand, llvm::Value* index) { + auto element_address = [&](int64 operand, llvm::Value* index) { keys_index[dimension_to_sort] = index; - return params[operand].EmitReadArrayElement(keys_index, b); + return values_arrays[operand].EmitArrayElementAddress(keys_index, b); }; auto write_element = [&](int64 operand, llvm::Value* index, llvm::Value* value) { keys_index[dimension_to_sort] = index; - params[operand].EmitWriteArrayElement(keys_index, value, b); + values_arrays[operand].EmitWriteArrayElement(keys_index, value, b); }; - EmitCompareLoopBody(dimension_to_sort_bound, keys_shape.element_type(), - values_arrays.size(), iota_values_parameter_index, - tiles_index[rank - 1], xor_masks[0], - tiles_index.GetType(), read_element, write_element, - b); + TF_RETURN_IF_ERROR(EmitCompareLoopBody( + dimension_to_sort_bound, values_arrays.size(), tiles_index[rank - 1], + xor_masks[0], tiles_index.GetType(), element_address, write_element, + emit_compare_callback, b)); } return Status::OK(); }; diff --git a/tensorflow/compiler/xla/service/llvm_ir/sort_util.h b/tensorflow/compiler/xla/service/llvm_ir/sort_util.h index 685f9383acba416f51681270e4037d56abb4b6ea..b9341a34d1f2203db6e02c3df5d607174b6d0f74 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/sort_util.h +++ b/tensorflow/compiler/xla/service/llvm_ir/sort_util.h @@ -28,19 +28,18 @@ limitations under the License. namespace xla { namespace llvm_ir { +using EmitCallToNestedComputationCallback = + std::function, llvm::Value*)>; // Emits llvm IR to do pairwise comparisons/swaps in the 'dimension_to_sort' -// dimension of 'keys_array'. All other dimensions are kept as-is. This -// implements the inner loop of BitonicSort. It is assumed that 'xor_masks' -// contains only powers of 2, or values 2^k - 1 (k > 0). If -// 'iota_values_parameter_index' is >= 0, it points at a 'values_arrays' operand -// that is a iota and can be used to make the sorting stable. -Status EmitSortInPlace(int64 dimension_to_sort, const IrArray& keys_array, - const std::vector& values_arrays, - int64 iota_values_parameter_index, - absl::string_view name, - absl::Span xor_masks, llvm::IRBuilder<>* b, - const gpu::LaunchDimensions& launch_dimensions, - int64 num_iterations_in_sort_dim, int64 tile_size); +// dimension of each array in 'values_arrays'. All other dimensions are kept +// as-is. This implements the inner loop of BitonicSort. It is assumed that +// 'xor_masks' contains only powers of 2, or values 2^k - 1 (k > 0). +Status EmitSortInPlace( + int64 dimension_to_sort, const std::vector& values_arrays, + absl::string_view name, absl::Span xor_masks, + llvm::IRBuilder<>* b, const gpu::LaunchDimensions& launch_dimensions, + int64 num_iterations_in_sort_dim, int64 tile_size, + const EmitCallToNestedComputationCallback& emit_compare_callback); } // namespace llvm_ir } // namespace xla diff --git a/tensorflow/compiler/xla/service/local_service.cc b/tensorflow/compiler/xla/service/local_service.cc index 600b069ecdbabf6b05e6abb3a6b8d9b1a4b0ecf4..3470fe5b2c34bf832207ed546fad176319446f31 100644 --- a/tensorflow/compiler/xla/service/local_service.cc +++ b/tensorflow/compiler/xla/service/local_service.cc @@ -110,6 +110,7 @@ ExecutionOptions CreateExecutionOptions( *execution_options.mutable_shape_with_output_layout() = result_shape.ToProto(); } + execution_options.set_num_replicas(build_options.num_replicas()); return execution_options; } diff --git a/tensorflow/compiler/xla/service/multi_output_fusion.cc b/tensorflow/compiler/xla/service/multi_output_fusion.cc index 9ccdd7d8d818b9fa3aa77cdd10d37ca18928b448..53d52d9a3d918fa6dee093668923fcfff963d084 100644 --- a/tensorflow/compiler/xla/service/multi_output_fusion.cc +++ b/tensorflow/compiler/xla/service/multi_output_fusion.cc @@ -198,7 +198,7 @@ void MultiOutputFusion::Update(HloInstruction* instr1, HloInstruction* instr2) { if (instr == fusion || is_fused(instr) || is_connected(fusion, instr)) { continue; } - if (in_list.count(instr) > 0) { + if (in_list.contains(instr)) { continue; } int64 profit = GetProfit(instr, fusion); diff --git a/tensorflow/compiler/xla/service/name_uniquer.cc b/tensorflow/compiler/xla/service/name_uniquer.cc index daa718879ddd45afb02725b557380b2f49fe833e..e55b83d17e90bc2ca0053a0421cf80ef6edd5bca 100644 --- a/tensorflow/compiler/xla/service/name_uniquer.cc +++ b/tensorflow/compiler/xla/service/name_uniquer.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/name_uniquer.h" +#include "absl/strings/ascii.h" #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" #include "tensorflow/compiler/xla/primitive_util.h" @@ -28,13 +29,13 @@ namespace { bool IsAllowed(char character) { auto c = static_cast(character); - return (isalnum(c) != 0) || c == '_' || c == '.' || c == '-'; + return (absl::ascii_isalnum(c) != 0) || c == '_' || c == '.' || c == '-'; } } // namespace NameUniquer::NameUniquer(const string& separator) { - CHECK(std::all_of(separator.begin(), separator.end(), IsAllowed)) + CHECK(absl::c_all_of(separator, IsAllowed)) << "separator should comprises allowed characters only"; separator_ = separator; } @@ -46,7 +47,7 @@ NameUniquer::NameUniquer(const string& separator) { string result = name; char c = static_cast(result[0]); - if (!isalpha(c) && c != '_') { + if (!absl::ascii_isalpha(c) && c != '_') { result[0] = '_'; } for (int i = 1; i < result.length(); i++) { diff --git a/tensorflow/compiler/xla/service/op_expander_pass.cc b/tensorflow/compiler/xla/service/op_expander_pass.cc new file mode 100644 index 0000000000000000000000000000000000000000..87f0886a9737807bdf6f00921b813d21b69a18cc --- /dev/null +++ b/tensorflow/compiler/xla/service/op_expander_pass.cc @@ -0,0 +1,44 @@ +/* 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/op_expander_pass.h" + +#include + +#include "absl/algorithm/container.h" +#include "tensorflow/compiler/xla/service/hlo_creation_utils.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/util.h" + +namespace xla { + +StatusOr OpExpanderPass::Run(HloModule* module) { + std::vector matching_instructions; + for (HloComputation* computation : module->MakeNonfusionComputations()) { + absl::c_copy_if( + computation->instructions(), std::back_inserter(matching_instructions), + [&](HloInstruction* inst) { return InstructionMatchesPattern(inst); }); + } + + for (HloInstruction* inst : matching_instructions) { + TF_ASSIGN_OR_RETURN(HloInstruction * expanded_root, + ExpandInstruction(inst)); + TF_RETURN_IF_ERROR(inst->parent()->ReplaceInstruction(inst, expanded_root)); + } + + return !matching_instructions.empty(); +} +} // namespace xla diff --git a/tensorflow/compiler/xla/service/op_expander_pass.h b/tensorflow/compiler/xla/service/op_expander_pass.h new file mode 100644 index 0000000000000000000000000000000000000000..794849d354bef6e2b0b79e6d03af4ed851dfbdb3 --- /dev/null +++ b/tensorflow/compiler/xla/service/op_expander_pass.h @@ -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. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_OP_EXPANDER_PASS_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_OP_EXPANDER_PASS_H_ + +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" + +namespace xla { + +// This pass is an abstract superclass for passes that replace operations that +// match a pattern. It is intended to be subclassed, not used directly. +// +// This pass is useful for legalizing HLO instructions that a particular backend +// does not support into other HLO instructions. +class OpExpanderPass : public HloModulePass { + public: + StatusOr Run(HloModule* module) override; + + protected: + // Returns `true` if `instruction` should be expanded by this pass. + virtual bool InstructionMatchesPattern(HloInstruction* instruction) = 0; + + // Returns a replacement for `instruction`. + virtual StatusOr ExpandInstruction( + HloInstruction* instruction) = 0; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_OP_EXPANDER_PASS_H_ diff --git a/tensorflow/compiler/xla/service/optimize_input_output_buffer_alias.cc b/tensorflow/compiler/xla/service/optimize_input_output_buffer_alias.cc new file mode 100644 index 0000000000000000000000000000000000000000..701c629add52a217f16877a085b9ef2d096623d9 --- /dev/null +++ b/tensorflow/compiler/xla/service/optimize_input_output_buffer_alias.cc @@ -0,0 +1,106 @@ +/* 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/optimize_input_output_buffer_alias.h" + +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/memory/memory.h" +#include "absl/strings/str_format.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { +namespace { + +// Returns true if the given shape is a non-nested tuple. +bool IsNonNestedTuple(const Shape& shape) { + return shape.IsTuple() && !ShapeUtil::IsNestedTuple(shape); +} + +} // namespace + +StatusOr OptimizeInputOutputBufferAlias::Build( + const Shape& input_shape, const Shape& output_shape, + HloInputOutputAliasConfig* alias_config) { + bool changed = false; + TF_RET_CHECK(LayoutUtil::HasLayout(input_shape)); + TF_RET_CHECK(LayoutUtil::HasLayout(output_shape)); + VLOG(1) << "input_shape:" << input_shape.ToString(); + VLOG(1) << "output_shape:" << output_shape.ToString(); + + // For all buffers defined by the parameter, build a map from the byte + // size to the list of the buffers of that size. + absl::flat_hash_map> size_to_input_index; + ShapeUtil::ForEachSubshape( + input_shape, [&](const Shape& subshape, const ShapeIndex& index) { + if (subshape.IsTuple()) { + return; + } + int64 bytes = size_func_(subshape); + size_to_input_index[bytes].push(index); + }); + + // For each result buffer shape index, take the first unused parameter + // buffer that matches the size. + TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( + output_shape, [&](const Shape& subshape, const ShapeIndex& index) { + if (subshape.IsTuple()) { + return Status::OK(); + } + int64 bytes = size_func_(subshape); + + auto it = size_to_input_index.find(bytes); + if (it != size_to_input_index.end() && !it->second.empty()) { + changed = true; + const ShapeIndex& input_index = it->second.front(); + const ShapeIndex& output_index = index; + if (!alias_config->ParameterHasAlias(0, input_index) && + !alias_config->OutputHasAlias(output_index)) { + TF_RETURN_IF_ERROR(alias_config->SetUpAlias( + output_index, 0, input_index, + HloInputOutputAliasConfig::AliasKind::kSystemAlias)); + } + VLOG(3) << "Set up alias from with param index " + << it->second.front().ToString() << ", shape size " << bytes + << " and result subshape " + << ShapeUtil::HumanStringWithLayout(subshape) << " at index " + << index.ToString(); + it->second.pop(); + } + return Status::OK(); + })); + return changed; +} + +StatusOr OptimizeInputOutputBufferAlias::Run(HloModule* module) { + // User buffer alias only work for modules with 1 parameter. + if (module->entry_computation()->num_parameters() != 1) { + return false; + } + + HloInputOutputAliasConfig* alias_config = + &module->input_output_alias_config(); + + return Build(module->entry_computation()->parameter_instruction(0)->shape(), + module->entry_computation()->root_instruction()->shape(), + alias_config); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/optimize_input_output_buffer_alias.h b/tensorflow/compiler/xla/service/optimize_input_output_buffer_alias.h new file mode 100644 index 0000000000000000000000000000000000000000..79ce468e975300ed703ae0fd780f4b9d5328a4b3 --- /dev/null +++ b/tensorflow/compiler/xla/service/optimize_input_output_buffer_alias.h @@ -0,0 +1,71 @@ +/* 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_OPTIMIZE_INPUT_OUTPUT_BUFFER_ALIAS_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_OPTIMIZE_INPUT_OUTPUT_BUFFER_ALIAS_H_ + +#include + +#include "absl/container/flat_hash_map.h" +#include "tensorflow/compiler/xla/service/hlo_input_output_alias_config.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" +#include "tensorflow/compiler/xla/shape_tree.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/gtl/flatmap.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +// This pass opportunistically finds input and output buffers that can be +// aliased, and writes the alias config into the HloModule. +// +// The input and the output buffers can be in any shape, and each output buffer +// can alias with an input buffer with the same size. Each input buffer may only +// alias with a single output buffer. For example, for the following parameter +// and the output buffers, +// +// Parameters : { P1(2MiB), P2(4MiB), P3(8MiB), P4(4MiB), P5(4MiB), ... } +// Outputs : { O1(4MiB), O2(2MiB), O3(4MiB), O4(6MiB), O5(4MiB), ... } +// +// one potential aliasing would be (O1, P2), (O2, P1), (O3, P4), (O5, P5), .. +class OptimizeInputOutputBufferAlias : public HloModulePass { + using ShapeSizeFunction = std::function; + + public: + OptimizeInputOutputBufferAlias(ShapeSizeFunction size_func) + : size_func_(size_func) {} + ~OptimizeInputOutputBufferAlias() override = default; + + absl::string_view name() const override { + return "optimize_input_output_buffer_alias.h"; + } + + StatusOr Run(HloModule* module) override; + + private: + friend class OptimizeInputOutputBufferAliasTest; + + StatusOr Build(const Shape& input_shape, const Shape& output_shape, + HloInputOutputAliasConfig* alias_config); + ShapeSizeFunction size_func_ = nullptr; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_OPTIMIZE_INPUT_OUTPUT_BUFFER_ALIAS_H_ diff --git a/tensorflow/compiler/xla/service/optimize_input_output_buffer_alias_test.cc b/tensorflow/compiler/xla/service/optimize_input_output_buffer_alias_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..41e90f9b6931619fd9824e2eda25e12e4c7197b0 --- /dev/null +++ b/tensorflow/compiler/xla/service/optimize_input_output_buffer_alias_test.cc @@ -0,0 +1,145 @@ +/* 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/optimize_input_output_buffer_alias.h" + +#include + +#include "absl/memory/memory.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_utils.h" +#include "tensorflow/core/platform/test.h" + +namespace xla { + +// Tests that UserBufferAlias properly maps input and output buffer indices of +// various shapes for aliasing. +class OptimizeInputOutputBufferAliasTest : public HloTestBase { + protected: + OptimizeInputOutputBufferAliasTest() { + r1f32_ = ShapeUtil::MakeShape(F32, {4}); + r2f32_ = ShapeUtil::MakeShape(F32, {4, 5}); + r3f32_ = ShapeUtil::MakeShape(F32, {4, 5, 6}); + r4f32_ = ShapeUtil::MakeShape(F32, {4, 5, 6, 7}); + + auto size_func = [](const Shape& shape) { + return ShapeUtil::ByteSizeOf(shape); + }; + + optimize_pass_ = + absl::make_unique(size_func); + } + + // Returns the number of output indices that aliases with the input. + int64 AliasCount() { + int64 count = 0; + + config_.ForEachAlias( + [&](const ShapeIndex&, const HloInputOutputAliasConfig::Alias&) { + count++; + }); + return count; + } + + bool BuildAliasConfig(const Shape& input_shape, const Shape& output_shape) { + config_ = HloInputOutputAliasConfig(output_shape); + auto changed = optimize_pass_->Build(input_shape, output_shape, &config_); + TF_CHECK_OK(changed.status()); + + return changed.ValueOrDie(); + } + + std::unique_ptr optimize_pass_; + + HloInputOutputAliasConfig config_; + + Shape r1f32_; + Shape r2f32_; + Shape r3f32_; + Shape r4f32_; +}; + +// All shapes are different, so no aliasing is available. +TEST_F(OptimizeInputOutputBufferAliasTest, AllDifferentBufferSizes) { + Shape input = ShapeUtil::MakeTupleShape({r1f32_, r2f32_}); + Shape output = ShapeUtil::MakeTupleShape({r3f32_, r4f32_}); + bool changed = BuildAliasConfig(input, output); + EXPECT_FALSE(changed); + EXPECT_EQ(AliasCount(), 0); +} + +// Input and output shapes are equal, so buffers can alias at the same index. +TEST_F(OptimizeInputOutputBufferAliasTest, OrderedNonNestedTuple) { + Shape input = ShapeUtil::MakeTupleShape({r1f32_, r2f32_, r3f32_, r4f32_}); + Shape output = ShapeUtil::MakeTupleShape({r1f32_, r2f32_, r3f32_, r4f32_}); + bool changed = BuildAliasConfig(input, output); + EXPECT_TRUE(changed); + EXPECT_EQ(AliasCount(), 4); + + EXPECT_EQ(config_.GetAliasedOutput(0, {0}), ShapeIndex{0}); + EXPECT_EQ(config_.GetAliasedOutput(0, {1}), ShapeIndex{1}); + EXPECT_EQ(config_.GetAliasedOutput(0, {2}), ShapeIndex{2}); + EXPECT_EQ(config_.GetAliasedOutput(0, {3}), ShapeIndex{3}); +} + +// Only a subset of the tuple element shapes match between the input and the +// output. +TEST_F(OptimizeInputOutputBufferAliasTest, PartialReuseNonNestedTuple) { + Shape input = ShapeUtil::MakeTupleShape({r1f32_, r1f32_, r2f32_, r2f32_}); + Shape output = ShapeUtil::MakeTupleShape({r1f32_, r2f32_, r3f32_, r4f32_}); + bool changed = BuildAliasConfig(input, output); + EXPECT_TRUE(changed); + + EXPECT_EQ(AliasCount(), 2); + + EXPECT_EQ(config_.GetAliasedOutput(0, {0}), ShapeIndex{0}); + EXPECT_EQ(config_.GetAliasedOutput(0, {2}), ShapeIndex{1}); +} + +// The output shape is reverse of the input shape, but we can still reuse all +// the buffers. +TEST_F(OptimizeInputOutputBufferAliasTest, UnorderedNonNestedTuple) { + Shape input = ShapeUtil::MakeTupleShape({r1f32_, r2f32_, r3f32_, r4f32_}); + Shape output = ShapeUtil::MakeTupleShape({r4f32_, r3f32_, r2f32_, r1f32_}); + bool changed = BuildAliasConfig(input, output); + EXPECT_TRUE(changed); + + EXPECT_EQ(AliasCount(), 4); + + EXPECT_EQ(config_.GetAliasedOutput(0, {0}), ShapeIndex{3}); + EXPECT_EQ(config_.GetAliasedOutput(0, {1}), ShapeIndex{2}); + EXPECT_EQ(config_.GetAliasedOutput(0, {2}), ShapeIndex{1}); + EXPECT_EQ(config_.GetAliasedOutput(0, {3}), ShapeIndex{0}); +} + +TEST_F(OptimizeInputOutputBufferAliasTest, UnorderedNestedTuple) { + Shape input = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeTupleShape({r1f32_}), r2f32_, r3f32_, r4f32_}); + Shape output = ShapeUtil::MakeTupleShape( + {r1f32_, ShapeUtil::MakeTupleShape({r3f32_, r2f32_}), r2f32_}); + bool changed = BuildAliasConfig(input, output); + EXPECT_TRUE(changed); + + EXPECT_EQ(AliasCount(), 3); + + EXPECT_EQ(config_.GetAliasedOutput(0, {0, 0}), ShapeIndex{0}); + EXPECT_EQ(config_.GetAliasedOutput(0, {1}), ShapeIndex({1, 1})); + EXPECT_EQ(config_.GetAliasedOutput(0, {2}), ShapeIndex({1, 0})); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/pattern_matcher.h b/tensorflow/compiler/xla/service/pattern_matcher.h index 802b2c0c8533eabb22c73fb214463b2b6dc2484b..9e3d1060210790f60243195a1c1dff13f1fc7fc5 100644 --- a/tensorflow/compiler/xla/service/pattern_matcher.h +++ b/tensorflow/compiler/xla/service/pattern_matcher.h @@ -2118,7 +2118,6 @@ XLA_BINOP_PATTERN(Divide) XLA_BINOP_PATTERN(Complex) XLA_BINOP_PATTERN(Convolution) XLA_BINOP_PATTERN(Dot) -XLA_BINOP_PATTERN(DynamicSlice) XLA_COMMUTATIVE_BINOP_PATTERN(Eq) XLA_BINOP_PATTERN(Gather) XLA_BINOP_PATTERN(Ge) @@ -2235,6 +2234,7 @@ inline auto WithOperands(Matcher&& m, int64 operand_num, FirstArg&& first_arg, XLA_VARIADIC_OP_PATTERN(AfterAll); XLA_VARIADIC_OP_PATTERN(Concatenate); XLA_VARIADIC_OP_PATTERN(CustomCall); +XLA_VARIADIC_OP_PATTERN(DynamicSlice) XLA_VARIADIC_OP_PATTERN(Map) XLA_VARIADIC_OP_PATTERN(Reduce); XLA_VARIADIC_OP_PATTERN(Sort); diff --git a/tensorflow/compiler/xla/service/pattern_matcher_gmock_test.cc b/tensorflow/compiler/xla/service/pattern_matcher_gmock_test.cc index 9ca2fb05c1f7ef093c58237cf21fbc7c813a592a..f51a18b13894d75300c46835fabd82a4ce0699af 100644 --- a/tensorflow/compiler/xla/service/pattern_matcher_gmock_test.cc +++ b/tensorflow/compiler/xla/service/pattern_matcher_gmock_test.cc @@ -23,7 +23,6 @@ namespace xla { namespace { namespace m = ::xla::match; -using ::testing::Eq; using ::testing::Not; template diff --git a/tensorflow/compiler/xla/service/platform_util.cc b/tensorflow/compiler/xla/service/platform_util.cc index 896b73cda41cb21b539b586aa4701c5bad43f8b9..886a0545624927fa77528141f61d8ecb6bec180a 100644 --- a/tensorflow/compiler/xla/service/platform_util.cc +++ b/tensorflow/compiler/xla/service/platform_util.cc @@ -70,6 +70,9 @@ PlatformUtil::GetSupportedPlatforms() { for (se::Platform* platform : all_platforms) { auto compiler_status = Compiler::GetForPlatform(platform); if (compiler_status.ok()) { + if (!platform->Initialized()) { + TF_RETURN_IF_ERROR(platform->Initialize({})); + } platforms.push_back(platform); } else { LOG(INFO) << "platform " << platform->Name() << " present but no " @@ -260,8 +263,8 @@ PlatformUtil::GetStreamExecutors( // Block here in thread_pool destructor until all devices are initialized. } VLOG(1) << "Device initialization complete"; - if (std::all_of(stream_executors.begin(), stream_executors.end(), - [](se::StreamExecutor* s) { return s == nullptr; })) { + if (absl::c_all_of(stream_executors, + [](se::StreamExecutor* s) { return s == nullptr; })) { return InternalError("no supported devices found for platform %s", platform->Name()); } diff --git a/tensorflow/compiler/xla/service/scatter_expander.cc b/tensorflow/compiler/xla/service/scatter_expander.cc index 8b9955faf8e4bb9271ae28cf5b3f9fc80e365cc0..acad871c4d427b174ffce3a462a0a3918a1e0c33 100644 --- a/tensorflow/compiler/xla/service/scatter_expander.cc +++ b/tensorflow/compiler/xla/service/scatter_expander.cc @@ -59,7 +59,7 @@ static StatusOr CanonicalizeScatterIndices( TF_ASSIGN_OR_RETURN( HloInstruction * transposed_scatter_indices, TransposeIndexVectorDimToLast(scatter_indices, index_vector_dim)); - if (ShapeUtil::Rank(scatter_indices->shape()) == index_vector_dim + 1 && + if (scatter_indices->shape().rank() == index_vector_dim + 1 && scatter_indices->shape().dimensions(index_vector_dim) == 1) { auto new_shape = ShapeUtil::DeleteDimension(index_vector_dim, scatter_indices->shape()); @@ -171,10 +171,9 @@ static StatusOr CheckIndexValidity( // Valid range for the index: [0, operand_dims - window_sizes] // Check if the index has any negative values. - TF_ASSIGN_OR_RETURN( - HloInstruction * zero_index, + HloInstruction* zero_index = BroadcastZeros(computation, index->shape().element_type(), - AsInt64Slice(index->shape().dimensions()))); + AsInt64Slice(index->shape().dimensions())); TF_ASSIGN_OR_RETURN(HloInstruction * negative_index_check, MakeBinaryHlo(HloOpcode::kLe, zero_index, index)); @@ -222,10 +221,9 @@ static StatusOr> ScatterLoopBody( bool has_scalar_indices = scatter_indices->shape().dimensions_size() == 1; // Build a vector form of the induction variable of the while loop. - TF_ASSIGN_OR_RETURN( - HloInstruction * induction_var_as_vector, + HloInstruction* induction_var_as_vector = MakeBroadcastHlo(induction_var, /*broadcast_dimensions=*/{}, - /*result_shape_bounds=*/{1})); + /*result_shape_bounds=*/{1}); // Pick the index to scatter from scatter_indices based on the induction_var // and transform that to an index into the `operand` space. diff --git a/tensorflow/compiler/xla/service/service.cc b/tensorflow/compiler/xla/service/service.cc index efefe9003fb253aa1fe1bedbac9978fbd3b126fe..9bda6fba3aabfed78ae724545387e86bad36c886 100644 --- a/tensorflow/compiler/xla/service/service.cc +++ b/tensorflow/compiler/xla/service/service.cc @@ -29,6 +29,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/service/computation_layout.h" #include "tensorflow/compiler/xla/service/device_memory_allocator.h" +#include "tensorflow/compiler/xla/service/dynamic_dimension_inference.h" #include "tensorflow/compiler/xla/service/executable.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_cost_analysis.h" @@ -295,11 +296,16 @@ StatusOr> Service::CreateModuleConfig( computation_layout->mutable_result_layout()->SetToDefaultLayout(); } - config->set_replica_count(options_.number_of_replicas()); if (execution_options != nullptr) { + if (execution_options->num_replicas() > 0) { + config->set_replica_count(execution_options->num_replicas()); + } else { + config->set_replica_count(options_.number_of_replicas()); + } config->set_seed(execution_options->seed()); config->set_debug_options(execution_options->debug_options()); } else { + config->set_replica_count(options_.number_of_replicas()); config->set_debug_options(GetDebugOptionsFromFlags()); } @@ -362,6 +368,7 @@ StatusOr>> Service::BuildExecutables( const HloModuleProto* proto = module_protos[i]; const HloModuleConfig& config = *module_configs[i]; TF_ASSIGN_OR_RETURN(auto module, CreateModuleFromProto(*proto, config)); + TF_RETURN_IF_ERROR(MaybeDumpUnoptimizedHloModule(*module)); module_group->push_back(std::move(module)); } @@ -523,13 +530,13 @@ Service::ExecuteParallelAndRegisterResult( StatusOr Service::ExecuteAndRegisterResult( Executable* executable, - const absl::Span> arguments, - Backend* backend, const string& result_tag, ExecutionProfile* profile) { + absl::Span> arguments, + Backend* backend, const DeviceHandle& device_handle, + const string& result_tag, ExecutionProfile* profile) { // Set up streams. std::vector streams; - TF_ASSIGN_OR_RETURN(auto replicas, - Replicas(*backend, SingleComputationDeviceHandle())); + TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*backend, device_handle)); TF_RET_CHECK(!replicas.empty()); for (se::StreamExecutor* executor : replicas) { TF_ASSIGN_OR_RETURN(StreamPool::Ptr stream, @@ -537,10 +544,11 @@ StatusOr Service::ExecuteAndRegisterResult( streams.push_back(std::move(stream)); } - TF_ASSIGN_OR_RETURN(DeviceAssignment device_assignment, - backend->computation_placer()->AssignDevices( - options_.number_of_replicas(), - /*computation_count=*/1)); + DeviceAssignment device_assignment(options_.number_of_replicas(), + /*computation_count=*/1); + for (int64 replica = 0; replica < replicas.size(); ++replica) { + device_assignment(replica, 0) = replicas[replica]->device_ordinal(); + } // Set up run options. std::vector run_options; @@ -552,9 +560,7 @@ StatusOr Service::ExecuteAndRegisterResult( options.set_intra_op_thread_pool( backend->eigen_intra_op_thread_pool_device()); options.set_device_assignment(&device_assignment); - run_options.emplace_back( - options, backend->StreamBorrower(), - /*xla_intra_op_thread_pool=*/backend->eigen_intra_op_thread_pool()); + run_options.emplace_back(options, backend->StreamBorrower()); } if (options_.number_of_replicas() == 1) { @@ -711,14 +717,33 @@ Status Service::ExecuteGraphParallel(const ExecuteGraphParallelRequest* arg, } } - // Execute the generated executables in parallel and return the device - // handles for each computation's output. + // If we have multiple executables to run, execute them all in parallel. But + // if we only have one executable, execute it using the vanilla, non-parallel + // call. + // + // We do this because the Client API uses ExecuteGraphParallel when it wants + // to compile and run one computation without caching the executable, but not + // all backends support the async StreamExecutor API required by + // ExecuteParallelAndRegisterResult. + // + // TODO(b/122731460): Consolidate Execute{,Parallel}AndRegisterResult; they do + // basically the same thing. ExecutionProfile profile; - TF_ASSIGN_OR_RETURN( - std::vector outputs, - ExecuteParallelAndRegisterResult(executable_ptrs, all_arguments, - execute_backend_.get(), device_handles, - computation_names, &profile)); + std::vector outputs; + if (executable_ptrs.size() == 1) { + TF_ASSIGN_OR_RETURN( + auto output, + ExecuteAndRegisterResult(executable_ptrs[0], all_arguments[0], + execute_backend_.get(), device_handles[0], + computation_names[0], &profile)); + outputs.push_back(std::move(output)); + } else { + TF_ASSIGN_OR_RETURN( + outputs, ExecuteParallelAndRegisterResult( + executable_ptrs, all_arguments, execute_backend_.get(), + device_handles, computation_names, &profile)); + } + for (const GlobalDataHandle& output : outputs) { ExecuteResponse response; *response.mutable_output() = output; @@ -904,6 +929,7 @@ Status Service::Execute(const ExecuteRequest* arg, ExecuteResponse* result) { *result->mutable_output(), ExecuteAndRegisterResult(executable.get(), replicated_arguments, execute_backend_.get(), + SingleComputationDeviceHandle(), "result of " + executable->module().name(), result->mutable_profile())); @@ -1097,7 +1123,11 @@ Status Service::ComputeConstantGraph(const ComputeConstantGraphRequest* arg, TF_ASSIGN_OR_RETURN(std::unique_ptr module, CreateModuleFromProto(arg->computation(), config)); + TF_ASSIGN_OR_RETURN(DynamicDimensionInference dynamic_dimension_inference, + DynamicDimensionInference::Run(module.get())); + HloEvaluator evaluator; + evaluator.set_dynamic_dimension_inference(&dynamic_dimension_inference); TF_ASSIGN_OR_RETURN(auto result_literal, evaluator.Evaluate(*module, {})); // Since the result layout is non-effective to the Evaluator results, explicit diff --git a/tensorflow/compiler/xla/service/service.h b/tensorflow/compiler/xla/service/service.h index abd3ee5a059ac0910d6acc8076899950498b4c43..fd907d07daef9e8337aeed198ef4fd23d069df21 100644 --- a/tensorflow/compiler/xla/service/service.h +++ b/tensorflow/compiler/xla/service/service.h @@ -53,7 +53,7 @@ class ServiceOptions { ServiceOptions& set_platform(se::Platform* platform); se::Platform* platform() const; - // Set the number of replicas to use when compiling replicated + // Set the default number of replicas to use when compiling replicated // programs. ServiceOptions& set_number_of_replicas(int number_of_replicas); int number_of_replicas() const; @@ -250,8 +250,9 @@ class Service : public ServiceInterface { // ExecutionProfile object which will be filled in with profile data. StatusOr ExecuteAndRegisterResult( Executable* executable, - const absl::Span> arguments, - Backend* backend, const string& result_tag, ExecutionProfile* profile); + absl::Span> arguments, + Backend* backend, const DeviceHandle& device_handle, + const string& result_tag, ExecutionProfile* profile); // Runs the given executables with the given arguments and register the result // from each executable in the allocation tracker. The handles of the result diff --git a/tensorflow/compiler/xla/service/service_executable_run_options.h b/tensorflow/compiler/xla/service/service_executable_run_options.h index dbfed628bfcabffe66bef41a82e0e2430897d80d..6bee671056552b83014367889320b748659bbfdf 100644 --- a/tensorflow/compiler/xla/service/service_executable_run_options.h +++ b/tensorflow/compiler/xla/service/service_executable_run_options.h @@ -32,12 +32,10 @@ class ServiceExecutableRunOptions { ServiceExecutableRunOptions() : ServiceExecutableRunOptions(ExecutableRunOptions()) {} - explicit ServiceExecutableRunOptions( - ExecutableRunOptions run_options, StreamBorrower borrow_stream = nullptr, - tensorflow::thread::ThreadPool* xla_intra_op_thread_pool = nullptr) + explicit ServiceExecutableRunOptions(ExecutableRunOptions run_options, + StreamBorrower borrow_stream = nullptr) : run_options_(std::move(run_options)), - borrow_stream_(std::move(borrow_stream)), - xla_intra_op_thread_pool_(xla_intra_op_thread_pool) {} + borrow_stream_(std::move(borrow_stream)) {} // Returns reference or pointer to `ExecutableRunOptions` member. const ExecutableRunOptions& run_options() const { return run_options_; } @@ -56,15 +54,9 @@ class ServiceExecutableRunOptions { : Status(tensorflow::error::UNIMPLEMENTED, "No stream cache"); } - // Returns reference to thread pool for execution of XLA ops on CPU backend. - tensorflow::thread::ThreadPool* xla_intra_op_thread_pool() const { - return xla_intra_op_thread_pool_; - } - private: ExecutableRunOptions run_options_; StreamBorrower borrow_stream_; - tensorflow::thread::ThreadPool* xla_intra_op_thread_pool_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/shape_inference.cc b/tensorflow/compiler/xla/service/shape_inference.cc index 4a2dd097425f3a5fef83f3ecf226cead4e6e9e9d..a570ee346d2f50e6eb2a592452bdec423556a916 100644 --- a/tensorflow/compiler/xla/service/shape_inference.cc +++ b/tensorflow/compiler/xla/service/shape_inference.cc @@ -156,6 +156,14 @@ Status VerifyReducerShape(const ProgramShape& reducer_shape, return Status::OK(); } +bool IsTrivialWindowDimension(const WindowDimension& window_dimension) { + return window_dimension.size() == 1 && window_dimension.stride() == 1 && + window_dimension.padding_low() == 0 && + window_dimension.padding_high() == 0 && + window_dimension.window_dilation() == 1 && + window_dimension.base_dilation() == 1; +} + StatusOr InferWindowOutputShape(const Shape& base_shape, const Window& window, PrimitiveType element_type, @@ -167,6 +175,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, } std::vector output_dimensions(window.dimensions_size()); + std::vector output_is_dynamic(window.dimensions_size()); for (int64 i = 0; i < window.dimensions_size(); ++i) { const auto& dim = window.dimensions(i); if (dim.size() <= 0) { @@ -196,6 +205,12 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, window.DebugString()); } + if (base_shape.is_dynamic_dimension(i) && !IsTrivialWindowDimension(dim)) { + return Unimplemented( + "Dynamic shape is not supported for non trivial window: %s", + window_util::ToString(window)); + } + const int64 dilated_base = window_util::DilatedBound( ShapeUtil::GetDimension(base_shape, i), dim.base_dilation()); const int64 padded_dilated_base = @@ -205,9 +220,11 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, output_dimensions[i] = window_util::StridedBound( padded_dilated_base, dilated_window, dim.stride()); + output_is_dynamic[i] = base_shape.is_dynamic_dimension(i); } - return ShapeUtil::MakeValidatedShape(element_type, output_dimensions); + return ShapeUtil::MakeValidatedShape(element_type, output_dimensions, + output_is_dynamic); } } // namespace @@ -500,17 +517,33 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, padding_config.ShortDebugString()); } + if (!padding_value_shape.is_static()) { + return InvalidArgument("Dynamic padding value is not supported"); + } + std::vector dimensions(operand_shape.rank()); + std::vector is_dynamic(operand_shape.rank()); for (int64 i = 0; i < operand_shape.dimensions_size(); ++i) { const auto& p = padding_config.dimensions(i); + if (operand_shape.is_dynamic_dimension(i) && p.edge_padding_high() != 0 && + p.edge_padding_low() != 0 && p.interior_padding() != 0) { + return InvalidArgument( + "Dynamic dimension on padding dimension is not supported."); + } dimensions[i] = operand_shape.dimensions(i) + p.edge_padding_low() + p.edge_padding_high() + std::max(operand_shape.dimensions(i) - 1, 0LL) * p.interior_padding(); + if (dimensions[i] < 0) { + return InvalidArgument("Padding result in negative size for dimension %d", + i); + } + is_dynamic[i] = operand_shape.is_dynamic_dimension(i); } + return ShapeUtil::MakeShape( ShapeUtil::HigherPrecisionElementType(operand_shape, padding_value_shape), - dimensions); + dimensions, is_dynamic); } // Current DotDimensionNumbers Requirements: @@ -534,9 +567,8 @@ Status ValidateDotDimensionNumbers( absl::Span contracting_dims, absl::Span batch_dims) -> bool { auto in_range = [&rank](int64 i) -> bool { return 0 <= i && i < rank; }; - return std::all_of(contracting_dims.begin(), contracting_dims.end(), - in_range) && - std::all_of(batch_dims.begin(), batch_dims.end(), in_range); + return absl::c_all_of(contracting_dims, in_range) && + absl::c_all_of(batch_dims, in_range); }; absl::Span lhs_contracting_dimensions = @@ -563,9 +595,8 @@ Status ValidateDotDimensionNumbers( auto is_unique = [&dim_set](int64 i) -> bool { return dim_set.insert(i).second; }; - return std::all_of(contracting_dims.begin(), contracting_dims.end(), - is_unique) && - std::all_of(batch_dims.begin(), batch_dims.end(), is_unique); + return absl::c_all_of(contracting_dims, is_unique) && + absl::c_all_of(batch_dims, is_unique); }; if (!dims_unique(lhs_contracting_dimensions, lhs_batch_dimensions) || @@ -622,7 +653,9 @@ Status ValidateDotDimensionNumbers( const int64 rhs_contracting_dimension = dimension_numbers.rhs_contracting_dimensions(i); if (lhs.dimensions(lhs_contracting_dimension) != - rhs.dimensions(rhs_contracting_dimension)) { + rhs.dimensions(rhs_contracting_dimension) || + lhs.is_dynamic_dimension(lhs_contracting_dimension) != + rhs.is_dynamic_dimension(rhs_contracting_dimension)) { return fail("Contracting dimension sizes do not match."); } } @@ -636,7 +669,10 @@ Status ValidateDotDimensionNumbers( // Check that batch dimension numbers and sizes match. for (int64 i = 0; i < dimension_numbers.lhs_batch_dimensions_size(); ++i) { if (lhs.dimensions(dimension_numbers.lhs_batch_dimensions(i)) != - rhs.dimensions(dimension_numbers.rhs_batch_dimensions(i))) { + rhs.dimensions(dimension_numbers.rhs_batch_dimensions(i)) || + lhs.is_dynamic_dimension(dimension_numbers.lhs_batch_dimensions(i)) != + rhs.is_dynamic_dimension( + dimension_numbers.rhs_batch_dimensions(i))) { return fail("Batch dimension sizes must match for lhs/rhs."); } } @@ -647,14 +683,17 @@ Status ValidateDotDimensionNumbers( // Generate the result dimensions in order, rhs dimensions followed by lhs // dimensions except the contracted and batch dimensions. std::vector dimensions; + std::vector is_dynamic; for (int64 lhs_dim : dimension_numbers.lhs_batch_dimensions()) { dimensions.push_back(lhs.dimensions(lhs_dim)); + is_dynamic.push_back(lhs.is_dynamic_dimension(lhs_dim)); } for (int64 i = 0; i < lhs.rank(); i++) { if (!absl::c_linear_search(dimension_numbers.lhs_contracting_dimensions(), i) && !absl::c_linear_search(dimension_numbers.lhs_batch_dimensions(), i)) { dimensions.push_back(lhs.dimensions(i)); + is_dynamic.push_back(lhs.is_dynamic_dimension(i)); } } for (int64 i = 0; i < rhs.rank(); i++) { @@ -662,10 +701,11 @@ Status ValidateDotDimensionNumbers( i) && !absl::c_linear_search(dimension_numbers.rhs_batch_dimensions(), i)) { dimensions.push_back(rhs.dimensions(i)); + is_dynamic.push_back(rhs.is_dynamic_dimension(i)); } } Shape result = ShapeUtil::MakeShape( - ShapeUtil::HigherPrecisionElementType(lhs, rhs), dimensions); + ShapeUtil::HigherPrecisionElementType(lhs, rhs), dimensions, is_dynamic); TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(result)); VLOG(2) << "inferred dot shape: " << ShapeUtil::HumanString(result); @@ -683,13 +723,17 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, // dimension). In that case, the output shape has the non-1 dimension size // from the lhs/rhs pair in every index. std::vector output_dimensions(lhs.rank()); + std::vector output_dimensions_is_dynamic(lhs.rank()); for (int64 i = 0; i < lhs.rank(); ++i) { if (lhs.dimensions(i) == rhs.dimensions(i)) { output_dimensions[i] = lhs.dimensions(i); + output_dimensions_is_dynamic[i] = lhs.is_dynamic_dimension(i); } else if (lhs.dimensions(i) == 1) { output_dimensions[i] = rhs.dimensions(i); + output_dimensions_is_dynamic[i] = rhs.is_dynamic_dimension(i); } else if (rhs.dimensions(i) == 1) { output_dimensions[i] = lhs.dimensions(i); + output_dimensions_is_dynamic[i] = lhs.is_dynamic_dimension(i); } else { return InvalidArgument( "Binary op %s with incompatible shapes: %s and %s.", @@ -698,7 +742,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, } } return ShapeUtil::MakeShape(ShapeUtil::HigherPrecisionElementType(lhs, rhs), - output_dimensions); + output_dimensions, output_dimensions_is_dynamic); } /* static */ StatusOr ShapeInference::InferInDimBroadcastShape( @@ -777,6 +821,9 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, } int64 small_dimension_size = smaller_shape.dimensions(i); int64 large_dimension_size = larger_shape.dimensions(dimension_to_match); + bool small_is_dynamic = smaller_shape.is_dynamic_dimension(i); + bool large_is_dynamic = + larger_shape.is_dynamic_dimension(dimension_to_match); // Dimension sizes must be compatible: match or be degenerate (degenerate // case is handled by degenerate dimension broadcasting which occurs after // InDim broadcasting). @@ -788,6 +835,17 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, ShapeUtil::HumanString(smaller_shape), ShapeUtil::HumanString(larger_shape)); } + if (small_is_dynamic != large_is_dynamic) { + if ((small_dimension_size == 1 && !small_is_dynamic) || + (large_dimension_size == 1 && !large_is_dynamic)) { + // Do nothing. It's OK when the size-1 dimension is not static. + } else { + return InvalidArgument( + "Broadcast dimension %d dynamism mismatch: %s and %s.", i, + ShapeUtil::HumanString(smaller_shape), + ShapeUtil::HumanString(larger_shape)); + } + } // Make sure the broadcast dimensions are listed in a strictly increasing // order. if (i > 0 && broadcast_dimensions.at(i - 1) >= dimension_to_match) { @@ -797,6 +855,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, } output_shape.set_dimensions(dimension_to_match, small_dimension_size); + output_shape.set_dynamic_dimension(dimension_to_match, small_is_dynamic); } return output_shape; @@ -908,6 +967,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, broadcast_dimensions)); if (lhs.element_type() == F32 && rhs.element_type() == F32) { return ShapeUtil::ChangeElementType(shape, C64); + } else if (lhs.element_type() == F64 && rhs.element_type() == F64) { + return ShapeUtil::ChangeElementType(shape, C128); } else { return Unimplemented("Complex component type is not implemented."); } @@ -1539,6 +1600,13 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, batch_group_count); } + if (batch_group_count > 1 && feature_group_count > 1) { + return InvalidArgument( + "both batch_group_count %d and feature_group_count %d cannot be " + "greater than 1", + batch_group_count, feature_group_count); + } + if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(lhs, rhs)) { return InvalidArgument( "Convolution with different element types: %s and %s.", @@ -1589,29 +1657,29 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, input_dnums[1] = dnums.input_feature_dimension(); std::copy(dnums.input_spatial_dimensions().begin(), dnums.input_spatial_dimensions().end(), input_dnums.begin() + 2); - std::sort(input_dnums.begin(), input_dnums.end()); + absl::c_sort(input_dnums); std::vector window_dnums(num_dims); window_dnums[0] = dnums.kernel_input_feature_dimension(); window_dnums[1] = dnums.kernel_output_feature_dimension(); std::copy(dnums.kernel_spatial_dimensions().begin(), dnums.kernel_spatial_dimensions().end(), window_dnums.begin() + 2); - std::sort(window_dnums.begin(), window_dnums.end()); + absl::c_sort(window_dnums); std::vector output_dnums(num_dims); output_dnums[0] = dnums.output_batch_dimension(); output_dnums[1] = dnums.output_feature_dimension(); std::copy(dnums.output_spatial_dimensions().begin(), dnums.output_spatial_dimensions().end(), output_dnums.begin() + 2); - std::sort(output_dnums.begin(), output_dnums.end()); + absl::c_sort(output_dnums); std::vector expected_dnums(num_dims); std::iota(expected_dnums.begin(), expected_dnums.end(), 0); const auto in_range = [num_dims](int64 i) { return 0 <= i && i < num_dims; }; - if (!std::all_of(input_dnums.begin(), input_dnums.end(), in_range) || - !std::all_of(window_dnums.begin(), window_dnums.end(), in_range) || - !std::all_of(output_dnums.begin(), output_dnums.end(), in_range)) { + if (!absl::c_all_of(input_dnums, in_range) || + !absl::c_all_of(window_dnums, in_range) || + !absl::c_all_of(output_dnums, in_range)) { return InvalidArgument( "A dimension number is out of range in convolution: %s.", dnums.DebugString()); @@ -1652,6 +1720,17 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, const int64 kernel_output_features = rhs.dimensions(dnums.kernel_output_feature_dimension()); + if (batch_group_count > 1 && input_batch % kernel_output_features != 0) { + return InvalidArgument( + "Expected output feature dimension (value %d) to be divisible by " + "input_batch (value %d) for batch group count %d; " + "got (%s, %s)\n" + "Dimension numbers: {%s}.", + kernel_output_features, input_batch, batch_group_count, + ShapeUtil::HumanString(lhs), ShapeUtil::HumanString(rhs), + dnums.DebugString()); + } + if (input_features % feature_group_count != 0 || input_features / feature_group_count != kernel_input_features) { return InvalidArgument( @@ -1713,8 +1792,33 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, dimensions[dnums.output_spatial_dimensions(i)] = window_output_shape.dimensions(i); } + std::vector is_dynamic(num_dims); + for (int i = 0; i < num_dims; i++) { + if (lhs.is_dynamic_dimension(i)) { + if (i == dnums.input_batch_dimension()) { + is_dynamic[dnums.output_batch_dimension()] = true; + } else if (i == dnums.input_feature_dimension()) { + // Input feature dimension is a contracting dimension, which does not + // affect the output dimension size. So we need to do nothing. + } else { + return InvalidArgument( + "Dynamic Spatial Convolution is not supported: lhs shape is %s ", + lhs.ToString()); + } + } + if (rhs.is_dynamic_dimension(i)) { + if (i == dnums.kernel_input_feature_dimension()) { + // Kernel feature dimension does not affect the output dimension size. + // So we need to do nothing. + } else { + return InvalidArgument( + "Dynamic Spatial Convolution is not supported: rhs shape is %s ", + rhs.ToString()); + } + } + } return ShapeUtil::MakeShape(ShapeUtil::HigherPrecisionElementType(lhs, rhs), - dimensions); + dimensions, is_dynamic); } /* static */ StatusOr ShapeInference::InferFftShape( @@ -1735,7 +1839,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, case FFT: case IFFT: if (in.element_type() != C64) { - return InvalidArgument("%s requires C64 input type, found %s.", + return InvalidArgument("%s requires complex input type, found %s.", FftType_Name(fft_type), PrimitiveType_Name(in.element_type())); } @@ -1758,6 +1862,9 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, fft_length[i]); } } + if (ShapeUtil::IsZeroElementArray(in)) { + return in; + } Shape result = ShapeUtil::ChangeElementType(in, C64); result.set_dimensions(result.dimensions_size() - 1, fft_length[fft_rank - 1] / 2 + 1); @@ -1799,6 +1906,49 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, #undef RET_CHECK_RANK } +/* static */ StatusOr ShapeInference::InferTriangularSolveShape( + const Shape& a, const Shape& b, const TriangularSolveOptions& options) { + if (a.rank() < 2) { + return InvalidArgument( + "The 'a' argument to TriangularSolve must have rank >= 2, got shape %s", + a.ToString()); + } + if (b.rank() != a.rank()) { + return InvalidArgument( + "Arguments to triangular solve must have equal rank; got %s and %s.", + b.ToString(), 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()); + } + if (a.dimensions(a.rank() - 1) != + b.dimensions(b.rank() - (options.left_side() ? 2 : 1))) { + return InvalidArgument( + "The shared dimension of 'a' and 'b' does not match, got shapes %s and " + "%s", + a.ToString(), b.ToString()); + } + absl::Span a_batch_dims(a.dimensions()); + absl::Span b_batch_dims(b.dimensions()); + a_batch_dims.remove_suffix(2); + b_batch_dims.remove_suffix(2); + if (a_batch_dims != b_batch_dims) { + return InvalidArgument( + "The leading batch dimensions of the arguments to triangular solve " + "must be equal; got %s and %s.", + b.ToString(), a.ToString()); + } + if (!TriangularSolveOptions_Transpose_IsValid(options.transpose_a()) || + options.transpose_a() == TriangularSolveOptions::TRANSPOSE_INVALID) { + return InvalidArgument( + "Invalid transpose option value for triangular solve (%d).\n", + options.transpose_a()); + } + return b; +} + /* static */ StatusOr ShapeInference::InferAllReduceShape( absl::Span operand_shapes) { for (const Shape* operand_shape : operand_shapes) { @@ -1886,7 +2036,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, for (int64 i = 1; i < num_reduced_args; ++i) { if (!ShapeUtil::SameDimensions(*reduced_args[0], *reduced_args[i])) { return InvalidArgument( - "All reduced tensors must have the sime dimension. Tensor 0 has " + "All reduced tensors must have the same dimension. Tensor 0 has " "shape %s, Tensor %d has shape %s", ShapeUtil::HumanString(*reduced_args[0]), i, ShapeUtil::HumanString(*reduced_args[i])); @@ -1915,20 +2065,22 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, std::set dimensions_to_reduce_set(dimensions_to_reduce.begin(), dimensions_to_reduce.end()); std::vector new_dimensions; + std::vector new_is_dynamic; for (int i = 0; i < arg.rank(); ++i) { if (dimensions_to_reduce_set.find(i) == dimensions_to_reduce_set.end()) { new_dimensions.push_back(arg.dimensions(i)); + new_is_dynamic.push_back(arg.is_dynamic_dimension(i)); } } if (ShapeUtil::IsScalar(to_apply.result())) { return ShapeUtil::MakeShape(to_apply.result().element_type(), - new_dimensions); + new_dimensions, new_is_dynamic); } else { std::vector result_subshapes; for (const Shape& subshape : to_apply.result().tuple_shapes()) { - result_subshapes.push_back( - ShapeUtil::MakeShape(subshape.element_type(), new_dimensions)); + result_subshapes.push_back(ShapeUtil::MakeShape( + subshape.element_type(), new_dimensions, new_is_dynamic)); } return ShapeUtil::MakeTupleShape(result_subshapes); } @@ -2002,6 +2154,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, ShapeUtil::HumanString(source_shape), ShapeUtil::HumanString(window_result_shape)); } + return operand_shape; } @@ -2101,15 +2254,15 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, } const Shape& start_indices_shape = start_index_shapes[0]; - TF_RETURN_IF_ERROR( - ExpectArray(start_indices_shape, "start indices of dynamic slice")); - VLOG(2) << StrFormat( "slicing shape %s at dynamic start_indices %s with slice_sizes={%s}", ShapeUtil::HumanString(operand_shape), ShapeUtil::HumanString(start_indices_shape), StrJoin(slice_sizes, ", ")); + TF_RETURN_IF_ERROR( + ExpectArray(start_indices_shape, "start indices of dynamic slice")); + if (start_indices_shape.rank() != 1) { return InvalidArgument( "Dynamic slice start indices of rank %d must be rank1.", @@ -2153,7 +2306,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, "Dynamic slice start indices must be of integral type."); } for (const Shape& index_shape : start_index_shapes) { - if (!ShapeUtil::Equal(first_index_shape, index_shape)) { + if (!ShapeUtil::Compatible(first_index_shape, index_shape)) { return InvalidArgument( "Dynamic slice start indices must all have the same shape, got " "mismatching indices with shapes %s and %s.", @@ -2242,8 +2395,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, if (operand_shape.rank() != number_of_indices) { return InvalidArgument( - "Dynamic update slice start number of dimensions %d must match rank " - "%d of slice input (%s).", + "Dynamic update slice start number of dimensions %d must match " + "rank %d of slice input (%s).", number_of_indices, operand_shape.rank(), ShapeUtil::HumanString(operand_shape)); } @@ -2260,7 +2413,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, "Dynamic update slice start indices must be of integral type."); } for (const Shape& index_shape : start_index_shapes) { - if (!ShapeUtil::Equal(first_index_shape, index_shape)) { + if (!ShapeUtil::Compatible(first_index_shape, index_shape)) { return InvalidArgument( "Dynamic update slice start indices must all have the same " "shape, got mismatching indices with shapes %s and %s.", @@ -2330,7 +2483,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, ShapeUtil::HumanString(arg)); } - if (index >= arg.tuple_shapes_size()) { + if (index < 0 || index >= arg.tuple_shapes_size()) { return InvalidArgument( "Cannot infer shape: attempt to index out of tuple bounds: %d " ">= %d in shape %s.", @@ -2360,7 +2513,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, }; // Check the shapes of computation parameters and return types. - if (!ShapeUtil::ShapeIs(condition.result(), PRED, {})) { + if (!ShapeUtil::Equal(condition.result(), ShapeUtil::MakeShape(PRED, {}))) { return InvalidArgument("Condition must return a boolean; got %s.", shape_string()); } @@ -2380,7 +2533,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, const Shape& predicate, const Shape& true_operand, const Shape& false_operand, const ProgramShape& true_computation, const ProgramShape& false_computation) { - if (!ShapeUtil::ShapeIs(predicate, PRED, {})) { + if (!ShapeUtil::Equal(predicate, ShapeUtil::MakeShape(PRED, {}))) { return InvalidArgument("Predicate must be a boolean; got %s.", ShapeUtil::HumanString(predicate)); } @@ -2479,11 +2632,17 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, operand_shape.dimensions(i) != 1) { return InvalidArgument( "Input dimension should be either 1 or equal to the output dimension " - "it's broadcasting into; the %lldth operand dimension is %lld, the " + "it is broadcasting into; the %lldth operand dimension is %lld, the " "%lldth output dimension is %lld.", i, operand_shape.dimensions(i), broadcast_dimensions[i], output_shape.dimensions(broadcast_dimensions[i])); } + if (operand_shape.is_dynamic_dimension(i) != + output_shape.is_dynamic_dimension(broadcast_dimensions[i])) { + return InvalidArgument( + "Broadcast input and output dynamism mismatch: %s and %s", + operand_shape.ToString(), output_shape.ToString()); + } // Make sure the broadcast dimensions are listed in a strictly increasing // order. if (i > 0 && broadcast_dimensions[i - 1] >= broadcast_dimensions[i]) { @@ -2526,6 +2685,14 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, StrJoin(dimensions, ","), ShapeUtil::HumanString(operand)); } + std::vector> unmodified_dims = + ShapeUtil::DimensionsUnmodifiedByReshape(operand, inferred_shape); + for (auto& unmodified : unmodified_dims) { + if (operand.is_dynamic_dimension(unmodified.first)) { + inferred_shape.set_dynamic_dimension(unmodified.second, true); + } + } + return inferred_shape; } @@ -2533,11 +2700,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, const Shape& operand, absl::Span dimensions) { TF_RETURN_IF_ERROR(ExpectArray(operand, "transpose")); - std::vector indices(operand.rank()); - std::iota(indices.begin(), indices.end(), 0); - if (dimensions.size() != operand.rank() || - !std::is_permutation(dimensions.begin(), dimensions.end(), - indices.begin())) { + if (!IsPermutation(dimensions, operand.rank())) { return InvalidArgument( "Transpose dimensions [%s] are not a permutation of the operand " "dimensions (operand shape is %s).", @@ -2599,19 +2762,31 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, "Select's pred operand must have PRED element type; got %s.", ShapeUtil::HumanString(pred)); } - if (ShapeUtil::CompatibleIgnoringElementType(pred, on_true) || + if (Shape::Equal() + .IgnoreElementType() + .IgnoreLayout() + .IgnoreDynamicDimension()(pred, on_true) || ShapeUtil::IsScalar(pred)) { // By this stage we know that pred's element type is PRED. Therefore, this // check restricts pred to be a PRED scalar, or a PRED array with the same // dimensions as on_true and on_false. - return ShapeUtil::ChangeElementType( + Shape inferred_shape = ShapeUtil::ChangeElementType( on_true, ShapeUtil::HigherPrecisionElementType(on_true, on_false)); - } else { - return InvalidArgument( - "Select operation with non-scalar predicate with dimensionality " - " different from the other operands: %s.", - ShapeUtil::HumanString(pred)); + + // Propagate dynamic dimensions if pred is not a scalar. + if (!ShapeUtil::IsScalar(pred)) { + for (int i = 0; i < inferred_shape.rank(); i++) { + if (pred.is_dynamic_dimension(i)) { + inferred_shape.set_dynamic_dimension(i, true); + } + } + } + return inferred_shape; } + return InvalidArgument( + "Select operation with non-scalar predicate with dimensionality " + "different from the other operands: %s.", + ShapeUtil::HumanString(pred)); } /* static */ StatusOr ShapeInference::InferTupleSelectShape( diff --git a/tensorflow/compiler/xla/service/shape_inference.h b/tensorflow/compiler/xla/service/shape_inference.h index e440e364c8f4edd32ea0eef612c43c5a1b3568a7..acb071ab18824472153fc608b812ad2d9c52651e 100644 --- a/tensorflow/compiler/xla/service/shape_inference.h +++ b/tensorflow/compiler/xla/service/shape_inference.h @@ -116,6 +116,10 @@ class ShapeInference { static StatusOr InferFftShape(const Shape& in, FftType fft_type, absl::Span fft_length); + // Infers the shape produced by the given triangular solve operation. + static StatusOr InferTriangularSolveShape( + const Shape& a, const Shape& b, const TriangularSolveOptions& options); + // Infers the shape produced by a cross replica sum with the given operand // shapes. static StatusOr InferAllReduceShape( @@ -177,14 +181,14 @@ class ShapeInference { // in 'slice_sizes', with dynamic start indices shape 'start_indices_shape'. static StatusOr InferDynamicSliceShape( const Shape& operand_shape, absl::Span start_index_shapes, - absl::Span slice_sizes, bool allow_scalar_indices = false); + absl::Span slice_sizes, bool allow_scalar_indices = true); // Infers the shape produced by a dynamic update slice operation based // on the shape of operand and update. static StatusOr InferDynamicUpdateSliceShape( const Shape& operand_shape, const Shape& update_shape, absl::Span start_index_shapes, - bool allow_scalar_indices = false); + bool allow_scalar_indices = true); // Infers the shape produced by doing a compile-time-constant indexing into // the given input shape. This is essential for operations on tuples, because diff --git a/tensorflow/compiler/xla/service/shape_inference_test.cc b/tensorflow/compiler/xla/service/shape_inference_test.cc index bb00bfd2eee39e0e87510245bde7f43f832120f4..f400ef51f07b006eef2ea674feff1dd72f836e77 100644 --- a/tensorflow/compiler/xla/service/shape_inference_test.cc +++ b/tensorflow/compiler/xla/service/shape_inference_test.cc @@ -35,6 +35,7 @@ class ShapeInferenceTest : public ::testing::Test { protected: // Some handy scalar shapes. const Shape s32_ = ShapeUtil::MakeShape(S32, {}); + const Shape f16_ = ShapeUtil::MakeShape(F16, {}); const Shape f32_ = ShapeUtil::MakeShape(F32, {}); const Shape f64_ = ShapeUtil::MakeShape(F64, {}); const Shape pred_ = ShapeUtil::MakeShape(PRED, {}); @@ -251,7 +252,7 @@ TEST_F(ShapeInferenceTest, ClampBadShapes) { TEST_F(ShapeInferenceTest, Complex) { auto complex_shape = [&](const Shape& lhs, const Shape& rhs, - const absl::Span& bcast) { + absl::Span bcast) { return ShapeInference::InferBinaryOpShape(HloOpcode::kComplex, lhs, rhs, bcast); }; @@ -260,8 +261,8 @@ TEST_F(ShapeInferenceTest, Complex) { ASSERT_FALSE(complex_shape(pred_, pred_, {}).ok()); // Component types must match. ASSERT_FALSE(complex_shape(f32_, f64_, {}).ok()); - // Only F32->C64 supported. - ASSERT_FALSE(complex_shape(f64_, f64_, {}).ok()); + // Only F32->C64 and F64->C128 supported. + ASSERT_FALSE(complex_shape(f16_, f16_, {}).ok()); // Validate correct uses. Shape c64_32 = ShapeUtil::MakeShape(C64, {32}); TF_ASSERT_OK_AND_ASSIGN(Shape result, complex_shape(f32_, f32_, {})); @@ -285,6 +286,9 @@ TEST_F(ShapeInferenceTest, Complex) { ASSERT_TRUE(ShapeUtil::Equal(result, c64_32_64)); TF_ASSERT_OK_AND_ASSIGN(result, complex_shape(matrix_32_64_, f32_, {})); ASSERT_TRUE(ShapeUtil::Equal(result, c64_32_64)); + + TF_ASSERT_OK_AND_ASSIGN(result, complex_shape(f64_, f64_, {})); + ASSERT_TRUE(ShapeUtil::Equal(result, ShapeUtil::MakeShape(C128, {}))); } TEST_F(ShapeInferenceTest, VariadicOpTuplify) { @@ -892,6 +896,20 @@ TEST_F(ShapeInferenceTest, InferConstIndexShape) { ASSERT_TRUE(ShapeUtil::Equal(s32_, inferred1_status.ValueOrDie())); } +TEST_F(ShapeInferenceTest, InferTupleElementShapeOutOfBound) { + Shape tuple_shape = ShapeUtil::MakeTupleShape({f32_, s32_}); + auto inferredNegative_status = + ShapeInference::InferGetTupleElementShape(tuple_shape, -1); + auto inferred2_status = + ShapeInference::InferGetTupleElementShape(tuple_shape, 2); + ASSERT_FALSE(inferredNegative_status.ok()); + ASSERT_FALSE(inferred2_status.ok()); + EXPECT_THAT(inferredNegative_status.status().error_message(), + HasSubstr("attempt to index out of tuple bounds")); + EXPECT_THAT(inferred2_status.status().error_message(), + HasSubstr("attempt to index out of tuple bounds")); +} + TEST_F(ShapeInferenceTest, InferPowShape) { auto ten_floats = ShapeUtil::MakeShape(F32, {10}); auto inferred_status = ShapeInference::InferBinaryOpShape( @@ -1463,6 +1481,14 @@ TEST_F(ShapeInferenceTest, Pad) { Shape inferred_shape = inferred_status.ValueOrDie(); ASSERT_TRUE( ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {39, 31}), inferred_shape)); + + dimension1->set_edge_padding_low(-20); + dimension1->set_edge_padding_high(-10); + auto negative_dimension_size = ShapeInference::InferPadShape( + input_shape, padding_value_shape, padding_config); + ASSERT_FALSE(negative_dimension_size.ok()); + ASSERT_THAT(negative_dimension_size.status().error_message(), + HasSubstr("negative size for dimension 1")); } TEST_F(ShapeInferenceTest, Reverse) { @@ -1546,6 +1572,16 @@ TEST_F(ShapeInferenceTest, Transpose) { ShapeUtil::MakeShape(F32, {3, 4, 5, 2}))); } +TEST_F(ShapeInferenceTest, Rank1Transpose) { + Shape a_shape = ShapeUtil::MakeShape(F32, {5}); + auto inferred_shape_and_status = + ShapeInference::InferTransposeShape(a_shape, {0}); + EXPECT_IS_OK(inferred_shape_and_status); + Shape inferred_shape = inferred_shape_and_status.ValueOrDie(); + EXPECT_TRUE( + ShapeUtil::Compatible(inferred_shape, ShapeUtil::MakeShape(F32, {5}))); +} + TEST_F(ShapeInferenceTest, Conditional) { auto inferred_status0 = ShapeInference::InferConditionalShape( pred_, vector_32_, vector_64_, diff --git a/tensorflow/compiler/xla/service/sort_simplifier.cc b/tensorflow/compiler/xla/service/sort_simplifier.cc new file mode 100644 index 0000000000000000000000000000000000000000..122366a0f322a66963b364e1b19629cbd2d9aabe --- /dev/null +++ b/tensorflow/compiler/xla/service/sort_simplifier.cc @@ -0,0 +1,165 @@ +/* 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/sort_simplifier.h" + +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/statusor.h" + +namespace xla { +namespace { + +// If the sort instruction has a tuple shape then looks for unused output +// values and removes them from the sort instruction. Returns true if the +// graph has been modified. +StatusOr RemoveUnusedOperandFromSort(HloInstruction* sort) { + if (!sort->shape().IsTuple()) { + return false; + } + + HloComputation* computation = sort->parent(); + + if (computation->root_instruction() == sort) { + // Can't analyse users of the root instruction. + return false; + } + + absl::flat_hash_set used_indices; + for (const HloInstruction* user : sort->users()) { + if (user->opcode() != HloOpcode::kGetTupleElement) { + // Can't analyse users other then get-tuple-element. + return false; + } + used_indices.insert(user->tuple_index()); + } + + // Also note which parameters are used by the comparator computation. + auto comparator = sort->to_apply(); + for (int64 i = 0; i < sort->operand_count() * 2; ++i) { + if (comparator->parameter_instruction(i)->user_count() > 0) { + // operand i corresponds to parameters 2 * i and 2 * i + 1 of the + // computation. + used_indices.insert(i / 2); + } + } + + if (used_indices.size() == sort->operand_count()) { + // All operands are used. + return false; + } + + std::vector operands; + std::vector new_shapes; + for (int64 i = 0; i < sort->operand_count(); ++i) { + if (used_indices.contains(i)) { + operands.push_back(sort->mutable_operand(i)); + new_shapes.push_back(sort->operand(i)->shape()); + } + } + + Shape new_sort_shape = new_shapes.size() == 1 + ? new_shapes[0] + : ShapeUtil::MakeTupleShape(new_shapes); + HloInstruction* new_sort = computation->AddInstruction( + sort->CloneWithNewOperands(new_sort_shape, operands)); + absl::flat_hash_map> + replacements; + int64 parameter_number = 0; + for (int64 i = 0; i < sort->operand_count(); ++i) { + auto* old_lhs_parameter = comparator->parameter_instruction(i * 2); + auto* old_rhs_parameter = comparator->parameter_instruction(i * 2 + 1); + if (used_indices.contains(i)) { + Shape scalar_shape = + ShapeUtil::MakeShape(sort->operand(i)->shape().element_type(), {}); + replacements[old_lhs_parameter] = HloInstruction::CreateParameter( + parameter_number, scalar_shape, + absl::StrCat("p.", parameter_number / 2, ".lhs")); + ++parameter_number; + replacements[old_rhs_parameter] = HloInstruction::CreateParameter( + parameter_number, scalar_shape, + absl::StrCat("p.", parameter_number / 2, ".rhs")); + ++parameter_number; + } else { + replacements[old_lhs_parameter] = nullptr; + replacements[old_rhs_parameter] = nullptr; + } + } + HloModule* module = sort->GetModule(); + HloComputation* new_compare = module->AddEmbeddedComputation( + comparator->CloneWithReplacements(std::move(replacements))); + new_sort->set_to_apply(new_compare); + + // Map from original get-tuple-element tuple index to new HLO instruction + absl::flat_hash_map result_map; + if (new_sort->shape().IsTuple()) { + // Old sort key maps to new sort key. + int64 new_index = 0; + for (int64 i = 0; i < sort->operand_count(); ++i) { + if (used_indices.count(i)) { + result_map[i] = + computation->AddInstruction(HloInstruction::CreateGetTupleElement( + new_shapes[new_index], new_sort, new_index)); + ++new_index; + } + } + } else { + CHECK_EQ(used_indices.size(), 1); + result_map[*used_indices.begin()] = new_sort; + } + std::vector users(sort->users().begin(), + sort->users().end()); + for (HloInstruction* user : users) { + TF_RETURN_IF_ERROR( + user->ReplaceAllUsesWith(result_map.at(user->tuple_index()))); + TF_RETURN_IF_ERROR(computation->RemoveInstructionAndUnusedOperands(user)); + } + return true; +} +} // namespace + +StatusOr SortSimplifier::Run(HloModule* module) { + VLOG(2) << "HLO module before SortSimplifier:"; + XLA_VLOG_LINES(2, module->ToString()); + + bool changed = false; + std::vector sort_instrs; + for (auto* comp : module->MakeNonfusionComputations()) { + absl::c_copy_if(comp->instructions(), std::back_inserter(sort_instrs), + [](const HloInstruction* instr) { + return instr->opcode() == HloOpcode::kSort; + }); + } + + for (HloInstruction* sort_instr : sort_instrs) { + TF_ASSIGN_OR_RETURN(bool result, RemoveUnusedOperandFromSort(sort_instr)); + changed |= result; + } + + if (changed) { + VLOG(2) << "HLO module after SortSimplifier:"; + XLA_VLOG_LINES(2, module->ToString()); + } else { + VLOG(2) << "HLO module unchanged after SortSimplifier"; + } + + return changed; +} +} // namespace xla diff --git a/tensorflow/compiler/xla/service/sort_simplifier.h b/tensorflow/compiler/xla/service/sort_simplifier.h new file mode 100644 index 0000000000000000000000000000000000000000..8c6f313aa04f51e14a14450bc72fc622d74133a4 --- /dev/null +++ b/tensorflow/compiler/xla/service/sort_simplifier.h @@ -0,0 +1,35 @@ +/* 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_SORT_SIMPLIFIER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_SORT_SIMPLIFIER_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 { + +// HLO pass which removes unused operands from sort, where an unused operand is +// defined as an operand at some index 'x' at which the output is not used. +class SortSimplifier : public HloModulePass { + public: + absl::string_view name() const override { return "simplify-sorts"; } + StatusOr Run(HloModule* module) override; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_SORT_SIMPLIFIER_H_ diff --git a/tensorflow/compiler/xla/service/sort_simplifier_test.cc b/tensorflow/compiler/xla/service/sort_simplifier_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..696ac1b465848894f8dcb1c88bc48c6a5b268ef4 --- /dev/null +++ b/tensorflow/compiler/xla/service/sort_simplifier_test.cc @@ -0,0 +1,160 @@ +/* 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/sort_simplifier.h" + +#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/service/hlo_parser.h" +#include "tensorflow/compiler/xla/service/pattern_matcher.h" +#include "tensorflow/compiler/xla/service/pattern_matcher_gmock.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 { + +namespace m = match; + +using SortSimplifierTest = HloTestBase; + +TEST_F(SortSimplifierTest, RemoveUnusedSortOperandArrayResult) { + const char* hlo_string = R"( + HloModule permutation_sort + + compare { + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + ROOT lt = pred[] less-than(p.0.lhs, p.0.rhs) + } + + ENTRY sort_computation { + keys = f32[64,8732]{1,0} parameter(0) + values = s32[64,8732]{1,0} parameter(1) + sort = (f32[64,8732]{1,0}, s32[64,8732]{1,0}) sort(keys, values), + dimensions={1}, to_apply=compare + ROOT gte = f32[64,8732]{1,0} get-tuple-element(sort), index=0 + })"; + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); + + SortSimplifier simplifier; + uint64 num_executions = 0; + do { + num_executions++; + } while (simplifier.Run(module.get()).ValueOrDie()); + EXPECT_EQ(num_executions, 2); + auto root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, GmockMatch(m::Sort(m::Parameter(0)))); +} + +TEST_F(SortSimplifierTest, RemoveUnusedSortOperandTuple) { + const char* hlo_string = R"( + HloModule permutation_sort + + compare { + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + p.2.lhs = u32[] parameter(4) + p.2.rhs = u32[] parameter(5) + ROOT lt = pred[] less-than(p.0.lhs, p.0.rhs) + } + + ENTRY sort_computation { + keys = f32[64,87] parameter(0) + values.0 = s32[64,87] parameter(1) + values.1 = u32[64,87] parameter(2) + sort = (f32[64,87], s32[64,87], u32[64,87]) sort( + keys, values.0, values.1), + dimensions={1}, to_apply=compare + gte.0 = f32[64,87] get-tuple-element(sort), index=0 + gte.1 = u32[64,87] get-tuple-element(sort), index=2 + ROOT tuple = (f32[64,87], u32[64,87]) tuple(gte.0, gte.1) + })"; + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); + + SortSimplifier simplifier; + EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + auto root = module->entry_computation()->root_instruction(); + EXPECT_THAT( + root, + GmockMatch(m::Tuple( + m::GetTupleElement(m::Sort(m::Parameter(0), m::Parameter(2)), 0), + m::GetTupleElement(m::Sort(m::Parameter(0), m::Parameter(2)), 1)))); +} + +TEST_F(SortSimplifierTest, DontRemoveUnusedSortKey) { + const char* hlo_string = R"( + HloModule permutation_sort + + compare { + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + ROOT lt = pred[] less-than(p.0.lhs, p.0.rhs) + } + + ENTRY sort_computation { + keys = f32[64,8732]{1,0} parameter(0) + values = s32[64,8732]{1,0} parameter(1) + sort = (f32[64,8732]{1,0}, s32[64,8732]{1,0}) sort(keys, values), dimensions={1}, to_apply=compare + ROOT gte = s32[64,8732]{1,0} get-tuple-element(sort), index=1 + })"; + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); + + SortSimplifier simplifier; + EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); +} + +TEST_F(SortSimplifierTest, RemoveUnusedFirstOperand) { + const char* hlo_string = R"( + HloModule permutation_sort + + compare { + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + ROOT lt = pred[] less-than(p.1.lhs, p.1.rhs) + } + + ENTRY sort_computation { + keys = f32[64,8732]{1,0} parameter(0) + values = s32[64,8732]{1,0} parameter(1) + sort = (f32[64,8732]{1,0}, s32[64,8732]{1,0}) sort(keys, values), + dimensions={1}, to_apply=compare + ROOT gte = s32[64,8732]{1,0} get-tuple-element(sort), index=1 + })"; + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); + + SortSimplifier simplifier; + uint64 num_executions = 0; + do { + num_executions++; + } while (simplifier.Run(module.get()).ValueOrDie()); + EXPECT_EQ(num_executions, 2); + auto root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, GmockMatch(m::Sort(m::Parameter(1)))); +} +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/transpose_folding.cc b/tensorflow/compiler/xla/service/transpose_folding.cc index 15eb46bac0ac4e03d476ee21be6834174798526c..a95ca2bf2a8fcd700eb9234cafbfce9b62f2370c 100644 --- a/tensorflow/compiler/xla/service/transpose_folding.cc +++ b/tensorflow/compiler/xla/service/transpose_folding.cc @@ -130,8 +130,7 @@ bool FoldTransposeIntoConvolution(InstructionOperandsPair pair) { HloInstruction* new_lhs; const int64 kLhsIdx = 0; - if (std::find(operand_indices.begin(), operand_indices.end(), kLhsIdx) != - operand_indices.end()) { + if (absl::c_linear_search(operand_indices, kLhsIdx)) { HloInstruction& transpose = *convolution.mutable_operand(kLhsIdx); const auto& transpose_dimensions = transpose.dimensions(); HloInstruction& transpose_operand = *transpose.mutable_operand(0); @@ -154,8 +153,7 @@ bool FoldTransposeIntoConvolution(InstructionOperandsPair pair) { HloInstruction* new_rhs; const int64 kRhsIdx = 1; - if (std::find(operand_indices.begin(), operand_indices.end(), kRhsIdx) != - operand_indices.end()) { + if (absl::c_linear_search(operand_indices, kRhsIdx)) { HloInstruction& transpose = *convolution.mutable_operand(kRhsIdx); const auto& transpose_dimensions = transpose.dimensions(); HloInstruction& transpose_operand = *transpose.mutable_operand(0); diff --git a/tensorflow/compiler/xla/client/lib/triangular_solve.cc b/tensorflow/compiler/xla/service/triangular_solve_expander.cc similarity index 80% rename from tensorflow/compiler/xla/client/lib/triangular_solve.cc rename to tensorflow/compiler/xla/service/triangular_solve_expander.cc index 6061e64656e859adcc52c250e7775a878399f4e6..b26cdc1db59b30d82b9ac58a8a2ac762220086be 100644 --- a/tensorflow/compiler/xla/client/lib/triangular_solve.cc +++ b/tensorflow/compiler/xla/service/triangular_solve_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/triangular_solve.h" +#include "tensorflow/compiler/xla/service/triangular_solve_expander.h" #include #include @@ -33,6 +33,8 @@ limitations under the License. namespace xla { +namespace { + // Get the diagonal blocks of the coefficient matrix XlaOp DiagonalBlocks(XlaOp a, int64 block_size) { XlaBuilder* builder = a.builder(); @@ -140,9 +142,7 @@ XlaOp InvertDiagonalBlocks(XlaOp diag_blocks, bool lower, bool transpose_a, // zero (which can happen if the last block was padded) otherwise it will // introduce nans which will propagate auto diags = GetMatrixDiagonal(diag_blocks); - TF_ASSIGN_OR_RETURN(Shape diags_shape, builder->GetShape(diags)); - auto one = ScalarLike(diags, 1); - auto ones = Broadcast(one, AsInt64Slice(diags_shape.dimensions())); + auto ones = FullLike(diags, 1); diags = Select(Eq(diags, Zero(builder, shape.element_type())), ones, diags); auto scaled_diag_blocks = Div(diag_blocks, diags, {0, 2}); @@ -165,10 +165,10 @@ XlaOp InvertDiagonalBlocks(XlaOp diag_blocks, bool lower, bool transpose_a, // The first or last diagonal element should be set to 1 instead of -1 // though, since we never update it auto pos_one = Reshape(One(builder, shape.element_type()), {1, 1}); - auto start_index = (lower) ? 0 : block_size - 1; - auto output_block = DynamicUpdateSlice( - neg_identity, pos_one, - /*start_indices=*/ConstantR1(builder, 2, start_index)); + auto start_index = ConstantR0(builder, (lower) ? 0 : block_size - 1); + auto output_block = + DynamicUpdateSlice(neg_identity, pos_one, + /*start_indices=*/{start_index, start_index}); // Broadcast diag([1, -1, -1, ...]) to every block XlaOp output = Broadcast(output_block, @@ -211,12 +211,10 @@ XlaOp InvertDiagonalBlocks(XlaOp diag_blocks, bool lower, bool transpose_a, auto body_out = GetTupleElement(input_tuple, 1); auto body_input = GetTupleElement(input_tuple, 2); - auto zero = ConstantR1(bodyb.get(), 1, 0); + auto zero = ConstantR0(bodyb.get(), 0); auto j = (lower) ? i : ScalarLike(i, block_size - 1) - i; - auto start_indices = - ConcatInDim(bodyb.get(), {zero, Reshape(j, {1}), zero}, 0); auto input_row = - DynamicSlice(body_input, start_indices, + DynamicSlice(body_input, {zero, j, zero}, /*slice_sizes=*/{num_blocks, 1, block_size}); // We want -L21 L11^{-1} @@ -230,7 +228,7 @@ XlaOp InvertDiagonalBlocks(XlaOp diag_blocks, bool lower, bool transpose_a, precision_proto.add_operand_precision(precision); auto update = -DotGeneral(input_row, body_out, dnums, &precision_proto); - body_out = DynamicUpdateSlice(body_out, update, start_indices); + body_out = DynamicUpdateSlice(body_out, update, {zero, j, zero}); auto next_i = i + ScalarLike(i, 1); Tuple(bodyb.get(), {next_i, body_out, body_input}); @@ -349,9 +347,10 @@ XlaOp SolveWithInvertedDiagonalBlocks(XlaOp a, XlaOp b, XlaOp inv_diag_blocks, }); } -XlaOp TriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower, - bool transpose_a, bool conjugate_a, int64 block_size, - PrecisionConfig::Precision precision) { +XlaOp BuildTriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower, + bool transpose_a, bool conjugate_a, + bool unit_diagonal, int64 block_size, + PrecisionConfig::Precision precision) { XlaBuilder* builder = a.builder(); return builder->ReportErrorOrReturn([&]() -> StatusOr { TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); @@ -410,6 +409,20 @@ XlaOp TriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower, return b; } + // TODO(phawkins): consider pushing triangle masking into + // InvertDiagonalBlocks. + if (unit_diagonal) { + // Mask everything but the subdiagonal/superdiagonal elements. + a = lower ? Select(TriangleMask(a, -1), a, ZerosLike(a)) + : Select(TriangleMask(a, 0), ZerosLike(a), a); + int64 k = ShapeUtil::GetDimension(a_shape, -1); + a = xla::Add(a, IdentityMatrix(builder, a_shape.element_type(), k, k), + /*broadcast_dimensions=*/{ndims - 2, ndims - 1}); + } else { + // Mask off the ignored elements of the triangular matrix a. + a = Triangle(a, lower); + } + // We find the diagonal blocks of the coefficient matrix auto diag_blocks = DiagonalBlocks(a, block_size); @@ -417,11 +430,6 @@ XlaOp TriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower, auto inv_diag_blocks = InvertDiagonalBlocks(diag_blocks, lower, transpose_a, conjugate_a, precision); - // Mask off the ignored elements of the triangular matrix a. - // TODO(phawkins): it would probably be preferable to perform this masking - // block by block inside SolveWithInvertedDiagonalBlocks. - a = Triangle(a, lower); - // We now find the solution using GEMMs auto x = SolveWithInvertedDiagonalBlocks(a, b, inv_diag_blocks, left_side, lower, @@ -431,4 +439,66 @@ XlaOp TriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower, }); } +} // namespace + +bool TriangularSolveExpander::InstructionMatchesPattern( + HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kTriangularSolve; +} + +StatusOr TriangularSolveExpander::ExpandInstruction( + HloInstruction* instruction) { + const TriangularSolveOptions& options = + instruction->triangular_solve_options(); + const string name = absl::StrFormat( + "xla.triangular_solve_%s_%s_%s_%s_%s_%s", + instruction->operand(0)->shape().ToString(), + instruction->operand(1)->shape().ToString(), + options.left_side() ? "left" : "right", + options.lower() ? "lower" : "upper", + TriangularSolveOptions_Transpose_Name(options.transpose_a()), + options.unit_diagonal() ? "unit" : "nonunit"); + + HloModule* module = instruction->parent()->parent(); + + HloComputation*& computation = + computation_cache_.emplace(name, nullptr).first->second; + if (!computation) { + // Builds a new expansion. + // + // 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 b = Parameter(&builder, 1, instruction->operand(1)->shape(), "b"); + bool transpose_a = + options.transpose_a() != TriangularSolveOptions::NO_TRANSPOSE; + bool conjugate_a = options.transpose_a() == TriangularSolveOptions::ADJOINT; + + BuildTriangularSolve(a, b, options.left_side(), options.lower(), + transpose_a, conjugate_a, options.unit_diagonal(), + /*block_size=*/128, + /*precision=*/PrecisionConfig::HIGHEST); + 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/triangular_solve_expander.h b/tensorflow/compiler/xla/service/triangular_solve_expander.h new file mode 100644 index 0000000000000000000000000000000000000000..be2374ef8c86254d8db5ac1acac385aa0de7d3a5 --- /dev/null +++ b/tensorflow/compiler/xla/service/triangular_solve_expander.h @@ -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. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_TRIANGULAR_SOLVE_EXPANDER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_TRIANGULAR_SOLVE_EXPANDER_H_ + +#include "absl/container/flat_hash_map.h" +#include "tensorflow/compiler/xla/service/op_expander_pass.h" + +namespace xla { + +class TriangularSolveExpander : public OpExpanderPass { + public: + absl::string_view name() const override { + return "triangular_solve_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_TRIANGULAR_SOLVE_EXPANDER_H_ diff --git a/tensorflow/compiler/xla/service/tuple_points_to_analysis.cc b/tensorflow/compiler/xla/service/tuple_points_to_analysis.cc index b1f0672c6089035153aa48c853bface116f70026..cc82e9bb0287b5a586fb21fee35d3124a6d6f121 100644 --- a/tensorflow/compiler/xla/service/tuple_points_to_analysis.cc +++ b/tensorflow/compiler/xla/service/tuple_points_to_analysis.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_set.h" #include "absl/memory/memory.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" @@ -55,11 +56,10 @@ bool PointsToSet::IsAmbiguous() const { bool PointsToSet::IsDistinct() const { bool distinct = true; - std::set all_points_to; - ForEachElement([&distinct, &all_points_to](const ShapeIndex& /*index*/, - const BufferList& points_to) { + absl::flat_hash_set all_points_to; + ForEachElement([&](const ShapeIndex& /*index*/, const BufferList& points_to) { for (auto& buffer : points_to) { - if (all_points_to.count(buffer) != 0) { + if (all_points_to.contains(buffer)) { distinct = false; } all_points_to.insert(buffer); @@ -87,9 +87,7 @@ bool PointsToSet::ContainsBuffer(const LogicalBuffer& buffer) const { bool found = false; ForEachElement([&found, &buffer](const ShapeIndex& /*index*/, const BufferList& pointed_to_buffers) { - if (!found && - std::find(pointed_to_buffers.begin(), pointed_to_buffers.end(), - &buffer) != pointed_to_buffers.end()) { + if (!found && absl::c_linear_search(pointed_to_buffers, &buffer)) { found = true; } }); @@ -99,8 +97,7 @@ bool PointsToSet::ContainsBuffer(const LogicalBuffer& buffer) const { bool PointsToSet::ContainsBufferAtIndex(const LogicalBuffer& buffer, const ShapeIndex& index) const { const auto& pointed_to_buffers = element(index); - return std::find(pointed_to_buffers.begin(), pointed_to_buffers.end(), - &buffer) != pointed_to_buffers.end(); + return absl::c_linear_search(pointed_to_buffers, &buffer); } void PointsToSet::AddPointedToBuffer(const LogicalBuffer& buffer, @@ -604,9 +601,8 @@ bool TuplePointsToAnalysis::DoesNotUseOperandBuffer( } else if (user->opcode() == HloOpcode::kFusion && user->fusion_kind() == HloInstruction::FusionKind::kLoop) { // Find fusion parameter associated with 'operand'. - auto it = std::find_if( - user->fused_parameters().begin(), user->fused_parameters().end(), - [=](HloInstruction* fused_param) { + auto it = absl::c_find_if( + user->fused_parameters(), [&](HloInstruction* fused_param) { return user->operand(fused_param->parameter_number()) == operand; }); CHECK(it != user->fused_parameters().end()); @@ -672,9 +668,8 @@ bool TuplePointsToAnalysis::HasUniqueFusedUseOfOperandAt( } // Find fusion parameter associated with 'operand'. const auto& fused_params = fusion->fused_parameters(); - auto fused_param_it = std::find_if( - fused_params.begin(), fused_params.end(), - [&](HloInstruction* fused_param) { + auto fused_param_it = + absl::c_find_if(fused_params, [&](HloInstruction* fused_param) { return fusion->operand(fused_param->parameter_number()) == operand; }); if (fused_param_it == fused_params.end()) { @@ -704,6 +699,8 @@ bool TuplePointsToAnalysis::HasUniqueFusedUseOfOperandAt( // (4) The 'user' of 'operand' is DynamicUpdateSlice or While at operand index // 0. // (5) The 'user' of 'operand' is Sort, and it is the only user. +// (6) The 'user' of 'operand' is TriangularSolve, it is the second operand, +// and it is the only user. // // (2) and (3) can only be determined if points-to analysis is available. bool TuplePointsToAnalysis::CanShareOperandBufferWithUser( @@ -743,11 +740,10 @@ bool TuplePointsToAnalysis::CanShareOperandBufferWithUser( // Check if one operand of kAdd fused root is kDot or kConvolution. auto* add = user->fused_expression_root(); auto add_operand_it = - std::find_if(add->operands().begin(), add->operands().end(), - [&](HloInstruction* operand) { - return operand->opcode() == HloOpcode::kConvolution || - operand->opcode() == HloOpcode::kDot; - }); + absl::c_find_if(add->operands(), [&](HloInstruction* operand) { + return operand->opcode() == HloOpcode::kConvolution || + operand->opcode() == HloOpcode::kDot; + }); if (add_operand_it == add->operands().end()) { return false; } @@ -785,6 +781,14 @@ bool TuplePointsToAnalysis::CanShareOperandBufferWithUser( std::vector operand_indices = user->OperandIndices(operand); return operand_indices.size() == 1 && user_index[0] == operand_indices[0]; } + if (user->opcode() == HloOpcode::kTriangularSolve) { + // Only valid if there are no other users. + if (operand->users().size() != 1) { + return false; + } + std::vector operand_indices = user->OperandIndices(operand); + return operand_indices.size() == 1 && operand_indices[0] == 1; + } if (user->opcode() == HloOpcode::kCall) { // TODO(b/62548313): Remove when buffer assignment is module scoped and // does not assign buffers to calls. diff --git a/tensorflow/compiler/xla/service/tuple_points_to_analysis_test.cc b/tensorflow/compiler/xla/service/tuple_points_to_analysis_test.cc index d8875ca74716420ba0d54e4aba53c99001989996..551602613927671c9c37a4e8685df76d6a4ca9cf 100644 --- a/tensorflow/compiler/xla/service/tuple_points_to_analysis_test.cc +++ b/tensorflow/compiler/xla/service/tuple_points_to_analysis_test.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/hlo_creation_utils.h" #include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/instruction_fusion.h" @@ -623,7 +624,7 @@ class FusionPointsToAnalysisTest : public TuplePointsToAnalysisTest { void Run(const bool add_additional_gte0_user) { Shape input_shape = ShapeUtil::MakeShape(F32, {8}); Shape update_shape = ShapeUtil::MakeShape(F32, {3}); - Shape starts_shape = ShapeUtil::MakeShape(S32, {1}); + Shape starts_shape = ShapeUtil::MakeShape(S32, {}); Shape tuple_shape = ShapeUtil::MakeTupleShape({input_shape, update_shape, starts_shape}); @@ -657,7 +658,7 @@ class FusionPointsToAnalysisTest : public TuplePointsToAnalysisTest { HloInstruction::CreateGetTupleElement(starts_shape, tuple_param0, 2)); // Update 'input' with 'update' at dynamic 'starts' indices. builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - input_shape, input, update, starts)); + input_shape, input, update, {starts})); // Build computation and add it to module as entry computation. BuildModule(builder.Build()); @@ -721,9 +722,8 @@ class FusionPointsToAnalysisTest : public TuplePointsToAnalysisTest { // to fusion 'operand'. HloInstruction* GetFusionParameterForOperand(HloInstruction* fusion, HloInstruction* operand) { - auto it = std::find_if( - fusion->fused_instructions().begin(), - fusion->fused_instructions().end(), [=](const HloInstruction* fused) { + auto it = absl::c_find_if( + fusion->fused_instructions(), [&](const HloInstruction* fused) { return fused->opcode() == HloOpcode::kParameter && fusion->operand(fused->parameter_number()) == operand; }); @@ -883,12 +883,12 @@ TEST_F(DoesNotUseOperandBufferTest, FusedDynamicUpdateSlice) { // Create a DynamicUpdateSlice instruction of tuple element 1. auto starts = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({2}))); + HloInstruction::CreateConstant(LiteralUtil::CreateR0(2))); auto update = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({2.f, 2.f, 2.f}))); auto dynamic_update_slice = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - data_shape, gte1, update, starts)); + data_shape, gte1, update, {starts})); builder.AddInstruction( HloInstruction::CreateTuple({gte0, dynamic_update_slice})); @@ -977,12 +977,12 @@ TEST_F(CanShareOperandBufferWithUserTest, FusedDynamicUpdateSlice) { // Create a DynamicUpdateSlice instruction of tuple element 1. auto starts = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({2}))); + HloInstruction::CreateConstant(LiteralUtil::CreateR0(2))); auto update = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({2.f, 2.f, 2.f}))); auto dynamic_update_slice = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - data_shape, gte1, update, starts)); + data_shape, gte1, update, {starts})); builder.AddInstruction( HloInstruction::CreateTuple({gte0, dynamic_update_slice})); @@ -1004,7 +1004,7 @@ TEST_F(CanShareOperandBufferWithUserTest, DynamicUpdateSliceCanShare) { Shape data_shape = ShapeUtil::MakeShape(F32, {8}); Shape update_shape = ShapeUtil::MakeShape(F32, {4}); - Shape starts_shape = ShapeUtil::MakeShape(S32, {1}); + Shape starts_shape = ShapeUtil::MakeShape(S32, {}); auto data = builder.AddInstruction( HloInstruction::CreateParameter(0, data_shape, "data")); auto update = builder.AddInstruction( @@ -1012,7 +1012,7 @@ TEST_F(CanShareOperandBufferWithUserTest, DynamicUpdateSliceCanShare) { auto starts = builder.AddInstruction( HloInstruction::CreateParameter(2, starts_shape, "starts")); auto dus = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( - data_shape, data, update, starts)); + data_shape, data, update, {starts})); BuildModuleAndRunAnalysis(builder.Build()); @@ -1066,14 +1066,16 @@ TEST_F(CanShareOperandBufferWithUserTest, ScatterCanShare) { TEST_F(CanShareOperandBufferWithUserTest, SortCanShare) { auto builder = HloComputation::Builder(TestName()); + module_ = CreateNewVerifiedModule(); Shape keys_shape = ShapeUtil::MakeShape(F32, {8}); auto keys = builder.AddInstruction( HloInstruction::CreateParameter(0, keys_shape, "keys")); - auto sort = - builder.AddInstruction(HloInstruction::CreateSort(keys_shape, 0, keys)); + TF_ASSERT_OK_AND_ASSIGN( + auto* sort, MakeSortHlo(keys_shape, {keys}, 0, &builder, module_.get())); - BuildModuleAndRunAnalysis(builder.Build()); + computation_ = module_->AddEntryComputation(builder.Build()); + RunAnalysis(); EXPECT_TRUE( points_to_analysis_->CanShareOperandBufferWithUser(keys, {}, sort, {})); @@ -1081,6 +1083,7 @@ TEST_F(CanShareOperandBufferWithUserTest, SortCanShare) { TEST_F(CanShareOperandBufferWithUserTest, SortCanShareWithTupleUser) { auto builder = HloComputation::Builder(TestName()); + module_ = CreateNewVerifiedModule(); Shape keys_shape = ShapeUtil::MakeShape(F32, {8}); Shape values_shape = ShapeUtil::MakeShape(F32, {8}); @@ -1088,11 +1091,13 @@ TEST_F(CanShareOperandBufferWithUserTest, SortCanShareWithTupleUser) { HloInstruction::CreateParameter(0, keys_shape, "keys")); auto values = builder.AddInstruction( HloInstruction::CreateParameter(1, values_shape, "values")); - auto sort = builder.AddInstruction(HloInstruction::CreateSort( - ShapeUtil::MakeTupleShape({keys_shape, values_shape}), 0, keys, - {values})); + TF_ASSERT_OK_AND_ASSIGN( + auto* sort, + MakeSortHlo(ShapeUtil::MakeTupleShape({keys_shape, values_shape}), + {keys, values}, 0, &builder, module_.get())); - BuildModuleAndRunAnalysis(builder.Build()); + computation_ = module_->AddEntryComputation(builder.Build()); + RunAnalysis(); // The buffer for the keys can be shared with the first tuple entry. EXPECT_TRUE( diff --git a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc index a1c627a319f19c83fd676000f9de971efe377dbd..69cc8feb3f31ad782b9d3437d81d0ab8ce10aadb 100644 --- a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc +++ b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc @@ -89,7 +89,7 @@ static void CreateLoopInvariantCopy( HloInstruction* next_operand = frame->instruction->mutable_operand(frame->operand_index++); - if (hoisted_instructions->count(next_operand) || + if (hoisted_instructions->contains(next_operand) || next_operand == while_body_param) { continue; } @@ -241,7 +241,7 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody( auto is_invariant = [&](HloInstruction* op) { return hoisted_instructions.find(op) != hoisted_instructions.end() || - unhoisted_invariant_instructions.count(op) || + unhoisted_invariant_instructions.contains(op) || op->opcode() == HloOpcode::kConstant; }; diff --git a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion_test.cc b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion_test.cc index 8e7c4bc8828552e197b41f874c070d496b85a382..3587c016b4420163a607422b1acc838646fab83a 100644 --- a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion_test.cc +++ b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion_test.cc @@ -299,7 +299,7 @@ TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistBitcastAlone) { // bitcast either. auto m = CreateNewVerifiedModule(); auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); - auto scalar_f32 = ShapeUtil::MakeShape(F32, {}); + auto effective_scalar_s32 = ShapeUtil::MakeShape(S32, {1}); auto token_shape = ShapeUtil::MakeTokenShape(); Shape while_shape = ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32, token_shape}); @@ -314,10 +314,12 @@ TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistBitcastAlone) { HloInstruction::CreateGetTupleElement(scalar_s32, param, 1)); HloInstruction* in_token = builder.AddInstruction( HloInstruction::CreateGetTupleElement(token_shape, param, 2)); - HloInstruction* bitcast_inst = builder.AddInstruction( - HloInstruction::CreateUnary(scalar_f32, HloOpcode::kBitcast, gte_0)); - HloInstruction* out_token = builder.AddInstruction( - HloInstruction::CreateOutfeed(scalar_f32, bitcast_inst, in_token, "")); + HloInstruction* bitcast_inst = + builder.AddInstruction(HloInstruction::CreateUnary( + effective_scalar_s32, HloOpcode::kBitcast, gte_0)); + HloInstruction* out_token = + builder.AddInstruction(HloInstruction::CreateOutfeed( + effective_scalar_s32, bitcast_inst, in_token, "")); builder.AddInstruction( HloInstruction::CreateTuple({gte_0, gte_1, out_token})); @@ -352,9 +354,9 @@ TEST_F(WhileLoopInvariantCodeMotionTest, HoistBitcastIfNeeded) { // The bitcast's user can be hoisted, so hoist the bitcast too. auto m = CreateNewVerifiedModule(); auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); - auto scalar_f32 = ShapeUtil::MakeShape(F32, {}); - Shape while_shape = - ShapeUtil::MakeTupleShape({scalar_s32, scalar_f32, scalar_f32}); + auto effective_scalar_s32 = ShapeUtil::MakeShape(S32, {1}); + Shape while_shape = ShapeUtil::MakeTupleShape( + {scalar_s32, effective_scalar_s32, effective_scalar_s32}); HloComputation* while_body = [&]() { HloComputation::Builder builder(TestName() + ".while_body"); @@ -363,12 +365,13 @@ TEST_F(WhileLoopInvariantCodeMotionTest, HoistBitcastIfNeeded) { HloInstruction* gte_0 = builder.AddInstruction( HloInstruction::CreateGetTupleElement(scalar_s32, param, 0)); HloInstruction* gte_1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_f32, param, 1)); - HloInstruction* bitcast_inst = builder.AddInstruction( - HloInstruction::CreateUnary(scalar_f32, HloOpcode::kBitcast, gte_0)); + HloInstruction::CreateGetTupleElement(effective_scalar_s32, param, 1)); + HloInstruction* bitcast_inst = + builder.AddInstruction(HloInstruction::CreateUnary( + effective_scalar_s32, HloOpcode::kBitcast, gte_0)); HloInstruction* add_inst = builder.AddInstruction(HloInstruction::CreateBinary( - scalar_f32, HloOpcode::kAdd, bitcast_inst, gte_1)); + effective_scalar_s32, HloOpcode::kAdd, bitcast_inst, gte_1)); builder.AddInstruction( HloInstruction::CreateTuple({gte_0, gte_1, add_inst})); diff --git a/tensorflow/compiler/xla/service/while_loop_simplifier.cc b/tensorflow/compiler/xla/service/while_loop_simplifier.cc index 772faa25e299ebd96d46545905bbc66b6e23c82e..386ffb995477ff1b4aef73080b6a6fd988dd1980 100644 --- a/tensorflow/compiler/xla/service/while_loop_simplifier.cc +++ b/tensorflow/compiler/xla/service/while_loop_simplifier.cc @@ -109,8 +109,7 @@ static StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { // operand appears in, but it may appear more than once! if (user->user_count() == 1 && user->users().front() == while_body_root && while_body_root->operand_index(user) == user->tuple_index() && - std::count(while_body_root->operands().begin(), - while_body_root->operands().end(), user) == 1) { + absl::c_count(while_body_root->operands(), user) == 1) { continue; } @@ -127,7 +126,7 @@ static StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { // through to the while body's root, count that element as "used", since // removing that element would be observable. for (int64 i = 0; i < while_body_root->operand_count(); ++i) { - if (used_tuple_indices.count(i)) { + if (used_tuple_indices.contains(i)) { continue; } @@ -158,7 +157,7 @@ static StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { // Build up maps from the old/new to the new/old tuple indices. std::vector new_to_old_tuple_idx(used_tuple_indices.begin(), used_tuple_indices.end()); - std::sort(new_to_old_tuple_idx.begin(), new_to_old_tuple_idx.end()); + absl::c_sort(new_to_old_tuple_idx); absl::flat_hash_map old_to_new_tuple_idx; for (int64 new_idx = 0; new_idx < new_to_old_tuple_idx.size(); ++new_idx) { @@ -181,7 +180,7 @@ static StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { // replace the old instructions after we remove unused elements from the while // tuple. auto make_while_computation_replacements = [&](const HloComputation* comp) { - std::unordered_map> + absl::flat_hash_map> replacements; auto* param = comp->parameter_instruction(0); @@ -233,7 +232,7 @@ static StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { while_cond->CloneWithReplacements( make_while_computation_replacements(while_cond)); - std::unordered_map> + absl::flat_hash_map> while_body_replacements = make_while_computation_replacements(while_body); std::vector new_while_body_root_elems; new_while_body_root_elems.reserve(new_to_old_tuple_idx.size()); diff --git a/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc b/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc index 3713989ca2f64ee1d94c9f77255017909d957da2..ecca76b1e86d833c73fbb9bad6a341660a7d2669 100644 --- a/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc +++ b/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc @@ -407,13 +407,12 @@ TEST_F(WhileLoopSimplifierTest, RemoveUnusedLoopOperands) { // The original while instruction is still left in the module as a dead // instruction, find a while instruction with a different name as the new // while instruction. + const auto& instrs = m->entry_computation()->instructions(); HloInstruction* new_while_op = - *std::find_if(m->entry_computation()->instructions().begin(), - m->entry_computation()->instructions().end(), - [&](const HloInstruction* instr) { - return (instr->opcode() == HloOpcode::kWhile && - instr->name() != "while"); - }); + *absl::c_find_if(instrs, [&](const HloInstruction* instr) { + return (instr->opcode() == HloOpcode::kWhile && + instr->name() != "while"); + }); auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); EXPECT_TRUE( diff --git a/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.cc b/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.cc index a76fbf3f66adae0a5e5357178bc576bbc74701c7..661b7aa7d99ca549da6a509812760a1665d60919 100644 --- a/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.cc +++ b/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.cc @@ -37,9 +37,15 @@ StatusOr ZeroSizedHloElimination::Run(HloModule* module) { } if (comp->IsRemovable(instruction) && ShapeUtil::IsZeroElementArray(instruction->shape())) { + // If the instruction doesn't have a layout, use a default layout for + // the literal. + Shape shape = instruction->shape(); + if (!LayoutUtil::HasLayout(shape)) { + LayoutUtil::SetToDefaultLayout(&shape); + } TF_RETURN_IF_ERROR(comp->ReplaceWithNewInstruction( - instruction, HloInstruction::CreateConstant( - Literal::CreateFromShape(instruction->shape())))); + instruction, + HloInstruction::CreateConstant(Literal::CreateFromShape(shape)))); changed = true; } } diff --git a/tensorflow/compiler/xla/service/zero_sized_hlo_elimination_test.cc b/tensorflow/compiler/xla/service/zero_sized_hlo_elimination_test.cc index a546a6d39cc55d1f327b8449c7d26cd4c95dbf98..572a79609e7a912277af0fd2ba43f9a1e14a6f52 100644 --- a/tensorflow/compiler/xla/service/zero_sized_hlo_elimination_test.cc +++ b/tensorflow/compiler/xla/service/zero_sized_hlo_elimination_test.cc @@ -82,5 +82,18 @@ TEST_F(ZeroSizedHloEliminationTest, DoesNotEliminateConstant) { EXPECT_FALSE(changed); } +TEST_F(ZeroSizedHloEliminationTest, ZeroSizedInstructionWithoutLayoutFolded) { + Shape op_shape = ShapeUtil::MakeShape(F32, {4, 0}); + op_shape.clear_layout(); + HloInstruction* param1 = builder_.AddInstruction( + HloInstruction::CreateParameter(1, op_shape, "zero sized param 1")); + HloInstruction* param2 = builder_.AddInstruction( + HloInstruction::CreateParameter(2, op_shape, "zero sized param 2")); + builder_.AddInstruction( + HloInstruction::CreateBinary(op_shape, HloOpcode::kAdd, param1, param2)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, RunZeroSizedElimination()); + EXPECT_TRUE(changed); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/shape.cc b/tensorflow/compiler/xla/shape.cc index c04a91615437d315934866448b70846ccf3bad40..93d630b8f736f6c41d4014ef6415e80eac5a65ec 100644 --- a/tensorflow/compiler/xla/shape.cc +++ b/tensorflow/compiler/xla/shape.cc @@ -27,7 +27,23 @@ Shape::Shape(const ShapeProto& shape_proto) { for (const int64 dimension : shape_proto.dimensions()) { add_dimensions(dimension); } - for (int i = 0; i < shape_proto.is_dynamic_dimension_size(); i++) { + // A malformed proto may have different is_dynamic_dimension_size and + // dimensions_size. Since C++ is evil, and we have no good way of bailing out + // in a constructor, conservatively trim the is_dynamic_dimension size. + // TODO(b/120111794): Make this a hard error when we have a factory method + // instead of a constructor. + if (shape_proto.dimensions_size() != + shape_proto.is_dynamic_dimension_size()) { + if (shape_proto.is_dynamic_dimension_size() != 0) { + LOG(ERROR) << "Malformed shape proto: number of is_dynamic_dimension " + "fields does not match number of dimension fields"; + } else { + LOG(WARNING) << "Malformed shape proto: is_dynamic_dimension is empty"; + } + } + int64 num_dynamic_dimension_fields = std::min( + shape_proto.dimensions_size(), shape_proto.is_dynamic_dimension_size()); + for (int i = 0; i < num_dynamic_dimension_fields; i++) { dynamic_dimensions_[i] = shape_proto.is_dynamic_dimension(i); } tuple_shapes_.reserve(shape_proto.tuple_shapes_size()); @@ -68,19 +84,18 @@ string Shape::ToString(bool print_layout) const { } bool Shape::is_static() const { - if (ShapeUtil::IsTuple(*this)) { + if (IsTuple()) { for (const Shape& subshape : tuple_shapes_) { if (!subshape.is_static()) { return false; } } } - return !std::any_of(dynamic_dimensions_.begin(), dynamic_dimensions_.end(), - [](bool b) { return b; }); + return !absl::c_any_of(dynamic_dimensions_, [](bool b) { return b; }); } void Shape::DeleteDimension(int64 dim_to_delete) { - CHECK(ShapeUtil::IsArray(*this)); + CHECK(IsArray()); CHECK_GE(dim_to_delete, 0); CHECK_LT(dim_to_delete, dimensions_.size()); dimensions_.erase(dimensions_.begin() + dim_to_delete); @@ -101,6 +116,61 @@ void Shape::DeleteDimension(int64 dim_to_delete) { } } +bool Shape::Equal::operator()(const Shape& lhs, const Shape& rhs) { + if (lhs.IsTuple()) { + return rhs.IsTuple() && + absl::c_equal( + lhs.tuple_shapes(), rhs.tuple_shapes(), + [=](const Shape& l, const Shape& r) { return (*this)(l, r); }); + } else if (!lhs.IsArray()) { + // Non-tuple, non-array tupes such as opaque and token types are trivially + // the same. + return lhs.element_type() == rhs.element_type(); + } + + if (!rhs.IsArray()) { + return false; + } + + if (!ignore_element_type_) { + if ((ignore_fp_precision_ && + !ShapeUtil::SameElementTypeIgnoringFpPrecision(lhs, rhs)) || + (!ignore_fp_precision_ && !ShapeUtil::SameElementType(lhs, rhs))) { + VLOG(3) << "CompareShapes: lhs element type != rhs element type"; + return false; + } + } + + if (!ignore_layout_) { + if (lhs.layout().format() != rhs.layout().format()) { + VLOG(3) << "CompareShapes: lhs layout format != rhs layout format"; + return false; + } + if (LayoutUtil::IsDenseArray(lhs)) { + if (lhs.layout() != rhs.layout()) { + VLOG(3) << "CompareShapes: lhs layout != rhs layout"; + return false; + } + } + } + + if (!ShapeUtil::SameDimensions(lhs, rhs)) { + VLOG(3) << "CompareShapes: lhs dimensions != rhs dimensions"; + return false; + } + + if (!ignore_dynamic_dimension_) { + for (int i = 0; i < lhs.rank(); ++i) { + if (lhs.is_dynamic_dimension(i) != rhs.is_dynamic_dimension(i)) { + VLOG(3) + << "CompareShapes: lhs and rhs have different dynamic dimensions."; + return false; + } + } + } + return true; +} + std::ostream& operator<<(std::ostream& out, const Shape& shape) { out << shape.ToString(/*print_layout=*/true); return out; diff --git a/tensorflow/compiler/xla/shape.h b/tensorflow/compiler/xla/shape.h index dc4cdc31a74d43471b72a71d9d436408e0e62deb..1d594904e0b9e6f1779674e75b41b7a597788bac 100644 --- a/tensorflow/compiler/xla/shape.h +++ b/tensorflow/compiler/xla/shape.h @@ -72,6 +72,10 @@ class Shape { dynamic_dimensions_[dimension] = is_dynamic; } + const std::vector& dynamic_dimensions() const { + return dynamic_dimensions_; + } + // Add dimension_upper_bound(). // Removes the given dimension form the shape. Layout, if it exists, is @@ -138,6 +142,49 @@ class Shape { string ShortDebugString() const { return ToProto().ShortDebugString(); } string DebugString() const { return ToProto().DebugString(); } + // Equal is a configurable functor to check the equality of two shapes. + // + // Examples: + // + // - Comparing two shapes ignoring they layout difference: + // Equal().IgnoreLayout()(shape1, shape2); + // + // - Comparing two shapes ignoring they layout and element type difference: + // Equal().IgnoreLayout().IgnoreElementType()(shape1, shape2); + class Equal { + public: + Equal() = default; + + bool operator()(const Shape& lhs, const Shape& rhs); + + Equal& IgnoreLayout() { + ignore_layout_ = true; + return *this; + } + Equal& IgnoreElementType() { + ignore_element_type_ = true; + return *this; + } + Equal& IgnoreFpPrecision() { + ignore_fp_precision_ = true; + return *this; + } + Equal& IgnoreDynamicDimension() { + ignore_dynamic_dimension_ = true; + return *this; + } + + public: + bool ignore_layout_ = false; + bool ignore_element_type_ = false; + bool ignore_fp_precision_ = false; + bool ignore_dynamic_dimension_ = false; + }; + + // Test that all fields of the shape are the same, equivalent to Equal(). + bool operator==(const Shape& other) const { return Equal()(*this, other); } + bool operator!=(const Shape& other) const { return !(*this == other); } + private: // The element type of this shape (tuple, array, etc). PrimitiveType element_type_ = PRIMITIVE_TYPE_INVALID; diff --git a/tensorflow/compiler/xla/shape_test.cc b/tensorflow/compiler/xla/shape_test.cc index 55ce5fe884e98e474253be9ef694f1b8137b4b01..526abafea5cc244418a4ec05db7da6203716b483 100644 --- a/tensorflow/compiler/xla/shape_test.cc +++ b/tensorflow/compiler/xla/shape_test.cc @@ -85,6 +85,24 @@ TEST_F(ShapeTest, DynamicShapeToString) { EXPECT_EQ("f32[<=23,44,55]", array_shape.ToString()); } +TEST_F(ShapeTest, EqualityTest) { + // Different layouts. + EXPECT_NE(ShapeUtil::MakeShapeWithLayout(F32, {23, 44}, {1, 0}), + ShapeUtil::MakeShapeWithLayout(F32, {23, 44}, {0, 1})); + + // Different dims. + EXPECT_NE(ShapeUtil::MakeShapeWithLayout(F32, {44, 23}, {1, 0}), + ShapeUtil::MakeShapeWithLayout(F32, {23, 44}, {1, 0})); + + // Different elements. + EXPECT_NE(ShapeUtil::MakeShapeWithLayout(S32, {44, 23}, {1, 0}), + ShapeUtil::MakeShapeWithLayout(F32, {23, 44}, {1, 0})); + + // Equal shapes. + EXPECT_EQ(ShapeUtil::MakeShapeWithLayout(F32, {23, 44}, {1, 0}), + ShapeUtil::MakeShapeWithLayout(F32, {23, 44}, {1, 0})); +} + TEST_F(ShapeTest, IsStatic) { EXPECT_TRUE(opaque_.is_static()); EXPECT_TRUE(token_.is_static()); diff --git a/tensorflow/compiler/xla/shape_util.cc b/tensorflow/compiler/xla/shape_util.cc index 7f68c0ae1fb1e3de12b0243829f8f02cc1137894..e6273c4e7f8ed3f8feab0ecd540ad1081f653c8b 100644 --- a/tensorflow/compiler/xla/shape_util.cc +++ b/tensorflow/compiler/xla/shape_util.cc @@ -85,81 +85,12 @@ bool ShapeIndexView::StartsWith(ShapeIndexView prefix) const { } namespace { - -// Recursive helper for comparing the equality of two shapes. Returns true if -// the shapes are the same. If compare_layouts is true, then layouts must also -// match. -bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts, - bool ignore_fp_precision) { - if ((ignore_fp_precision && - !ShapeUtil::SameElementTypeIgnoringFpPrecision(lhs, rhs)) || - (!ignore_fp_precision && !ShapeUtil::SameElementType(lhs, rhs))) { - VLOG(3) << "CompareShapes: lhs element type != rhs element type"; - return false; - } - - if (lhs.IsTuple()) { - return absl::c_equal(lhs.tuple_shapes(), rhs.tuple_shapes(), - [=](const Shape& l, const Shape& r) { - return CompareShapes(l, r, compare_layouts, - ignore_fp_precision); - }); - } else if (!lhs.IsArray()) { - // Non-tuple, non-array tupes such as opaque and token types are trivially - // the same. - return true; - } - - if (compare_layouts) { - if (lhs.layout().format() != rhs.layout().format()) { - return false; - } - if (LayoutUtil::IsDenseArray(lhs)) { - if (!absl::c_equal(LayoutUtil::MinorToMajor(lhs), - LayoutUtil::MinorToMajor(rhs))) { - VLOG(3) << "CompareShapes: lhs layout != rhs layout"; - return false; - } - - const auto& lhs_tiles = lhs.layout().tiles(); - const auto& rhs_tiles = rhs.layout().tiles(); - if (lhs_tiles.size() != rhs_tiles.size()) { - return false; - } - for (int64 i = 0; i < lhs_tiles.size(); i++) { - if (!absl::c_equal(lhs_tiles[i].dimensions(), - rhs_tiles[i].dimensions())) { - return false; - } - } - - if (lhs.layout().element_size_in_bits() != - rhs.layout().element_size_in_bits()) { - return false; - } - } - } - - if (!ShapeUtil::SameDimensions(lhs, rhs)) { - VLOG(3) << "CompareShapes: lhs dimensions != rhs dimensions"; - return false; - } - - for (int i = 0; i < ShapeUtil::Rank(lhs); ++i) { - if (lhs.is_dynamic_dimension(i) != rhs.is_dynamic_dimension(i)) { - VLOG(3) - << "CompareShapes: lhs and rhs have different dynamic dimensions."; - return false; - } - } - return true; -} - // Constructs and returns the new shape with the given minor_to_major order in // its Layout. StatusOr MakeShapeWithLayoutInternal( PrimitiveType element_type, absl::Span dimensions, - absl::Span minor_to_major) { + absl::Span minor_to_major, absl::Span tiles, + int64 element_size_in_bits) { if (dimensions.size() != minor_to_major.size()) { return InvalidArgument("Dimensions size is %ld, but layout size is %ld.", dimensions.size(), minor_to_major.size()); @@ -170,23 +101,19 @@ StatusOr MakeShapeWithLayoutInternal( } TF_ASSIGN_OR_RETURN(Shape shape, ShapeUtil::MakeValidatedShape(element_type, dimensions)); - auto min2maj = shape.mutable_layout()->mutable_minor_to_major(); - min2maj->clear(); - for (int64 value : minor_to_major) { - min2maj->push_back(value); - } + *shape.mutable_layout() = + LayoutUtil::MakeLayout(minor_to_major, tiles, element_size_in_bits); if (!shape.has_layout()) { return InvalidArgument("Shape has no layout."); } TF_RETURN_IF_ERROR(ShapeUtil::ValidateShape(shape)); return shape; } - } // namespace /* static */ bool ShapeUtil::Equal(const Shape& lhs, const Shape& rhs) { - bool equal = CompareShapes(lhs, rhs, /*compare_layouts=*/true, - /*ignore_fp_precision=*/false); + bool equal = Shape::Equal()(lhs, rhs); + if (!equal && VLOG_IS_ON(3)) { VLOG(3) << "ShapeUtil::Equal differ: lhs = " << lhs.ShortDebugString() << ", rhs = " << rhs.ShortDebugString(); @@ -197,8 +124,7 @@ StatusOr MakeShapeWithLayoutInternal( /* static */ bool ShapeUtil::EqualIgnoringFpPrecision(const Shape& lhs, const Shape& rhs) { - bool equal = CompareShapes(lhs, rhs, /*compare_layouts=*/true, - /*ignore_fp_precision=*/true); + bool equal = Shape::Equal().IgnoreFpPrecision()(lhs, rhs); if (!equal && VLOG_IS_ON(3)) { VLOG(3) << "ShapeUtil::EqualIgnoringFpPrecision differ: lhs = " << lhs.ShortDebugString() << ", rhs = " << rhs.ShortDebugString(); @@ -207,11 +133,6 @@ StatusOr MakeShapeWithLayoutInternal( return equal; } -/* static */ int64 ShapeUtil::Rank(const Shape& shape) { - CHECK(shape.IsArray()) << "Non-arrays do not have a rank, shape: " << shape; - return shape.dimensions_size(); -} - /* static */ int64 ShapeUtil::TrueRank(const Shape& shape) { int64 accum = 0; for (int64 dimension : shape.dimensions()) { @@ -266,8 +187,10 @@ StatusOr MakeShapeWithLayoutInternal( /* static */ Shape ShapeUtil::MakeShapeWithLayout( PrimitiveType element_type, absl::Span dimensions, - absl::Span minor_to_major) { - return MakeShapeWithLayoutInternal(element_type, dimensions, minor_to_major) + absl::Span minor_to_major, absl::Span tiles, + int64 element_size_in_bits) { + return MakeShapeWithLayoutInternal(element_type, dimensions, minor_to_major, + tiles, element_size_in_bits) .ValueOrDie(); } @@ -343,7 +266,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( /* static */ void ShapeUtil::AppendMajorDimension(int bound, Shape* shape) { CHECK(LayoutUtil::IsDenseArray(*shape)); - shape->mutable_layout()->add_minor_to_major(Rank(*shape)); + shape->mutable_layout()->add_minor_to_major(shape->rank()); shape->add_dimensions(bound); TF_DCHECK_OK(ValidateShape(*shape)); } @@ -358,7 +281,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( } /* static */ bool ShapeUtil::ElementHasBitWidth(const Shape& shape, int bits) { - if (!IsArray(shape)) { + if (!shape.IsArray()) { return false; } return primitive_util::BitWidth(shape.element_type()) == bits; @@ -382,6 +305,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( case U32: case U64: case C64: + case C128: case TUPLE: case OPAQUE: case TOKEN: @@ -400,27 +324,24 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( return primitive_util::IsFloatingPointType(shape.element_type()); } -/* static */ bool ShapeUtil::IsArray(const Shape& shape) { - return IsArrayPrimitiveType(shape.element_type()); -} - /* static */ bool ShapeUtil::IsNestedTuple(const Shape& shape) { - return IsTuple(shape) && std::any_of(shape.tuple_shapes().begin(), - shape.tuple_shapes().end(), IsTuple); + return shape.IsTuple() && + absl::c_any_of(shape.tuple_shapes(), + [](const Shape& s) { return s.IsTuple(); }); } /* static */ bool ShapeUtil::IsEmptyTuple(const Shape& shape) { - return IsTuple(shape) && TupleElementCount(shape) == 0; + return shape.IsTuple() && TupleElementCount(shape) == 0; } /* static */ int64 ShapeUtil::TupleElementCount(const Shape& shape) { - CHECK(IsTuple(shape)) << HumanString(shape); + CHECK(shape.IsTuple()) << HumanString(shape); return shape.tuple_shapes_size(); } /* static */ const Shape& ShapeUtil::GetTupleElementShape(const Shape& shape, int64 index) { - CHECK(IsTuple(shape)); + CHECK(shape.IsTuple()); CHECK_GT(TupleElementCount(shape), index); TF_DCHECK_OK(ValidateShapeWithOptionalLayout(shape.tuple_shapes(index))); return shape.tuple_shapes(index); @@ -436,7 +357,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( /* static */ Shape ShapeUtil::SliceTuple(const Shape& tuple, int64 start, int64 limit) { TF_DCHECK_OK(ValidateShapeWithOptionalLayout(tuple)); - CHECK(IsTuple(tuple)); + CHECK(tuple.IsTuple()); CHECK_LE(start, TupleElementCount(tuple)); CHECK_LE(limit, TupleElementCount(tuple)); @@ -453,15 +374,9 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( complex_shape.element_type())); } -/* static */ bool ShapeUtil::ShapeIs(const Shape& shape, - PrimitiveType element_type, - std::initializer_list dimensions) { - return Equal(shape, MakeShape(element_type, dimensions)); -} - /* static */ int64 ShapeUtil::ElementsIn(const Shape& shape) { - DCHECK(IsArray(shape)) << ShapeUtil::HumanString(shape); - DCHECK_EQ(shape.dimensions_size(), Rank(shape)); + DCHECK(shape.IsArray()) << ShapeUtil::HumanString(shape); + DCHECK_EQ(shape.dimensions_size(), shape.rank()); if (shape.dimensions().size() == 1) { return shape.dimensions()[0]; } @@ -471,8 +386,8 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( } /* static */ int64 ShapeUtil::ElementsInRecursive(const Shape& shape) { - CHECK(IsArray(shape) || IsTuple(shape)); - if (IsArray(shape)) { + CHECK(shape.IsArray() || shape.IsTuple()); + if (shape.IsArray()) { return ElementsIn(shape); } int64 count = 0; @@ -505,7 +420,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( } /* static */ string ShapeUtil::HumanString(const Shape& shape) { - if (IsTuple(shape)) { + if (shape.IsTuple()) { string text = "("; const char* prefix = ""; for (const Shape& elem_shape : shape.tuple_shapes()) { @@ -529,7 +444,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( } /* static */ string ShapeUtil::HumanStringWithLayout(const Shape& shape) { - if (IsTuple(shape)) { + if (shape.IsTuple()) { string text = "("; const char* prefix = ""; for (const Shape& elem_shape : shape.tuple_shapes()) { @@ -542,10 +457,11 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( string result = StrCat( primitive_util::LowercasePrimitiveTypeName(shape.element_type()), "["); for (int i = 0; i < shape.dimensions().size(); i++) { - StrAppend(&result, (i > 0) ? "," : "", shape.dimensions(i)); + StrAppend(&result, (i > 0) ? "," : "", + shape.is_dynamic_dimension(i) ? "<=" : "", shape.dimensions(i)); } result += "]"; - if (!IsScalar(shape) && IsArray(shape)) { + if (!IsScalar(shape) && shape.IsArray()) { if (LayoutUtil::HasLayout(shape)) { StrAppend(&result, LayoutUtil::HumanString(shape.layout())); } @@ -574,37 +490,17 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( } /* static */ bool ShapeUtil::Compatible(const Shape& lhs, const Shape& rhs) { - return CompareShapes(lhs, rhs, /*compare_layouts=*/false, - /*ignore_fp_precision=*/false); + return Shape::Equal().IgnoreLayout()(lhs, rhs); } /* static */ bool ShapeUtil::CompatibleIgnoringElementType(const Shape& lhs, const Shape& rhs) { - if (IsArray(lhs)) { - return IsArray(rhs) && SameDimensions(lhs, rhs); - } else if (lhs.element_type() == TUPLE) { - return rhs.element_type() == TUPLE && - absl::c_equal(lhs.tuple_shapes(), rhs.tuple_shapes(), - CompatibleIgnoringElementType); - } else { - // Opaque, token, etc types are vacuously compatible. - return lhs.element_type() == rhs.element_type(); - } + return Shape::Equal().IgnoreElementType().IgnoreLayout()(lhs, rhs); } /* static */ bool ShapeUtil::CompatibleIgnoringFpPrecision(const Shape& lhs, const Shape& rhs) { - if (IsArray(lhs)) { - return IsArray(rhs) && SameElementTypeIgnoringFpPrecision(lhs, rhs) && - CompatibleIgnoringElementType(lhs, rhs); - } else if (lhs.element_type() == TUPLE) { - return rhs.element_type() == TUPLE && - absl::c_equal(lhs.tuple_shapes(), rhs.tuple_shapes(), - CompatibleIgnoringFpPrecision); - } else { - // Opaque, token, etc types are vacuously compatible. - return lhs.element_type() == rhs.element_type(); - } + return Shape::Equal().IgnoreFpPrecision().IgnoreLayout()(lhs, rhs); } /* static */ int64 ShapeUtil::GetDimension(const Shape& shape, @@ -615,7 +511,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( /* static */ int64 ShapeUtil::GetDimensionNumber(const Shape& shape, int64 dimension_number) { if (dimension_number < 0) { - dimension_number += Rank(shape); + dimension_number += shape.rank(); } CHECK_GE(dimension_number, 0); return dimension_number; @@ -652,6 +548,8 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( return sizeof(double); case C64: return sizeof(complex64); + case C128: + return sizeof(complex128); case TOKEN: // Tokens require no space. return 0; @@ -669,7 +567,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( TF_DCHECK_OK(ValidateShape(shape)); if (shape.element_type() == TUPLE) { return ByteSizeOfTupleIndexTable(shape, pointer_size); - } else if (IsArray(shape)) { + } else if (shape.IsArray()) { int64 byte_size = ByteSizeOfElements(shape); if (LayoutUtil::IsSparseArray(shape)) { byte_size += ByteSizeOfSparseIndices(shape); @@ -755,10 +653,10 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( return Status::OK(); } - if (LayoutUtil::IsSparseArray(shape) && Rank(shape) == 0) { + if (LayoutUtil::IsSparseArray(shape) && shape.rank() == 0) { return InvalidArgument("sparse arrays must have rank > 0"); } - for (int64 i = 0; i < Rank(shape); ++i) { + for (int64 i = 0; i < shape.rank(); ++i) { int64 dimension = shape.dimensions(i); if (dimension < 0) { return InvalidArgument( @@ -774,7 +672,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( /* static */ Status ShapeUtil::ValidateShapeSize(const Shape& shape) { VLOG(3) << "Validating shape size: " << ShapeUtil::HumanString(shape); - if (!IsArray(shape)) { + if (!shape.IsArray()) { return Status::OK(); } @@ -867,7 +765,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( ShapeIndexView index) { const Shape* subshape = &shape; for (auto i : index) { - if (!IsTuple(*subshape) || i >= subshape->tuple_shapes_size() || i < 0) { + if (!subshape->IsTuple() || i >= subshape->tuple_shapes_size() || i < 0) { return false; } subshape = &subshape->tuple_shapes(i); @@ -879,7 +777,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( ShapeIndexView index) { const Shape* return_shape = &shape; for (auto i : index) { - CHECK(IsTuple(*return_shape)) + CHECK(return_shape->IsTuple()) << "Invalid index " << index << " for shape " << shape; return_shape = &return_shape->tuple_shapes(i); } @@ -890,7 +788,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( const Shape& shape, ShapeIndexView index) { const Shape* return_shape = &shape; for (auto i : index) { - if (!IsTuple(*return_shape) || i < 0 || + if (!return_shape->IsTuple() || i < 0 || i >= return_shape->tuple_shapes_size()) { return InvalidArgument( "Shape index %s not a valid subshape index for tuple with shape %s", @@ -905,7 +803,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( ShapeIndexView index) { Shape* return_shape = shape; for (auto i : index) { - CHECK(IsTuple(*return_shape)); + CHECK(return_shape->IsTuple()); return_shape = return_shape->mutable_tuple_shapes(i); } return return_shape; @@ -913,11 +811,11 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( /* static */ bool ShapeUtil::IsLeafIndex(const Shape& shape, const ShapeIndex& index) { - return !IsTuple(GetSubshape(shape, index)); + return !GetSubshape(shape, index).IsTuple(); } /* static */ int64 ShapeUtil::GetLeafCount(const Shape& shape) { - if (!IsTuple(shape)) { + if (!shape.IsTuple()) { return 1; } int64 count = 0; @@ -1031,6 +929,10 @@ Status ForEachMutableSubshapeHelper( for (auto dim : Permute(permutation, shape.dimensions())) { new_shape.add_dimensions(dim); } + for (int64 i = 0; i < shape.rank(); i++) { + new_shape.set_dynamic_dimension(permutation[i], + shape.is_dynamic_dimension(i)); + } // If `shape` has a layout, by contract we choose a new layout such that the // transpose defined by this permutation is a bitcast. @@ -1081,8 +983,8 @@ Status ForEachMutableSubshapeHelper( /* static */ std::tuple, std::vector> ShapeUtil::InsertedOrDeleted1SizedDimensions(const Shape& shape_pre, const Shape& shape_post) { - CHECK(IsArray(shape_pre)); - CHECK(IsArray(shape_post)); + CHECK(shape_pre.IsArray()); + CHECK(shape_post.IsArray()); auto nil = std::make_tuple(false, std::vector(), std::vector()); @@ -1129,7 +1031,7 @@ ShapeUtil::InsertedOrDeleted1SizedDimensions(const Shape& shape_pre, auto unmodified_dim_pair = i < unmodified_dims.size() ? unmodified_dims[i] - : std::make_pair(Rank(shape_pre), Rank(shape_post)); + : std::make_pair(shape_pre.rank(), shape_post.rank()); if (!check_modified_dims(prior_unmodified_dim_pair, unmodified_dim_pair)) { return nil; } @@ -1141,8 +1043,8 @@ ShapeUtil::InsertedOrDeleted1SizedDimensions(const Shape& shape_pre, /* static */ std::vector> ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, const Shape& output_shape) { - CHECK(IsArray(input_shape)); - CHECK(IsArray(output_shape)); + CHECK(input_shape.IsArray()); + CHECK(output_shape.IsArray()); // Unmodified dimensions are merely common factors of rank 1. auto common_factors = CommonFactors(AsInt64Slice(input_shape.dimensions()), @@ -1192,8 +1094,8 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, /* static */ bool ShapeUtil::ReshapeIsBitcast(const Shape& input_shape, const Shape& output_shape) { - CHECK(IsArray(input_shape)); - CHECK(IsArray(output_shape)); + CHECK(input_shape.IsArray()); + CHECK(output_shape.IsArray()); CHECK(LayoutUtil::HasLayout(input_shape)); CHECK(LayoutUtil::HasLayout(output_shape)); @@ -1321,12 +1223,12 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, Shape output_shape_dim0_major = MakeShapeWithDescendingLayout( output_shape.element_type(), AsInt64Slice(output_shape.dimensions())); - for (int64 input_dim = 0; input_dim < Rank(input_shape); ++input_dim) { + for (int64 input_dim = 0; input_dim < input_shape.rank(); ++input_dim) { if (input_shape.dimensions(input_dim) <= 1) { continue; } - std::vector input_unit_index(Rank(input_shape), 0); + std::vector input_unit_index(input_shape.rank(), 0); input_unit_index[input_dim] = 1; int64 logical_linear_index = IndexUtil::MultidimensionalIndexToLinearIndex(input_shape_dim0_major, @@ -1352,11 +1254,11 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, /* static */ absl::optional ShapeUtil::AlignLayouts( const Shape& input_shape, const Shape& output_shape) { - CHECK(IsArray(input_shape)); - CHECK(IsArray(output_shape)); + CHECK(input_shape.IsArray()); + CHECK(output_shape.IsArray()); - int64 input_rank = Rank(input_shape); - int64 output_rank = Rank(output_shape); + int64 input_rank = input_shape.rank(); + int64 output_rank = output_shape.rank(); // First, calculate an alignment of the dimensions. A consecutive sequence of // input dimensions and output dimensions belong to the same alignment part if @@ -1493,14 +1395,14 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, /* static */ Shape ShapeUtil::DeleteDimension(int64 dim_to_delete, Shape shape) { - CHECK(IsArray(shape)); + CHECK(shape.IsArray()); shape.DeleteDimension(dim_to_delete); return shape; } /* static */ Shape ShapeUtil::FilterDimensions( const std::function& p, Shape shape) { - CHECK(IsArray(shape)); + CHECK(shape.IsArray()); std::vector dims_to_delete; for (int64 i = shape.dimensions().size() - 1; i >= 0; --i) { if (!p(i)) { diff --git a/tensorflow/compiler/xla/shape_util.h b/tensorflow/compiler/xla/shape_util.h index c8295e85ce15722a3da7ab5585e6d499d3b5efa2..7f610a6085d6fbe3d3143d5027cdc43d4b07bcbf 100644 --- a/tensorflow/compiler/xla/shape_util.h +++ b/tensorflow/compiler/xla/shape_util.h @@ -185,7 +185,7 @@ class ShapeUtil { // may not actually be able to store this number of elements. See // LayoutUtil::MaxSparseElements(shape) to obtain the maximum number of // elements that can be stored in a sparse shape. - // Precondition: IsArray(shape) + // Precondition: shape.IsArray() static int64 ElementsIn(const Shape& shape); // As ElementsIn(), but recurses through tuples. @@ -296,11 +296,6 @@ class ShapeUtil { // As Equal, but allow one of lhs and rhs to be F16 while the other is F32. static bool EqualIgnoringFpPrecision(const Shape& lhs, const Shape& rhs); - // Returns the rank (number of dimensions) of the given shape. - // Precondition: !IsTuple(shape) - ABSL_DEPRECATED("Use `Shape::rank` instead.") - static int64 Rank(const Shape& shape); - // Returns the number of dimensions for which the dimension is not (trivially) // 1. e.g., f32[2x1x1] has a true rank of 1D, the other dimensions are just // fluff. Note that zero dimensions are included in the true rank, e.g., @@ -314,10 +309,10 @@ class ShapeUtil { // Scalar-specific static bool IsScalar(const Shape& shape) { - return IsArray(shape) && Rank(shape) == 0; + return shape.IsArray() && shape.rank() == 0; } static bool IsEffectiveScalar(const Shape& shape) { - return IsArray(shape) && TrueRank(shape) == 0; + return shape.IsArray() && TrueRank(shape) == 0; } // Returns whether "shape" is a scalar (array) with the given element_type. @@ -403,7 +398,9 @@ class ShapeUtil { // Returns a value shape such that shape.has_layout(). static Shape MakeShapeWithLayout(PrimitiveType element_type, absl::Span dimensions, - absl::Span minor_to_major); + absl::Span minor_to_major, + absl::Span tiles = {}, + int64 element_size_in_bits = 0); static Shape MakeShapeWithSparseLayout(PrimitiveType element_type, absl::Span dimensions, @@ -457,31 +454,6 @@ class ShapeUtil { // that floating point numbers are signed. static bool ElementIsSigned(const Shape& shape); - // Returns whether the shape is a tuple. - ABSL_DEPRECATED("Use Shape::IsTuple instead.") - static bool IsTuple(const Shape& shape) { - return shape.element_type() == TUPLE; - } - - // Returns whether the shape is an opaque value (i.e. an 'existential' typed - // value that is passed to CustomCall operations). - ABSL_DEPRECATED("Use Shape::IsOpaque instead.") - static bool IsOpaque(const Shape& shape) { - return shape.element_type() == OPAQUE; - } - - // Returns whether the shape is an token value used for ordering - // side-effecting operations. - ABSL_DEPRECATED("Use Shape::IsToken instead.") - static bool IsToken(const Shape& shape) { - return shape.element_type() == TOKEN; - } - - // Returns whether the shape is an array. Note that scalars are considered - // arrays. - ABSL_DEPRECATED("Use Shape::IsArray instead.") - static bool IsArray(const Shape& shape); - // Returns whether the given primitive type corresponds to an array shape. static bool IsArrayPrimitiveType(PrimitiveType primitive_type); @@ -511,12 +483,6 @@ class ShapeUtil { // shape. static Shape ComplexComponentShape(const Shape& complex_shape); - // Shorthand for testing whether a shape is of a given element type and - // sequence of dimensions. - ABSL_DEPRECATED("Use Equal() instead.") - static bool ShapeIs(const Shape& shape, PrimitiveType element_type, - std::initializer_list dimensions); - // Returns true if the given shape has a subshape at the given index. static bool IndexIsValid(const Shape& shape, ShapeIndexView index); @@ -711,11 +677,9 @@ class ShapeUtil { template static void ForEachIndex(const Shape& shape, const FnType& visitor_function) { - ForEachIndexWithStatus(shape, - [&](absl::Span indices) { - return StatusOr(visitor_function(indices)); - }) - .IgnoreError(); + ForEachIndexWithStatus(shape, [&](absl::Span indices) { + return StatusOr(visitor_function(indices)); + }).IgnoreError(); } // A parallel version of ForEachIndex(WithStatus). This can only be used if @@ -764,7 +728,7 @@ class ShapeUtil { if (ShapeUtil::IsZeroElementArray(shape)) { return Status::OK(); } - CHECK_EQ(Rank(shape), base.size()); + CHECK_EQ(shape.rank(), base.size()); CHECK_EQ(incr.size(), base.size()); CHECK_EQ(count.size(), base.size()); const int64 rank = LayoutUtil::MinorToMajor(shape).size(); diff --git a/tensorflow/compiler/xla/shape_util_test.cc b/tensorflow/compiler/xla/shape_util_test.cc index 8e7c20819335bba161185bcf1237f709cfdb2a5d..126ae58293d12182e9b6e30f779f681829729526 100644 --- a/tensorflow/compiler/xla/shape_util_test.cc +++ b/tensorflow/compiler/xla/shape_util_test.cc @@ -538,10 +538,6 @@ TEST(ShapeUtilTest, InsertedOrDeleted1SizedDimensions) { ShapeUtil::InsertedOrDeleted1SizedDimensions(shape0, shape2))); } -TEST(ShapeUtilTest, ShapeIs) { - EXPECT_FALSE(ShapeUtil::ShapeIs(ShapeUtil::MakeShape(PRED, {2}), PRED, {})); -} - TEST(ShapeUtilTest, ForEachIndex) { struct ShapeDimensionAndNumberInvocations { std::vector dimensions; @@ -714,6 +710,26 @@ TEST(ShapeUtilTest, PermuteDimensionsLayout) { } while (std::next_permutation(layout.begin(), layout.end())); } +TEST(ShapeUtilTest, PermuteDynamicDimensions) { + Shape shape = + ShapeUtil::MakeShape(F32, {10, 100, 1000}, + /*dynamic_dimensions*/ {false, true, true}); + SCOPED_TRACE(absl::StrCat("shape=", shape.ToString())); + + std::vector permutation(3); + std::iota(permutation.begin(), permutation.end(), 0); + do { + SCOPED_TRACE(absl::StrCat("permutation=", absl::StrJoin(permutation, ","))); + + auto permuted = ShapeUtil::PermuteDimensions(permutation, shape); + for (int i = 0; i < shape.rank(); i++) { + EXPECT_EQ(permuted.dimensions(permutation[i]), shape.dimensions(i)); + EXPECT_EQ(permuted.is_dynamic_dimension(permutation[i]), + shape.is_dynamic_dimension(i)); + } + } while (std::next_permutation(permutation.begin(), permutation.end())); +} + TEST(AlgebraicSimplifierTest, ReshapeIsBitcast_3x2x2_6x2_Dim0IsMostMinor) { EXPECT_FALSE(ShapeUtil::ReshapeIsBitcast( ShapeUtil::MakeShapeWithLayout(F32, {3, 2, 2}, {0, 1, 2}), diff --git a/tensorflow/compiler/xla/sparse_index_array.h b/tensorflow/compiler/xla/sparse_index_array.h index a96d483462efd77ae4761541e8c79b2c84fa49f3..0c25355467da3fd346d80db790d78252869975ef 100644 --- a/tensorflow/compiler/xla/sparse_index_array.h +++ b/tensorflow/compiler/xla/sparse_index_array.h @@ -135,7 +135,7 @@ void SparseIndexArray::SortWithValues(absl::Span values) { auto sort_order_less = [this](int64 lhs, int64 rhs) { return IndexUtil::CompareIndices(At(lhs), At(rhs)) < 0; }; - std::sort(sort_order.begin(), sort_order.end(), sort_order_less); + absl::c_sort(sort_order, sort_order_less); // Reorder the array elements according to sort_order. Work through the array // and follow cycles so we can do the reorder in-place. diff --git a/tensorflow/compiler/xla/status_macros.cc b/tensorflow/compiler/xla/status_macros.cc index b88fe367d7416a26c1147fd5e10fb20772814fe5..aa7238f07d432aabb44d2cbed66786217e6a846c 100644 --- a/tensorflow/compiler/xla/status_macros.cc +++ b/tensorflow/compiler/xla/status_macros.cc @@ -25,6 +25,13 @@ limitations under the License. namespace xla { namespace status_macros { +ABSL_CONST_INIT const char kPossibleAutoJitAlternative[] = + "This error might be occurring with the use of xla.compile. If it is not " + "necessary that every Op be compiled with XLA, an alternative is to use " + "auto_jit with OptimizerOptions.global_jit_level = ON_2 or the environment " + "variable TF_XLA_FLAGS=\"tf_xla_auto_jit=2\" which will attempt to use xla " + "to compile as much of the graph as the compiler is able to."; + static Status MakeStatus(tensorflow::error::Code code, const string& message) { return Status(code, message); } diff --git a/tensorflow/compiler/xla/status_macros.h b/tensorflow/compiler/xla/status_macros.h index e51dd64e2a3dc7c359918cb08c6c94b2b4d9e91b..315136acc71670fa3ad48da4dc064e384ddadaa9 100644 --- a/tensorflow/compiler/xla/status_macros.h +++ b/tensorflow/compiler/xla/status_macros.h @@ -30,6 +30,10 @@ limitations under the License. namespace xla { namespace status_macros { +// This is a useful error message when encountering XLA Compiler errors that +// could be handled with the non-strict AutoJit mode. +extern const char kPossibleAutoJitAlternative[]; + // Stream object used to collect error messages in MAKE_ERROR macros // or append error messages with APPEND_ERROR. It accepts any // arguments with operator<< to build an error string, and then has an diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD index ee24d4d99cb1f7ce51a72c6258cbadd6adf12f81..db1c9274690583326b8a8d36413d725c14007aa3 100644 --- a/tensorflow/compiler/xla/tests/BUILD +++ b/tensorflow/compiler/xla/tests/BUILD @@ -71,6 +71,7 @@ cc_library( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_casting_utils", "//tensorflow/compiler/xla/service:hlo_dataflow_analysis", "//tensorflow/compiler/xla/service:hlo_verifier", "//tensorflow/compiler/xla/service:transfer_manager", @@ -276,9 +277,6 @@ cc_library( xla_test( name = "bad_rng_shape_validation_test", srcs = ["bad_rng_shape_validation_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:test", @@ -315,6 +313,31 @@ xla_test( ], ) +xla_test( + name = "conv_depthwise_backprop_filter_test", + timeout = "long", + srcs = ["conv_depthwise_backprop_filter_test.cc"], + # these backends do not natively handle batch group counts. + blacklisted_backends = [ + "gpu", + "cpu", + ], + shard_count = 6, + deps = [ + "//tensorflow/compiler/xla:execution_options_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla/client:xla_computation", + "//tensorflow/compiler/xla/service:bfloat16_normalization", + "//tensorflow/compiler/xla/service:despecializer", + "//tensorflow/compiler/xla/service:hlo_parser", + "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "@com_google_absl//absl/types:optional", + ], +) + xla_test( name = "grouped_convolution_test", timeout = "long", @@ -344,9 +367,6 @@ xla_test( xla_test( name = "check_execution_arity_test", srcs = ["check_execution_arity_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", @@ -367,9 +387,6 @@ xla_test( xla_test( name = "query_inferred_shape_test", srcs = ["query_inferred_shape_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", @@ -387,9 +404,6 @@ xla_test( xla_test( name = "while_test", srcs = ["while_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", @@ -413,6 +427,10 @@ xla_test( xla_test( name = "xla_hlo_profile_test", srcs = ["xla_hlo_profile_test.cc"], + blacklisted_backends = [ + # Hlo profiles are not supported on the interpreter backend. + "interpreter", + ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:shape_util", @@ -436,9 +454,6 @@ xla_test( xla_test( name = "axpy_simple_test", srcs = ["axpy_simple_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:xla_builder", @@ -453,7 +468,6 @@ xla_test( xla_test( name = "map_test", srcs = ["map_test.cc"], - tags = ["enable_for_xla_interpreter"], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:literal", @@ -506,9 +520,6 @@ xla_test( xla_test( name = "pred_test", srcs = ["pred_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla/client:local_client", @@ -524,9 +535,6 @@ xla_test( xla_test( name = "select_test", srcs = ["select_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:global_data", @@ -544,7 +552,6 @@ xla_test( xla_test( name = "conditional_test", srcs = ["conditional_test.cc"], - tags = ["enable_for_xla_interpreter"], deps = [ "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:global_data", @@ -562,7 +569,6 @@ xla_test( xla_test( name = "unary_op_test", srcs = ["unary_op_test.cc"], - tags = ["enable_for_xla_interpreter"], deps = [ "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:global_data", @@ -623,9 +629,6 @@ xla_test( xla_test( name = "deconstruct_tuple_test", srcs = ["deconstruct_tuple_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", @@ -648,7 +651,6 @@ xla_test( name = "array_elementwise_ops_test", srcs = ["array_elementwise_ops_test.cc"], shard_count = 25, - tags = ["enable_for_xla_interpreter"], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -673,22 +675,19 @@ xla_test( xla_test( name = "exhaustive_f32_elementwise_op_test", - size = "enormous", srcs = ["exhaustive_f32_elementwise_op_test.cc"], - backends = [ - "cpu", - "gpu", - ], + real_hardware_only = True, # Very slow on the interpreter. shard_count = 48, tags = [ - "broken", - "manual", - "notap", + "optonly", + # This is a big test that we skip for capacity reasons in OSS testing. + "nooss", ], deps = [ ":client_library_test_base", ":literal_test_util", "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/client/lib:math", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "@com_google_absl//absl/base", @@ -698,7 +697,6 @@ xla_test( xla_test( name = "reduce_precision_test", srcs = ["reduce_precision_test.cc"], - tags = ["enable_for_xla_interpreter"], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:literal", @@ -725,7 +723,6 @@ xla_test( srcs = ["dot_operation_test.cc"], shard_count = 20, tags = [ - "enable_for_xla_interpreter", "optonly", ], deps = [ @@ -735,7 +732,9 @@ xla_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/client/lib:matrix", "//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", @@ -792,7 +791,9 @@ xla_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/client/lib:matrix", "//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", @@ -806,9 +807,6 @@ xla_test( xla_test( name = "transpose_test", srcs = ["transpose_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:reference_util", @@ -828,9 +826,6 @@ xla_test( xla_test( name = "constants_test", srcs = ["constants_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -841,6 +836,7 @@ xla_test( "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:xla_builder", "//tensorflow/compiler/xla/client:xla_computation", + "//tensorflow/compiler/xla/client/lib:constants", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", @@ -951,6 +947,11 @@ xla_test( xla_test( name = "batch_normalization_test", srcs = ["batch_normalization_test.cc"], + blacklisted_backends = [ + # BatchNorm HLOs are not handled by the interpreter backend, and the + # BatchNorm expander is not run on the interpreter. + "interpreter", + ], shard_count = 40, deps = [ ":test_utils", @@ -1042,9 +1043,6 @@ xla_test( name = "slice_test", srcs = ["slice_test.cc"], shard_count = 40, - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:reference_util", @@ -1065,9 +1063,6 @@ xla_test( xla_test( name = "multidimensional_slice_test", srcs = ["multidimensional_slice_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -1085,9 +1080,6 @@ xla_test( name = "dynamic_ops_test", timeout = "moderate", srcs = ["dynamic_ops_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:reference_util", @@ -1113,9 +1105,6 @@ xla_test( xla_test( name = "tuple_test", srcs = ["tuple_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:literal", @@ -1139,9 +1128,6 @@ xla_test( xla_test( name = "vector_ops_reduce_test", srcs = ["vector_ops_reduce_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -1162,7 +1148,6 @@ xla_test( srcs = ["reduce_test.cc"], shard_count = 40, tags = [ - "enable_for_xla_interpreter", "optonly", ], deps = [ @@ -1229,7 +1214,6 @@ xla_test( srcs = [], shard_count = 20, tags = [ - "enable_for_xla_interpreter", "optonly", ], xla_test_library_deps = [":reduce_window_test_library"], @@ -1241,7 +1225,6 @@ xla_test( timeout = "long", srcs = ["select_and_scatter_test.cc"], tags = [ - "enable_for_xla_interpreter", "optonly", ], deps = [ @@ -1267,9 +1250,6 @@ xla_test( xla_test( name = "copy_test", srcs = ["copy_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ ":client_library_test_base", "//tensorflow/compiler/xla:array2d", @@ -1290,9 +1270,6 @@ xla_test( xla_test( name = "reduce_hlo_test", srcs = ["reduce_hlo_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla/service:hlo_parser", "//tensorflow/compiler/xla/tests:hlo_test_base", @@ -1306,9 +1283,6 @@ xla_test( xla_test( name = "token_hlo_test", srcs = ["token_hlo_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla/service:hlo_parser", "//tensorflow/compiler/xla/service:hlo_verifier", @@ -1323,9 +1297,6 @@ xla_test( xla_test( name = "call_test", srcs = ["call_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:literal_util", @@ -1368,9 +1339,6 @@ xla_test( xla_test( name = "binop_scaling_test", srcs = ["binop_scaling_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array4d", @@ -1388,9 +1356,6 @@ xla_test( xla_test( name = "broadcast_simple_test", srcs = ["broadcast_simple_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array4d", @@ -1410,9 +1375,6 @@ xla_test( xla_test( name = "pad_test", srcs = ["pad_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array4d", @@ -1434,9 +1396,6 @@ xla_test( xla_test( name = "fmax_test", srcs = ["fmax_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:xla_builder", @@ -1451,9 +1410,6 @@ xla_test( xla_test( name = "log_test", srcs = ["log_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:xla_builder", @@ -1468,9 +1424,6 @@ xla_test( xla_test( name = "matrix_ops_simple_test", srcs = ["matrix_ops_simple_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:literal", @@ -1517,9 +1470,6 @@ xla_test( name = "reshape_test", srcs = ["reshape_test.cc"], shard_count = 30, - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array4d", @@ -1545,9 +1495,6 @@ xla_test( xla_test( name = "reverse_test", srcs = ["reverse_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array4d", @@ -1566,9 +1513,6 @@ xla_test( xla_test( name = "vector_ops_simple_test", srcs = ["vector_ops_simple_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array4d", "//tensorflow/compiler/xla:shape_util", @@ -1592,9 +1536,6 @@ xla_test( xla_test( name = "concat_test", srcs = ["concat_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -1615,9 +1556,6 @@ xla_test( xla_test( name = "convert_test", srcs = ["convert_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:xla_data_proto", @@ -1637,6 +1575,10 @@ xla_test( xla_test( name = "all_reduce_test", srcs = ["all_reduce_test.cc"], + blacklisted_backends = [ + # All reduce is not supported on the interpreter backend. + "interpreter", + ], deps = [ "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", @@ -1661,9 +1603,6 @@ xla_test( xla_test( name = "bitcast_convert_test", srcs = ["bitcast_convert_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:xla_data_proto", @@ -1703,9 +1642,6 @@ xla_test( xla_test( name = "floor_ceil_test", srcs = ["floor_ceil_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:xla_builder", @@ -1767,6 +1703,10 @@ xla_test( xla_test( name = "execution_profile_test", srcs = ["execution_profile_test.cc"], + blacklisted_backends = [ + # Execution profiles are not supported on the interpreter backend. + "interpreter", + ], deps = [ ":client_library_test_base", "//tensorflow/compiler/xla/client:global_data", @@ -1781,6 +1721,10 @@ xla_test( name = "execution_profile_test_with_xla_hlo_profile", srcs = ["execution_profile_test.cc"], args = ["--xla_hlo_profile"], + blacklisted_backends = [ + # Hlo profiles are not supported on the interpreter backend. + "interpreter", + ], deps = [ ":client_library_test_base", "//tensorflow/compiler/xla/client:global_data", @@ -1794,9 +1738,6 @@ xla_test( xla_test( name = "replay_test", srcs = ["replay_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:protobuf_util", @@ -1819,9 +1760,6 @@ xla_test( xla_test( name = "broadcast_test", srcs = ["broadcast_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", @@ -1883,9 +1821,6 @@ xla_test( xla_test( name = "fusion_test", srcs = ["fusion_test.cc"], - tags = [ - "enable_for_xla_interpreter", - ], deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:literal", @@ -2003,6 +1938,10 @@ xla_test( xla_test( name = "outfeed_in_nested_computation_test", srcs = ["outfeed_in_nested_computation_test.cc"], + blacklisted_backends = [ + # Outfeed ops are not supported on the interpreter backend. + "interpreter", + ], deps = [ "//tensorflow/compiler/xla/tests:local_client_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", @@ -2179,7 +2118,6 @@ xla_test( srcs = ["iota_test.cc"], shard_count = 30, tags = [ - "enable_for_xla_interpreter", # Require optimized builds, iota_test_cpu is very slow in fastbuild. "optonly", ], @@ -2207,3 +2145,41 @@ tf_cc_test( "@com_google_absl//absl/synchronization", ], ) + +xla_test( + name = "ptxas_bug_120501638", + srcs = ["ptxas_bug_120501638.cc"], + tags = [ + # Disabled in OSS until nvidia publicly releases a fixed ptxas. + "no_oss", + ], + deps = [ + ":hlo_test_base", + ":xla_internal_test_main", # fixdeps: keep + "//tensorflow/compiler/xla:debug_options_flags", + "//tensorflow/compiler/xla:test", + ], +) + +xla_test( + name = "triangular_solve_test", + srcs = ["triangular_solve_test.cc"], + tags = [ + "enable_for_xla_interpreter", + "noasan", # sometimes times out, http://b/78650012 + ], + 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:math", + "//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/tests/array_elementwise_ops_test.cc b/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc index 915b456b52215f8d6a9eb6c5b933f3502f1d3d2c..acdd3c9da92efe8fae1336eaa861c01d5bb9b158 100644 --- a/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc +++ b/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc @@ -35,7 +35,6 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/compiler/xla/tests/test_macros.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/types.h" namespace xla { @@ -1443,6 +1442,27 @@ XLA_TEST_F(ArrayElementwiseOpTest, PowNonIntegerF32s) { error_spec_); } +XLA_TEST_F(ArrayElementwiseOpTest, PowC64s) { + SetFastMathDisabled(true); + XlaBuilder builder(TestName()); + auto lhs = + ConstantR1(&builder, {-2.0f, -0.6f, -0.6f, 0.0f, 0.0f, 0.0f}); + auto rhs = + ConstantR1(&builder, {0.5f, 0.6f, -0.6f, 0.5f, 0.6f, 0.0f}); + Pow(lhs, rhs); + + ComputeAndCompareR1(&builder, + { + {0, 1.41421356}, + {-2.27443288e-01, 0.69999846}, + {-4.19847531e-01, -1.29215783}, + {0, 0}, + {0, 0}, + {1, 0}, + }, + {}, error_spec_); +} + XLA_TEST_F(ArrayElementwiseOpTest, PowZeroElementF32s) { XlaBuilder builder(TestName()); auto lhs = ConstantR1(&builder, {}); @@ -2047,6 +2067,19 @@ XLA_TEST_F(ArrayElementwiseOpTest, NonNanClampF32) { error_spec_); } +XLA_TEST_F(ArrayElementwiseOpTest, ClampF32) { + SetFastMathDisabled(true); + XlaBuilder builder(TestName()); + auto minimum = ConstantR1(&builder, {1.0f, -6.5f, 1.0f, 2.25f, NAN}); + auto argument = + ConstantR1(&builder, {2.0f, 10.0f, -5.0f, 1.0f, 10.0f}); + auto maximum = ConstantR1(&builder, {3.0f, 0.5f, 25.5f, NAN, 123.0f}); + Clamp(minimum, argument, maximum); + + ComputeAndCompareR1(&builder, {2.0f, 0.5f, 1.0f, NAN, NAN}, {}, + error_spec_); +} + XLA_TEST_F(ArrayElementwiseOpTest, ClampF32Scalar) { XlaBuilder builder(TestName()); auto minimum = ConstantR0(&builder, 0.0f); diff --git a/tensorflow/compiler/xla/tests/bfloat16_test.cc b/tensorflow/compiler/xla/tests/bfloat16_test.cc index e9728e636f0ee032416b2da17a3ea83c5bb18083..63e48117056dec4af603cbc85e478fcb15ad0cec 100644 --- a/tensorflow/compiler/xla/tests/bfloat16_test.cc +++ b/tensorflow/compiler/xla/tests/bfloat16_test.cc @@ -76,7 +76,9 @@ XLA_TEST_F(Bfloat16Test, NegateScalarF16) { error_spec_); } -XLA_TEST_F(Bfloat16Test, BatchNormTraining) { +// Disabled on interpreter since BatchNormExanper is not run by default on the +// intepreter backend. +XLA_TEST_F(Bfloat16Test, DISABLED_ON_INTERPRETER(BatchNormTraining)) { const int kFeatureIndex = 2; XlaBuilder builder(TestName()); @@ -110,7 +112,9 @@ XLA_TEST_F(Bfloat16Test, BatchNormTraining) { ComputeAndCompareTuple(&builder, expected, {}, ErrorSpec(0.01, 0.02)); } -XLA_TEST_F(Bfloat16Test, BatchNormGrad) { +// Disabled on interpreter since BatchNormExanper is not run by default on the +// intepreter backend. +XLA_TEST_F(Bfloat16Test, DISABLED_ON_INTERPRETER(BatchNormGrad)) { const int kFeatureIndex = 2; XlaBuilder builder(TestName()); diff --git a/tensorflow/compiler/xla/tests/build_defs.bzl b/tensorflow/compiler/xla/tests/build_defs.bzl index 05d4d04034bf50c8bb840e59b28a590fce048c19..c14d279ac560db33066ae4fc68b6290f7499bb39 100644 --- a/tensorflow/compiler/xla/tests/build_defs.bzl +++ b/tensorflow/compiler/xla/tests/build_defs.bzl @@ -34,6 +34,7 @@ def xla_test( xla_test_library_deps = [], backends = [], blacklisted_backends = [], + real_hardware_only = False, args = [], tags = [], copts = [], @@ -108,6 +109,10 @@ def xla_test( use for that target. **kwargs: Additional keyword arguments to pass to native.cc_test. """ + + # All of the backends in all_backends are real hardware. + _ignore = [real_hardware_only] + test_names = [] if not backends: backends = all_backends diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.cc b/tensorflow/compiler/xla/tests/client_library_test_base.cc index edb95c973b70e30702ed8490c15a48d4d5604170..0e99ede5d01fcfa88c54c9cbc5a6a85bf8f15ddf 100644 --- a/tensorflow/compiler/xla/tests/client_library_test_base.cc +++ b/tensorflow/compiler/xla/tests/client_library_test_base.cc @@ -41,8 +41,9 @@ constexpr char kInterpreter[] = "interpreter"; // Wrapper function that creates a nicer error message (than a bare // ValueOrDie()) if the platform we intend to test is not available. -Client* GetOrCreateLocalClientOrDie(const LocalClientOptions& client_options) { - StatusOr result = +LocalClient* GetOrCreateLocalClientOrDie( + const LocalClientOptions& client_options) { + StatusOr result = ClientLibrary::GetOrCreateLocalClient(client_options); TF_CHECK_OK(result.status()) << " could not create local client for testing"; return result.ValueOrDie(); diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.h b/tensorflow/compiler/xla/tests/client_library_test_base.h index 65a23dd883594b9bf9c37494a37e9be39b197788..d700437ed355c144639f76d683055e211975fde9 100644 --- a/tensorflow/compiler/xla/tests/client_library_test_base.h +++ b/tensorflow/compiler/xla/tests/client_library_test_base.h @@ -385,8 +385,8 @@ class ClientLibraryTestBase : public ::testing::Test { StatusOr> ComputeValueAndReference( XlaBuilder* builder, absl::Span arguments); - Client* client_; - Client* ref_client_; // To compute reference result. + LocalClient* client_; + LocalClient* ref_client_; // To compute reference result. ExecutionOptions execution_options_; private: @@ -431,7 +431,8 @@ void ClientLibraryTestBase::ComputeAndCompareR0( std::is_same::value || std::is_same::value || std::is_same::value || - std::is_same::value, + std::is_same::value || + std::is_same::value, "Float or complex type required when specifying an ErrorSpec"); Literal expected_literal = LiteralUtil::CreateR0(expected); ClientLibraryTestBase::ComputeAndCompareLiteral(builder, expected_literal, @@ -455,7 +456,8 @@ void ClientLibraryTestBase::ComputeAndCompareR1( std::is_same::value || std::is_same::value || std::is_same::value || - std::is_same::value, + std::is_same::value || + std::is_same::value, "Float or complex type required when specifying an ErrorSpec"); Literal expected_literal = LiteralUtil::CreateR1(expected); ClientLibraryTestBase::ComputeAndCompareLiteral(builder, expected_literal, @@ -480,7 +482,8 @@ void ClientLibraryTestBase::ComputeAndCompareR2( std::is_same::value || std::is_same::value || std::is_same::value || - std::is_same::value, + std::is_same::value || + std::is_same::value, "Float or complex type required when specifying an ErrorSpec"); Literal expected_literal = LiteralUtil::CreateR2FromArray2D(expected); @@ -506,7 +509,8 @@ void ClientLibraryTestBase::ComputeAndCompareR3( std::is_same::value || std::is_same::value || std::is_same::value || - std::is_same::value, + std::is_same::value || + std::is_same::value, "Float or complex type required when specifying an ErrorSpec"); Literal expected_literal = LiteralUtil::CreateR3FromArray3D(expected); @@ -532,7 +536,8 @@ void ClientLibraryTestBase::ComputeAndCompareR4( std::is_same::value || std::is_same::value || std::is_same::value || - std::is_same::value, + std::is_same::value || + std::is_same::value, "Float or complex type required when specifying an ErrorSpec"); Literal expected_literal = LiteralUtil::CreateR4FromArray4D(expected); diff --git a/tensorflow/compiler/xla/tests/client_test.cc b/tensorflow/compiler/xla/tests/client_test.cc index d5b3a4d14f932d76b2e7b198ee233756c94a2694..247328b730f3af936d933f824da491b593b27c90 100644 --- a/tensorflow/compiler/xla/tests/client_test.cc +++ b/tensorflow/compiler/xla/tests/client_test.cc @@ -109,7 +109,10 @@ XLA_TEST_F(ClientTest, ExecuteWithTupleLayout) { /*minor_to_major=*/{1, 0}))); } -XLA_TEST_F(ClientTest, DISABLED_ON_GPU(ExecuteParallel)) { +// Disabled for interpreter since ExecuteAsyncOnStream is not implemented on +// interpreter backend. +XLA_TEST_F(ClientTest, + DISABLED_ON_INTERPRETER(DISABLED_ON_GPU(ExecuteParallel))) { XlaComputation add_with_one_arg, mul_with_two_args, dot_with_one_arg; Shape shape = ShapeUtil::MakeShape(S32, {2, 2}); diff --git a/tensorflow/compiler/xla/tests/compute_constant_test.cc b/tensorflow/compiler/xla/tests/compute_constant_test.cc index 3b0414a6045a7c5f4f75948d8ccf2775c575626e..ef800b8ef624bf1020ff1e6857c13b0387482cd3 100644 --- a/tensorflow/compiler/xla/tests/compute_constant_test.cc +++ b/tensorflow/compiler/xla/tests/compute_constant_test.cc @@ -151,19 +151,35 @@ TEST_F(ComputeConstantTest, DirectParamMissing) { } } -TEST_F(ComputeConstantTest, IndirectParamMissing) { +TEST_F(ComputeConstantTest, GetDimensionSize) { for (ClientType client_type : client_types) { Client* client = ClientOrDie(platform_, client_type); XlaBuilder b(TestName()); - auto computation = - Add(ConstantR0(&b, 1.0f), - Parameter(&b, 0, ShapeUtil::MakeShape(F32, {}), "param")); - EXPECT_FALSE(IsConstant(computation, &b)); + auto add = + Add(ConstantR1(&b, {1.0f}), ConstantR1(&b, {1.0f})); + auto get_dimension_size = GetDimensionSize(add, 0); + EXPECT_TRUE(IsConstant(get_dimension_size, &b)); + + TF_ASSERT_OK_AND_ASSIGN(auto value, ComputeConstantScalar( + client, get_dimension_size, &b)); + EXPECT_EQ(value, 1); + } +} - auto value = ComputeConstantScalar(client, computation, &b); - EXPECT_TRUE( - absl::StrContains(value.status().ToString(), "depends on a parameter")) - << value.status(); +TEST_F(ComputeConstantTest, MultipleGetDimensionSize) { + for (ClientType client_type : client_types) { + Client* client = ClientOrDie(platform_, client_type); + XlaBuilder b(TestName()); + auto add = + Add(ConstantR2(&b, {{1.0f}}), ConstantR2(&b, {{1.0f}})); + auto get_dimension_size = GetDimensionSize(add, 0); + auto get_dimension_size_2 = GetDimensionSize(add, 0); + auto add_2 = Add(get_dimension_size, get_dimension_size_2); + EXPECT_TRUE(IsConstant(add_2, &b)); + + TF_ASSERT_OK_AND_ASSIGN(auto value, + ComputeConstantScalar(client, add_2, &b)); + EXPECT_EQ(value, 2); } } diff --git a/tensorflow/compiler/xla/tests/constants_test.cc b/tensorflow/compiler/xla/tests/constants_test.cc index 9174f2651cb90b364f869364fe108cf208c11a84..6530007871ced1d0bbffe2b44ccc8cf9bddd79e1 100644 --- a/tensorflow/compiler/xla/tests/constants_test.cc +++ b/tensorflow/compiler/xla/tests/constants_test.cc @@ -21,6 +21,7 @@ limitations under the License. #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/array3d.h" #include "tensorflow/compiler/xla/array4d.h" +#include "tensorflow/compiler/xla/client/lib/constants.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal_util.h" @@ -180,6 +181,29 @@ TEST_F(ConstantsTest, Token) { TF_ASSERT_OK(Execute(&builder, {}).status()); } +TEST_F(ConstantsTest, FullLike) { + XlaBuilder b(TestName()); + auto val1 = Iota(&b, F32, 3); + auto val2 = FullLike(val1, 10); + val1 + val2; + ComputeAndCompareR1(&b, {10, 11, 12}, {}, error_spec_); +} + +TEST_F(ConstantsTest, IllegalFullLikeOnTuple) { + XlaBuilder b(TestName()); + auto tuple = Tuple(&b, {Iota(&b, F32, 3), Iota(&b, F32, 1)}); + FullLike(tuple, 10); // Illegal; can't do FullLike on a tuple. + EXPECT_FALSE(b.Build().ok()); +} + +TEST_F(ConstantsTest, FullLikeScalar) { + XlaBuilder b(TestName()); + auto scalar1 = ConstantR0WithType(&b, F32, 1); + auto scalar2 = FullLike(scalar1, 2); + scalar1 - scalar2; + ComputeAndCompareR0(&b, -1, {}, error_spec_); +} + class ConstantsHloTest : public HloTestBase {}; // TODO(b/121147351): Fails on GPU. Not clear if this is expected behavior. @@ -200,9 +224,7 @@ XLA_TEST_F(ConstantsHloTest, DISABLED_ON_GPU(BitcastOfConstant)) { ROOT result = s32[] call(parameter.0, constant-as-scalar), to_apply=func } )"; - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param = LiteralUtil::CreateR0(1); auto result = ExecuteNoHloPasses(std::move(module), {¶m}); EXPECT_TRUE(LiteralTestUtil::Equal(param, result)); diff --git a/tensorflow/compiler/xla/tests/conv_depthwise_backprop_filter_test.cc b/tensorflow/compiler/xla/tests/conv_depthwise_backprop_filter_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..dfbf0478e62713635446d11557367cfac6ab0dce --- /dev/null +++ b/tensorflow/compiler/xla/tests/conv_depthwise_backprop_filter_test.cc @@ -0,0 +1,178 @@ +/* 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 "absl/types/optional.h" +#include "tensorflow/compiler/xla/client/xla_computation.h" +#include "tensorflow/compiler/xla/execution_options_util.h" +#include "tensorflow/compiler/xla/service/bfloat16_normalization.h" +#include "tensorflow/compiler/xla/service/despecializer.h" +#include "tensorflow/compiler/xla/service/hlo_parser.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/tests/client_library_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/tests/test_macros.h" + +namespace xla { +namespace { + +string GetFloatDataType(bool use_bfloat16) { + return use_bfloat16 ? "bf16" : "f32"; +} + +struct BatchGroupedConvolution2DSpec { + int64 output_batch, window, window_dilation; + std::vector activation_dims; + std::vector kernel_dims; + std::vector output_dims; + std::vector activation_and_kernel_layout; + std::vector output_layout; +}; + +class BatchGroupedConvolution2DTest + : public HloTestBase, + public ::testing::WithParamInterface< + ::testing::tuple> {}; + +static std::vector GetConv2DTestCases() { + std::vector config_set; + std::vector> config_options = { + {8, 5, 3, 2}, {4, 5, 5, 2}, {8, 7, 4, 128}, + {16, 20, 20, 256}, {256, 7, 5, 4}, {256, 6, 6, 4}, + {256, 8, 8, 512}, {64, 7, 7, 960}, {64, 14, 14, 576}}; + + for (auto option : config_options) { + int64 feature = option[3]; + int64 activation_size = option[1]; + int64 kernel_size = option[2]; + int64 batch = option[0]; + + BatchGroupedConvolution2DSpec config; + config.window_dilation = 1; + config.output_batch = feature; + config.window = kernel_size; + + config.activation_dims = {batch, activation_size, activation_size, feature}; + + config.kernel_dims = {batch, kernel_size, kernel_size, feature}; + + int64 output_space_size = 3 + activation_size - kernel_size; + config.output_dims = {output_space_size, output_space_size, feature, 1}; + + config.activation_and_kernel_layout = {0, 3, 1, 2}; + config.output_layout = {2, 3, 0, 1}; + config_set.push_back(config); + + BatchGroupedConvolution2DSpec different_layout_config = config; + different_layout_config.activation_and_kernel_layout = {3, 0, 1, 2}; + config_set.push_back(different_layout_config); + + // Add configurations for window dilation cases. + if (activation_size % 2 == 0 && activation_size == kernel_size) { + BatchGroupedConvolution2DSpec config; + config.window_dilation = 2; + config.output_batch = feature; + config.window = kernel_size / 2; + config.activation_dims = {batch, activation_size, activation_size, + feature}; + config.kernel_dims = {batch, kernel_size / 2, kernel_size / 2, feature}; + config.activation_and_kernel_layout = {0, 3, 1, 2}; + config.output_layout = {2, 3, 0, 1}; + + int64 output_space_size = 5; + config.output_dims = {output_space_size, output_space_size, feature, 1}; + + config_set.push_back(config); + + BatchGroupedConvolution2DSpec different_layout_config = config; + different_layout_config.activation_and_kernel_layout = {3, 0, 1, 2}; + config_set.push_back(different_layout_config); + } + } + + return config_set; +} + +string BatchGroupedConvolution2DTestDataToString( + const ::testing::TestParamInfo< + ::testing::tuple>& data) { + const auto& spec = ::testing::get<0>(data.param); + const string data_type = GetFloatDataType(::testing::get<1>(data.param)); + string str = absl::StrCat( + "activation_dims_", absl::StrJoin(spec.activation_dims, "x"), + "_kernel_dims_", absl::StrJoin(spec.kernel_dims, "x"), + "_activation_layout_", + absl::StrJoin(spec.activation_and_kernel_layout, "_"), "_output_dims_", + absl::StrJoin(spec.output_dims, "x"), data_type, "_output_layout_", + absl::StrJoin(spec.output_layout, "_")); + + // Test names are not allowed to contain the '-' character. + absl::c_replace(str, '-', 'n'); + return str; +} + +string BuildHloTextBatchGroupedConvolution2D( + const BatchGroupedConvolution2DSpec& spec, bool use_bfloat16) { + const string data_type = GetFloatDataType(use_bfloat16); + return absl::StrFormat( + R"( + HloModule TensorFlowDepthwiseConv, is_scheduled=true + + ENTRY main { + activation = %s[%s]{%s} parameter(0) + kernel = %s[%s]{%s} parameter(1) + ROOT conv = %s[%s]{%s} convolution(%s[%s]{%s} activation, %s[%s]{%s} kernel), + window={size=%dx%d pad=1_%dx1_%d rhs_dilate=%dx%d}, dim_labels=f01b_i01o->01fb, + batch_group_count=%d + } + )", + data_type, absl::StrJoin(spec.activation_dims, ","), + absl::StrJoin(spec.activation_and_kernel_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.activation_and_kernel_layout, ","), data_type, + absl::StrJoin(spec.output_dims, ","), + absl::StrJoin(spec.output_layout, ","), data_type, + absl::StrJoin(spec.activation_dims, ","), + absl::StrJoin(spec.activation_and_kernel_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.activation_and_kernel_layout, ","), spec.window, + spec.window, spec.window_dilation, spec.window_dilation, + spec.window_dilation, spec.window_dilation, spec.output_batch); +} + +XLA_TEST_P(BatchGroupedConvolution2DTest, DoIt) { + const BatchGroupedConvolution2DSpec& spec = ::testing::get<0>(GetParam()); + bool use_bfloat16 = ::testing::get<1>(GetParam()); + const string hlo_text = + BuildHloTextBatchGroupedConvolution2D(spec, use_bfloat16); + + EXPECT_TRUE(RunAndCompareNoHloPasses( + hlo_text, ErrorSpec{0.01, 0.01}, [](HloModule* module) -> Status { + BFloat16MixedPrecisionRemoval remover; + TF_RETURN_IF_ERROR(remover.Run(module).status()); + Despecializer despecializer; + return despecializer.Run(module).status(); + })); +} + +INSTANTIATE_TEST_CASE_P( + BatchGroupedConvolution2DTestWithRandomIndices, + BatchGroupedConvolution2DTest, + ::testing::Combine(::testing::ValuesIn(GetConv2DTestCases()), + ::testing::Bool()), + BatchGroupedConvolution2DTestDataToString); + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/convolution_test.cc b/tensorflow/compiler/xla/tests/convolution_test.cc index 249693891290e14645ee5b4b4d97b2d506a01302..9db9f2563b636c4f929585eb13a9c7f809833eda 100644 --- a/tensorflow/compiler/xla/tests/convolution_test.cc +++ b/tensorflow/compiler/xla/tests/convolution_test.cc @@ -467,8 +467,8 @@ XLA_TEST_F(ConvolutionTest, Convolve3D_1x4x2x3x3_2x2x2x3x3_Valid) { // servers. The error message is missing the operator ++. template void iota_int_init_value(std::vector& values, int init_value) { - std::for_each(values.begin(), values.end(), - [&](T& value) { value = static_cast(init_value++); }); + absl::c_for_each(values, + [&](T& value) { value = static_cast(init_value++); }); } template diff --git a/tensorflow/compiler/xla/tests/dot_operation_test.cc b/tensorflow/compiler/xla/tests/dot_operation_test.cc index c5d8b663f4abe77e05ec213d2e4e075c260a8655..11343ddfd0b7f67f11704b8ef42889b7325cda9f 100644 --- a/tensorflow/compiler/xla/tests/dot_operation_test.cc +++ b/tensorflow/compiler/xla/tests/dot_operation_test.cc @@ -19,12 +19,14 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/array3d.h" +#include "tensorflow/compiler/xla/client/lib/matrix.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/compiler/xla/tests/test_macros.h" #include "tensorflow/compiler/xla/tests/test_utils.h" @@ -918,8 +920,9 @@ XLA_TEST_F(DotOperationTest, DotOfGatherOptimizationWithConstRHSClassicMM) { XlaBuilder builder(TestName()); auto lhs_constant = ConstantR2FromArray2D(&builder, *constant_lhs_array); auto rhs_constant = ConstantR2FromArray2D(&builder, *constant_rhs_array); - auto start_constant = ConstantR1(&builder, {1, 0}); - auto dynamic_slice = DynamicSlice(lhs_constant, start_constant, {1, 6}); + auto one = ConstantR0(&builder, 1); + auto zero = ConstantR0(&builder, 0); + auto dynamic_slice = DynamicSlice(lhs_constant, {one, zero}, {1, 6}); DotDimensionNumbers dot_dnums; dot_dnums.add_lhs_contracting_dimensions(1); @@ -945,8 +948,9 @@ XLA_TEST_F(DotOperationTest, DotOfGatherOptimizationWithConstLHSClassicMM) { XlaBuilder builder(TestName()); auto lhs_constant = ConstantR2FromArray2D(&builder, *constant_lhs_array); auto rhs_constant = ConstantR2FromArray2D(&builder, *constant_rhs_array); - auto start_constant = ConstantR1(&builder, {0, 1}); - auto dynamic_slice = DynamicSlice(rhs_constant, start_constant, {6, 1}); + auto zero = ConstantR0(&builder, 0); + auto one = ConstantR0(&builder, 1); + auto dynamic_slice = DynamicSlice(rhs_constant, {zero, one}, {6, 1}); DotDimensionNumbers dot_dnums; dot_dnums.add_lhs_contracting_dimensions(1); @@ -974,8 +978,9 @@ XLA_TEST_F(DotOperationTest, XlaBuilder builder(TestName()); auto lhs_constant = ConstantR2FromArray2D(&builder, *constant_lhs_array); auto rhs_constant = ConstantR2FromArray2D(&builder, *constant_rhs_array); - auto start_constant = ConstantR1(&builder, {0, 1}); - auto dynamic_slice = DynamicSlice(lhs_constant, start_constant, {6, 1}); + auto zero = ConstantR0(&builder, 0); + auto one = ConstantR0(&builder, 1); + auto dynamic_slice = DynamicSlice(lhs_constant, {zero, one}, {6, 1}); DotDimensionNumbers dot_dnums; dot_dnums.add_lhs_contracting_dimensions(0); @@ -1001,8 +1006,9 @@ XLA_TEST_F(DotOperationTest, DotOfGatherOptimizationWithConstLHSReverseMM) { XlaBuilder builder(TestName()); auto lhs_constant = ConstantR2FromArray2D(&builder, *constant_lhs_array); auto rhs_constant = ConstantR2FromArray2D(&builder, *constant_rhs_array); - auto start_constant = ConstantR1(&builder, {1, 0}); - auto dynamic_slice = DynamicSlice(rhs_constant, start_constant, {1, 6}); + auto zero = ConstantR0(&builder, 0); + auto one = ConstantR0(&builder, 1); + auto dynamic_slice = DynamicSlice(rhs_constant, {one, zero}, {1, 6}); DotDimensionNumbers dot_dnums; dot_dnums.add_lhs_contracting_dimensions(0); @@ -1033,8 +1039,9 @@ XLA_TEST_F(DotOperationTest, DotOfGatherOptimizationWithConstRHSRows) { XlaBuilder builder(TestName()); auto lhs_constant = ConstantR2FromArray2D(&builder, *constant_lhs_array); auto rhs_constant = ConstantR2FromArray2D(&builder, *constant_rhs_array); - auto start_constant = ConstantR1(&builder, {0, 1}); - auto dynamic_slice = DynamicSlice(lhs_constant, start_constant, {6, 1}); + auto zero = ConstantR0(&builder, 0); + auto one = ConstantR0(&builder, 1); + auto dynamic_slice = DynamicSlice(lhs_constant, {zero, one}, {6, 1}); DotDimensionNumbers dot_dnums; dot_dnums.add_lhs_contracting_dimensions(0); @@ -1065,8 +1072,9 @@ XLA_TEST_F(DotOperationTest, DotOfGatherOptimizationWithConstLHSRows) { XlaBuilder builder(TestName()); auto lhs_constant = ConstantR2FromArray2D(&builder, *constant_lhs_array); auto rhs_constant = ConstantR2FromArray2D(&builder, *constant_rhs_array); - auto start_constant = ConstantR1(&builder, {0, 1}); - auto dynamic_slice = DynamicSlice(rhs_constant, start_constant, {6, 1}); + auto zero = ConstantR0(&builder, 0); + auto one = ConstantR0(&builder, 1); + auto dynamic_slice = DynamicSlice(rhs_constant, {zero, one}, {6, 1}); DotDimensionNumbers dot_dnums; dot_dnums.add_lhs_contracting_dimensions(0); @@ -1089,8 +1097,9 @@ XLA_TEST_F(DotOperationTest, DotOfGatherOptimizationWithConstRHSCols) { XlaBuilder builder(TestName()); auto lhs_constant = ConstantR2FromArray2D(&builder, *constant_lhs_array); auto rhs_constant = ConstantR2FromArray2D(&builder, *constant_rhs_array); - auto start_constant = ConstantR1(&builder, {1, 0}); - auto dynamic_slice = DynamicSlice(lhs_constant, start_constant, {1, 6}); + auto zero = ConstantR0(&builder, 0); + auto one = ConstantR0(&builder, 1); + auto dynamic_slice = DynamicSlice(lhs_constant, {one, zero}, {1, 6}); DotDimensionNumbers dot_dnums; dot_dnums.add_lhs_contracting_dimensions(1); @@ -1113,8 +1122,9 @@ XLA_TEST_F(DotOperationTest, DotOfGatherOptimizationWithConstLHSCols) { XlaBuilder builder(TestName()); auto lhs_constant = ConstantR2FromArray2D(&builder, *constant_lhs_array); auto rhs_constant = ConstantR2FromArray2D(&builder, *constant_rhs_array); - auto start_constant = ConstantR1(&builder, {1, 0}); - auto dynamic_slice = DynamicSlice(rhs_constant, start_constant, {1, 6}); + auto zero = ConstantR0(&builder, 0); + auto one = ConstantR0(&builder, 1); + auto dynamic_slice = DynamicSlice(rhs_constant, {one, zero}, {1, 6}); DotDimensionNumbers dot_dnums; dot_dnums.add_lhs_contracting_dimensions(1); @@ -1147,5 +1157,113 @@ XLA_TEST_F(DotOperationTest, DotRank2AndRank2NonDefaultContractionDims) { ComputeAndCompareR2(&builder, expected, {}, error_spec_); } + +using EinsumParamType = + std::tuple, std::vector, string>; +class EinsumTest : public DotOperationTest, + public ::testing::WithParamInterface {}; +XLA_TEST_P(EinsumTest, SimpleEinsumTest) { + XlaBuilder builder(TestName()); + auto x = AddParam( + MakeFakeLiteral(ShapeUtil::MakeShape(F32, std::get<0>(GetParam()))) + .ValueOrDie(), + &builder); + auto y = AddParam( + MakeFakeLiteral(ShapeUtil::MakeShape(F32, std::get<1>(GetParam()))) + .ValueOrDie(), + &builder); + Einsum(x, y, std::get<2>(GetParam())); + ComputeAndCompare(&builder, {}, ErrorSpec{1e-3, 1e-3}); +} + +std::vector GetEinsumTestCases() { + using v = std::vector; + using p = EinsumParamType; + std::vector

test_cases = { + p{v{5, 6}, v{6, 7}, "mk,kn->mn"}, + p{v{5, 6}, v{6, 7}, "mk,kn->nm"}, + p{v{5, 6, 11}, v{6, 11, 7}, "mkB,kBn->nmB"}, + p{v{31, 55, 11}, v{55, 11, 29}, "mkB,kBn->nmB"}, + p{v{31, 55, 11}, v{55, 11, 29}, "mkB,kBn->Bnm"}, + p{v{8, 55, 11, 3}, v{55, 11, 3, 29}, "mkBC,kBCn->BCnm"}, + p{v{5, 6}, v{6, 7}, "ab,cd->dcba"}, + p{v{6}, v{6, 7}, "b,bc->c"}, + p{v{77}, v{77}, "a,a->a"}, + p{v{77}, v{77, 55}, "a,ab->ba"}, + p{v{2, 3, 77}, v{77, 2, 3, 55}, "ija,aijb->baij"}, + p{v{55}, v{}, "a,->a"}, + p{v{11, 111}, v{11}, "ab,a->ab"}, + p{v{16, 34}, v{16, 34}, "ab,ab->ab"}, + p{v{16, 3, 34}, v{3, 16, 34}, "abc,bac->abc"}, + p{v{5, 19}, v{}, "ab,->ab"}, + }; + return test_cases; +} + +INSTANTIATE_TEST_CASE_P(Einsum, EinsumTest, + ::testing::ValuesIn(GetEinsumTestCases())); + +class DotOperationTextTest : public HloTestBase {}; + +XLA_TEST_F(DotOperationTextTest, DotReorderedDotDims) { + absl::string_view hlo_string = + R"( +HloModule ComplexDotMultipleNonContracting + +ENTRY %test { + %lhs = f32[7,17,10,13]{3,2,1,0} parameter(0) + %rhs = f32[7,9,10,13,6]{4,3,2,1,0} parameter(1) + ROOT %dot = f32[10,7,17,9,6]{4,3,2,1,0} dot(%lhs, %rhs), lhs_batch_dims={2,0}, rhs_batch_dims={2,0}, lhs_contracting_dims={3}, rhs_contracting_dims={3} +} +)"; + + EXPECT_TRUE(RunAndCompare(hlo_string, ErrorSpec{1e-3, 1e-3})); +} + +XLA_TEST_F(DotOperationTextTest, DotReorderedDotDimsAndMultipleContracting) { + absl::string_view hlo_string = + R"( +HloModule ComplexDotMultipleNonContracting + +ENTRY %test { + %lhs = f32[7,5,17,10,13]{4,3,2,1,0} parameter(0) + %rhs = f32[7,9,10,13,6,5]{5,4,3,2,1,0} parameter(1) + ROOT %dot = f32[10,7,17,9,6]{4,3,2,1,0} dot(%lhs, %rhs), lhs_batch_dims={3,0}, rhs_batch_dims={2,0}, lhs_contracting_dims={1,4}, rhs_contracting_dims={5,3} +} +)"; + + EXPECT_TRUE(RunAndCompare(hlo_string, ErrorSpec{1e-3, 1e-3})); +} + +XLA_TEST_F(DotOperationTextTest, DotWithNoDnums) { + absl::string_view hlo_string = + R"( +HloModule DotWithNoDnums + +ENTRY %test { + %lhs = f32[2,3]{1,0} parameter(0) + %rhs = f32[4,5]{1,0} parameter(1) + ROOT %dot = f32[2,3,4,5]{3,2,1,0} dot(%lhs, %rhs) +} +)"; + + EXPECT_TRUE(RunAndCompare(hlo_string, ErrorSpec{1e-3, 1e-3})); +} + +XLA_TEST_F(DotOperationTextTest, Einsum) { + absl::string_view hlo_string = + R"( +HloModule Einsum + +ENTRY %test { + %lhs = f32[8,64,96]{2,1,0} parameter(0) + %rhs = f32[96,32,4]{2,1,0} parameter(1) + ROOT %dot = f32[8,64,32,4]{3,2,1,0} dot(%lhs, %rhs), lhs_contracting_dims={2}, rhs_contracting_dims={0} +} +)"; + + EXPECT_TRUE(RunAndCompare(hlo_string, ErrorSpec{4e-3, 4e-3})); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/dynamic_ops_test.cc b/tensorflow/compiler/xla/tests/dynamic_ops_test.cc index 7501c6d957e7afe99b8c530e5f0d575f818367da..82e2db36143b2552472fedae701f32389a9be108 100644 --- a/tensorflow/compiler/xla/tests/dynamic_ops_test.cc +++ b/tensorflow/compiler/xla/tests/dynamic_ops_test.cc @@ -135,11 +135,11 @@ class DynamicSliceTest : public ClientLibraryTestBase { XlaBuilder builder(TestName()); // Initialize and transfer dynamic slice start indices parameter. XlaOp starts; - std::unique_ptr start_data = CreateR1Parameter( - slice_starts, 0, "slice_starts", &builder, &starts); + std::unique_ptr start_data = CreateR0Parameter( + slice_starts[0], 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. auto input = ConstantLiteral(&builder, input_values); - DynamicSlice(input, starts, slice_sizes); + DynamicSlice(input, absl::Span({starts}), slice_sizes); // Run computation and compare against expected values. ComputeAndCompareLiteral(&builder, expected_values, {start_data.get()}); } @@ -160,14 +160,23 @@ class DynamicSliceTest : public ClientLibraryTestBase { XlaBuilder builder(TestName()); // Initialize and transfer dynamic slice start indices parameter. - XlaOp starts; - std::unique_ptr start_data = CreateR1Parameter( - slice_starts, 0, "slice_starts", &builder, &starts); + std::vector starts(2); + std::vector> start_data(2); + for (int i = 0; i < 2; ++i) { + start_data[i] = CreateR0Parameter( + slice_starts[i], i, "slice_starts", &builder, &starts[i]); + } + // Build dynamic slice computation. auto input = ConstantLiteral(&builder, input_values); DynamicSlice(input, starts, slice_sizes); // Run computation and compare against expected values. - ComputeAndCompareLiteral(&builder, expected_values, {start_data.get()}); + std::vector argument_ptrs; + absl::c_transform(start_data, std::back_inserter(argument_ptrs), + [](const std::unique_ptr& argument) { + return argument.get(); + }); + ComputeAndCompareLiteral(&builder, expected_values, argument_ptrs); } template @@ -186,14 +195,22 @@ class DynamicSliceTest : public ClientLibraryTestBase { XlaBuilder builder(TestName()); // Initialize and transfer dynamic slice start indices parameter. - XlaOp starts; - std::unique_ptr start_data = CreateR1Parameter( - slice_starts, 0, "slice_starts", &builder, &starts); + std::vector starts(3); + std::vector> start_data(3); + for (int i = 0; i < 3; ++i) { + start_data[i] = CreateR0Parameter( + slice_starts[i], i, "slice_starts", &builder, &starts[i]); + } // Build dynamic slice computation. auto input = ConstantLiteral(&builder, input_values); DynamicSlice(input, starts, slice_sizes); // Run computation and compare against expected values. - ComputeAndCompareLiteral(&builder, expected_values, {start_data.get()}); + std::vector argument_ptrs; + absl::c_transform(start_data, std::back_inserter(argument_ptrs), + [](const std::unique_ptr& argument) { + return argument.get(); + }); + ComputeAndCompareLiteral(&builder, expected_values, argument_ptrs); } }; @@ -372,16 +389,12 @@ class DynamicUpdateSliceTest : public ClientLibraryTestBase { .ValueOrDie()); XlaBuilder builder(TestName()); - // Initialize and transfer dynamic slice start indices parameter. - XlaOp starts; - std::unique_ptr start_data = CreateR1Parameter( - slice_starts, 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. auto input = ConstantLiteral(&builder, input_value); auto update = ConstantLiteral(&builder, update_value); - DynamicUpdateSlice(input, update, starts); + DynamicUpdateSlice(input, update, absl::Span({})); // Run computation and compare against expected values. - ComputeAndCompareLiteral(&builder, expected_value, {start_data.get()}); + ComputeAndCompareLiteral(&builder, expected_value, {}); } template @@ -405,12 +418,12 @@ class DynamicUpdateSliceTest : public ClientLibraryTestBase { XlaBuilder builder(TestName()); // Initialize and transfer dynamic slice start indices parameter. XlaOp starts; - std::unique_ptr start_data = CreateR1Parameter( - slice_starts, 0, "slice_starts", &builder, &starts); + std::unique_ptr start_data = CreateR0Parameter( + slice_starts[0], 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. auto input = ConstantLiteral(&builder, input_values); auto update = ConstantLiteral(&builder, update_values); - DynamicUpdateSlice(input, update, starts); + DynamicUpdateSlice(input, update, absl::Span({starts})); // Run computation and compare against expected values. ComputeAndCompareLiteral(&builder, expected_values, {start_data.get()}); } @@ -435,15 +448,23 @@ class DynamicUpdateSliceTest : public ClientLibraryTestBase { XlaBuilder builder(TestName()); // Initialize and transfer dynamic slice start indices parameter. - XlaOp starts; - std::unique_ptr start_data = CreateR1Parameter( - slice_starts, 0, "slice_starts", &builder, &starts); + std::vector starts(2); + std::vector> start_data(2); + for (int i = 0; i < 2; ++i) { + start_data[i] = CreateR0Parameter( + slice_starts[i], i, "slice_starts", &builder, &starts[i]); + } // Build dynamic slice computation. auto input = ConstantLiteral(&builder, input_values); auto update = ConstantLiteral(&builder, update_values); DynamicUpdateSlice(input, update, starts); // Run computation and compare against expected values. - ComputeAndCompareLiteral(&builder, expected_values, {start_data.get()}); + std::vector argument_ptrs; + absl::c_transform(start_data, std::back_inserter(argument_ptrs), + [](const std::unique_ptr& argument) { + return argument.get(); + }); + ComputeAndCompareLiteral(&builder, expected_values, argument_ptrs); } template @@ -466,15 +487,24 @@ class DynamicUpdateSliceTest : public ClientLibraryTestBase { XlaBuilder builder(TestName()); // Initialize and transfer dynamic slice start indices parameter. - XlaOp starts; - std::unique_ptr start_data = CreateR1Parameter( - slice_starts, 0, "slice_starts", &builder, &starts); + std::vector starts(3); + std::vector> start_data(3); + for (int i = 0; i < 3; ++i) { + start_data[i] = CreateR0Parameter( + slice_starts[i], i, "slice_starts", &builder, &starts[i]); + } + // Build dynamic slice computation. auto input = ConstantLiteral(&builder, input_values); auto update = ConstantLiteral(&builder, update_values); DynamicUpdateSlice(input, update, starts); // Run computation and compare against expected values. - ComputeAndCompareLiteral(&builder, expected_values, {start_data.get()}); + std::vector argument_ptrs; + absl::c_transform(start_data, std::back_inserter(argument_ptrs), + [](const std::unique_ptr& argument) { + return argument.get(); + }); + ComputeAndCompareLiteral(&builder, expected_values, argument_ptrs); } template @@ -518,8 +548,9 @@ class DynamicUpdateSliceTest : public ClientLibraryTestBase { XlaOp update; std::unique_ptr update_data = CreateR3Parameter( update_values, 1, "update_values", &builder, &update); - auto starts = ConstantR1(&builder, {index, 0, 0}); - DynamicUpdateSlice(input, update, starts); + auto constant_index = ConstantR0(&builder, index); + auto zero = ConstantR0(&builder, 0); + DynamicUpdateSlice(input, update, {constant_index, zero, zero}); // Run computation and compare against expected values. ComputeAndCompareR3(&builder, expected_values, @@ -720,46 +751,55 @@ void BM_DynamicSlice(int num_iters) { {{13, 14, 15, 16}, {17, 18, 19, 20}, {21, 22, 23, 24}}}}); auto input = ConstantLiteral(&builder, input_literal); + auto stream = + client->mutable_backend()->BorrowStream(device_ordinal).ValueOrDie(); + // Create dynamic slice start indices as a parameter: shape [4] - auto start_indices_shape = ShapeUtil::MakeShape(S32, {4}); - auto start_indices = - Parameter(&builder, 0, start_indices_shape, "start_indices"); + auto start_indices_shape = ShapeUtil::MakeShape(S32, {}); + std::vector start_indices(4); + std::vector shaped_buffers; + std::vector host_shapes(4); + for (int i = 0; i < 4; ++i) { + start_indices[i] = + Parameter(&builder, i, start_indices_shape, "start_indices"); + auto start_index_literal = LiteralUtil::CreateR0(i + 1); + // Initialize and transfer parameter buffer. + shaped_buffers.emplace_back( + client->backend() + .transfer_manager() + ->AllocateScopedShapedBuffer(start_indices_shape, &allocator, + /*device_ordinal=*/0) + .ConsumeValueOrDie()); + host_shapes[i] = &shaped_buffers[i].on_host_shape(); + ASSERT_IS_OK(transfer_manager->TransferLiteralToDevice( + stream.get(), start_index_literal, shaped_buffers[i])); + } + // Add DynamicSlice op to the computatation. DynamicSlice(input, start_indices, {1, 1, 1, 1}); auto computation = builder.Build().ConsumeValueOrDie(); - // Initialize and transfer parameter buffer. - auto buffer = client->backend() - .transfer_manager() - ->AllocateScopedShapedBuffer( - start_indices_shape, &allocator, /*device_ordinal=*/0) - .ConsumeValueOrDie(); - - auto start_indices_literal = LiteralUtil::CreateR1({0, 1, 2, 3}); - auto stream = - client->mutable_backend()->BorrowStream(device_ordinal).ValueOrDie(); - ASSERT_IS_OK(transfer_manager->TransferLiteralToDevice( - stream.get(), start_indices_literal, buffer)); - std::unique_ptr executable = - client - ->Compile(computation, {&buffer.on_host_shape()}, - ExecutableBuildOptions()) + client->Compile(computation, host_shapes, ExecutableBuildOptions()) .ConsumeValueOrDie(); // Run some warm-up executions. ExecutableRunOptions options; options.set_allocator(&allocator); const int kWarmups = 2; + std::vector shaped_buffer_ptrs; + absl::c_transform(shaped_buffers, std::back_inserter(shaped_buffer_ptrs), + [](const ScopedShapedBuffer& buffer) { return &buffer; }); + for (int i = 0; i < kWarmups; ++i) { - auto result = executable->Run({&buffer}, options); + auto result = executable->Run(shaped_buffer_ptrs, options); ASSERT_TRUE(result.ok()); } // Run benchmark. tensorflow::testing::StartTiming(); for (int i = 0; i < num_iters; ++i) { - auto result = executable->Run({&buffer}, options); + auto result = executable->Run(shaped_buffer_ptrs, options); ASSERT_TRUE(result.ok()); } } diff --git a/tensorflow/compiler/xla/tests/exhaustive_f32_elementwise_op_test.cc b/tensorflow/compiler/xla/tests/exhaustive_f32_elementwise_op_test.cc index c84973e17b234c24c84f02a369ce0185f5772cca..b961e6102692cb3b90976d621c62cb4cf18a9b6b 100644 --- a/tensorflow/compiler/xla/tests/exhaustive_f32_elementwise_op_test.cc +++ b/tensorflow/compiler/xla/tests/exhaustive_f32_elementwise_op_test.cc @@ -13,7 +13,9 @@ 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" @@ -21,66 +23,166 @@ limitations under the License. namespace xla { namespace { + class ExhaustiveF32ElementwiseOpTest : public ClientLibraryTestBase, public ::testing::WithParamInterface> { protected: - ErrorSpec error_spec_{0.0001, 0.0001, /*relaxed_nans=*/true}; + 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()); - Literal input_literal = - LiteralUtil::CreateFromDimensions(F32, {input_size}); - for (int64 i = begin; i < end; i++) { + 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) { - // 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. - input_literal.Set({i - begin}, 0.0f); - } else { - input_literal.Set({i - begin}, absl::bit_cast(i)); + return 0; } - } - - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr input_data, - client_->TransferToServer(input_literal)); + 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; + } - std::vector expected_result; - expected_result.reserve(input_size); - for (int64 i = 0; i < input_size; i++) { - expected_result.push_back(evaluate_op(input_literal.Get({i}))); - } + 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)); + } + } - ComputeAndCompareR1(&builder, expected_result, {input_data.get()}, - error_spec_); + 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) { -#ifdef XLA_TEST_BACKEND_CPU - // TODO(b/73141998): The vectorized Log 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 = {1, 8388608}; -#else - std::pair known_incorrect_range = {0, 0}; +#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); + /*known_incorrect_range=*/{0, 0}); } XLA_TEST_P(ExhaustiveF32ElementwiseOpTest, ExpF32) { @@ -105,6 +207,18 @@ XLA_TEST_P(ExhaustiveF32ElementwiseOpTest, TanhF32) { /*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. diff --git a/tensorflow/compiler/xla/tests/fusion_test.cc b/tensorflow/compiler/xla/tests/fusion_test.cc index d1fddf9d6b494a822610e41307fa103dc90bdef3..2178c9b3f3d39ac034c59585c6836d2bc59162c1 100644 --- a/tensorflow/compiler/xla/tests/fusion_test.cc +++ b/tensorflow/compiler/xla/tests/fusion_test.cc @@ -523,10 +523,10 @@ XLA_TEST_F(FusionTest, DynamicSliceNegate) { auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({1, 2, 3, 4}))); auto const1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({1}))); + HloInstruction::CreateConstant(LiteralUtil::CreateR0(1))); auto dynamic_slice2 = builder.AddInstruction(HloInstruction::CreateDynamicSlice( - ShapeUtil::MakeShape(S32, {2}), const0, const1, {2})); + ShapeUtil::MakeShape(S32, {2}), const0, {const1}, {2})); auto negate3 = builder.AddInstruction(HloInstruction::CreateUnary( ShapeUtil::MakeShape(S32, {2}), HloOpcode::kNegate, dynamic_slice2)); hlo_module->AddEntryComputation(builder.Build()) diff --git a/tensorflow/compiler/xla/tests/gather_operation_test.cc b/tensorflow/compiler/xla/tests/gather_operation_test.cc index daa89398a697af9149797d621c3bdca80a00aedd..d65b67a535d43553a3a94f76482ad4618f9b8aab 100644 --- a/tensorflow/compiler/xla/tests/gather_operation_test.cc +++ b/tensorflow/compiler/xla/tests/gather_operation_test.cc @@ -600,7 +600,9 @@ ENTRY main { class GatherClientLibraryTest : public ClientLibraryTestBase {}; -XLA_TEST_F(GatherClientLibraryTest, DISABLED_ON_GPU(Basic)) { +// Disabled on interpreter since ExectuteAsyncOnStream is not supported. +XLA_TEST_F(GatherClientLibraryTest, + DISABLED_ON_INTERPRETER(DISABLED_ON_GPU(Basic))) { // We create this HLO, but using the XlaBuilder API. // // ENTRY main { diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.cc b/tensorflow/compiler/xla/tests/hlo_test_base.cc index d57846e19bb80c5b9c87d50596da2915f9aef317..d9d54fd2556be01a56afd36c13fcc8cf2184ece8 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.cc +++ b/tensorflow/compiler/xla/tests/hlo_test_base.cc @@ -139,7 +139,8 @@ std::unique_ptr HloTestBase::CreateNewVerifiedModule( const string& name) { return absl::make_unique( name, GetModuleConfigForTest(), verifier_layout_sensitive_, - allow_mixed_precision_in_hlo_verifier_); + allow_mixed_precision_in_hlo_verifier_, + backend().compiler()->ShapeSizeBytesFunction()); } StatusOr> @@ -147,7 +148,8 @@ HloTestBase::ParseAndReturnVerifiedModule(absl::string_view hlo_text, const HloModuleConfig& config) { auto module = absl::make_unique( TestName(), config, verifier_layout_sensitive_, - allow_mixed_precision_in_hlo_verifier_); + allow_mixed_precision_in_hlo_verifier_, + backend().compiler()->ShapeSizeBytesFunction()); TF_RETURN_IF_ERROR(ParseHloString(hlo_text, module.get())); TF_RETURN_IF_ERROR(module->Verify()); return std::move(module); @@ -311,7 +313,10 @@ StatusOr<::testing::AssertionResult> HloTestBase::RunAndCompareInternal( reference_preprocessor); } -::testing::AssertionResult HloTestBase::Run(string_view hlo_string) { +::testing::AssertionResult HloTestBase::Run(string_view hlo_string, + bool run_hlo_passes, + ExecutionProfile* profile, + string backend_config) { auto module_or_status = HloRunner::CreateModuleFromString(hlo_string, GetDebugOptionsForTest()); if (!module_or_status.ok()) { @@ -319,19 +324,108 @@ StatusOr<::testing::AssertionResult> HloTestBase::RunAndCompareInternal( << "Error while parsing HLO text format: " << module_or_status.status().ToString(); } + + std::unique_ptr module = std::move(module_or_status.ValueOrDie()); const auto& fake_arguments = - MakeFakeArguments(module_or_status.ValueOrDie().get()) - .ConsumeValueOrDie(); + MakeFakeArguments(module.get()).ConsumeValueOrDie(); std::vector fake_argument_ptrs; absl::c_transform( fake_arguments, std::back_inserter(fake_argument_ptrs), [](const Literal& literal) { return const_cast(&literal); }); - return test_runner_ - .Execute(std::move(module_or_status.ValueOrDie()), - fake_argument_ptrs, /*run_hlo_passes=*/true) - .ok() + + if (profile != nullptr) { + // We have to enable HLO profiling since otherwise currently the + // ExecutionProfile is not correct. + // + // TODO(b/119432044): Fix collection of the ExecutionProfile + // so that this is not necessary. + HloModuleConfig config = module->config(); + DebugOptions debug_options = config.debug_options(); + debug_options.set_xla_hlo_profile(true); + config.set_debug_options(debug_options); + module->set_config(config); + } + + if (!backend_config.empty()) { + // Set backend configuration if it is given. + HloInstruction* instruction = + module->entry_computation()->root_instruction(); + instruction->set_raw_backend_config_string(backend_config); + } + + // return ::testing::AssertionSuccess(); + auto output = test_runner_.Execute(std::move(module), fake_argument_ptrs, + /*run_hlo_passes=*/run_hlo_passes, + /*profile=*/profile); + + return output.ok() ? ::testing::AssertionSuccess() - : ::testing::AssertionFailure(); + : ::testing::AssertionFailure() << output.status().error_message(); +} + +::testing::AssertionResult HloTestBase::RunMultipleTimes( + string_view hlo_string, bool run_hlo_passes, + std::vector* profiles, string backend_config) { + int n = profiles->size(); + std::vector> fake_argument_ptrs(n); + std::vector> fake_arguments(n); + std::vector> executables(n); + + for (int i = 0; i < n; ++i) { + auto module_or_status = + HloRunner::CreateModuleFromString(hlo_string, GetDebugOptionsForTest()); + if (!module_or_status.ok()) { + return ::testing::AssertionFailure() + << "Error while parsing HLO text format: " + << module_or_status.status().ToString(); + } + std::unique_ptr module = + std::move(module_or_status.ValueOrDie()); + + fake_arguments[i] = MakeFakeArguments(module.get()).ConsumeValueOrDie(); + absl::c_transform( + fake_arguments[i], std::back_inserter(fake_argument_ptrs[i]), + [](const Literal& literal) { return const_cast(&literal); }); + + if (profiles != nullptr) { + // We have to enable HLO profiling since otherwise currently the + // ExecutionProfile is not correct. + // + // TODO(b/119432044): Fix collection of the ExecutionProfile + // so that this is not necessary. + HloModuleConfig config = module->config(); + DebugOptions debug_options = config.debug_options(); + debug_options.set_xla_hlo_profile(true); + config.set_debug_options(debug_options); + module->set_config(config); + } + + if (!backend_config.empty()) { + // Set backend configuration if it is given. + HloInstruction* instruction = + module->entry_computation()->root_instruction(); + instruction->set_raw_backend_config_string(backend_config); + } + + auto executable = + test_runner_.CreateExecutable(std::move(module), run_hlo_passes); + if (!executable.ok()) { + return ::testing::AssertionFailure() + << executable.status().error_message(); + } + executables[i] = std::move(executable.ValueOrDie()); + } + + for (int i = 0; i < n; ++i) { + auto output = + test_runner_.Execute(std::move(executables[i]), fake_argument_ptrs[i], + /*profile=*/&((*profiles)[i])); + if (!output.ok()) { + return ::testing::AssertionFailure() << output.status().error_message(); + } + } + + return ::testing::AssertionSuccess(); } ::testing::AssertionResult HloTestBase::RunAndCompareFromFile( diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.h b/tensorflow/compiler/xla/tests/hlo_test_base.h index 1d1e7f437296a7493ef7da07039fcf6d273f35bc..78bdd336e0a96999440f6331a965987ee0cb6bf2 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.h +++ b/tensorflow/compiler/xla/tests/hlo_test_base.h @@ -46,10 +46,12 @@ class VerifiedHloModule : public HloModule { public: VerifiedHloModule(const string& name, const HloModuleConfig& config, bool verifier_layout_sensitive, - bool allow_mixed_precision_in_hlo_verifier) + bool allow_mixed_precision_in_hlo_verifier, + std::function shape_size_function) : HloModule(name, config), - verifier_(verifier_layout_sensitive, - allow_mixed_precision_in_hlo_verifier) {} + verifier_( + verifier_layout_sensitive, allow_mixed_precision_in_hlo_verifier, + /*instruction_can_change_layout_func=*/{}, shape_size_function) {} ~VerifiedHloModule() override { VerifyOrAddFailure("in destructor"); } @@ -219,8 +221,14 @@ class HloTestBase : public ::testing::Test { const absl::optional& error, const std::function& reference_preprocessor = nullptr) TF_MUST_USE_RESULT; - ::testing::AssertionResult Run(const absl::string_view hlo_string) - TF_MUST_USE_RESULT; + ::testing::AssertionResult Run(const absl::string_view hlo_string, + bool run_hlo_passes = true, + ExecutionProfile* profile = nullptr, + string backend_config = "") TF_MUST_USE_RESULT; + ::testing::AssertionResult RunMultipleTimes( + const absl::string_view hlo_string, bool run_hlo_passes, + std::vector* profiles, + string backend_config = "") TF_MUST_USE_RESULT; ::testing::AssertionResult RunAndCompareFromFile( const string& filename, const absl::optional& error, const std::function& reference_preprocessor = nullptr) diff --git a/tensorflow/compiler/xla/tests/literal_test_util.cc b/tensorflow/compiler/xla/tests/literal_test_util.cc index 554eb24d44168caa7d7252015e3d99f2d567df9b..a2fd6070731943f15c773265f428b16f520d02ee 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util.cc +++ b/tensorflow/compiler/xla/tests/literal_test_util.cc @@ -86,7 +86,7 @@ void OnMiscompare(const LiteralSlice& expected, const LiteralSlice& actual, /* static */ ::testing::AssertionResult LiteralTestUtil::Near( const LiteralSlice& expected, const LiteralSlice& actual, - const ErrorSpec& error_spec, bool detailed_message) { + const ErrorSpec& error_spec, absl::optional detailed_message) { return StatusToAssertion(literal_comparison::Near( expected, actual, error_spec, detailed_message, &OnMiscompare)); } @@ -97,7 +97,8 @@ void OnMiscompare(const LiteralSlice& expected, const LiteralSlice& actual, if (error.has_value()) { VLOG(1) << "Expects near"; return StatusToAssertion(literal_comparison::Near( - expected, actual, *error, /*detailed_message=*/false, &OnMiscompare)); + expected, actual, *error, /*detailed_message=*/absl::nullopt, + &OnMiscompare)); } VLOG(1) << "Expects equal"; return StatusToAssertion(literal_comparison::Equal(expected, actual)); diff --git a/tensorflow/compiler/xla/tests/literal_test_util.h b/tensorflow/compiler/xla/tests/literal_test_util.h index 43cca91f64b2c0fbfde5054a361cf0f95302c23d..d7cf9bed98a3eb7479b6deb6838dc388a0869360 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util.h +++ b/tensorflow/compiler/xla/tests/literal_test_util.h @@ -93,7 +93,7 @@ class LiteralTestUtil { static ::testing::AssertionResult Near( const LiteralSlice& expected, const LiteralSlice& actual, const ErrorSpec& error_spec, - bool detailed_message = false) TF_MUST_USE_RESULT; + absl::optional detailed_message = absl::nullopt) TF_MUST_USE_RESULT; // Asserts the given literal are within the given error bound of the given // expected values. Only supported for floating point values. diff --git a/tensorflow/compiler/xla/tests/local_client_execute_test.cc b/tensorflow/compiler/xla/tests/local_client_execute_test.cc index 6522c563252a0e8271bf733668f39bcf00a80d06..96527886b718bc1ea4ce8cc2d7dbeb2e3ef1d1eb 100644 --- a/tensorflow/compiler/xla/tests/local_client_execute_test.cc +++ b/tensorflow/compiler/xla/tests/local_client_execute_test.cc @@ -842,7 +842,8 @@ XLA_TEST_F(LocalClientExecuteTest, ShapeBufferToLiteralConversion64bit) { LiteralUtil::CreateR0(123456789000LL)})); } -XLA_TEST_F(LocalClientExecuteTest, InfeedTest) { +// Disabled on interpreter backend since infeed HLO is unsupported. +XLA_TEST_F(LocalClientExecuteTest, DISABLED_ON_INTERPRETER(InfeedTest)) { XlaBuilder builder(TestName()); const Shape shape = ShapeUtil::MakeShape(F32, {3}); auto in = Infeed(&builder, shape); @@ -867,7 +868,8 @@ XLA_TEST_F(LocalClientExecuteTest, InfeedTest) { LiteralTestUtil::ExpectR1Equal({-4.0, 125.0, 45.0}, result); } -XLA_TEST_F(LocalClientExecuteTest, InfeedOutfeedTest) { +// Disabled on interpreter backend since infeed/outfeed HLOs are unsupported. +XLA_TEST_F(LocalClientExecuteTest, DISABLED_ON_INTERPRETER(InfeedOutfeedTest)) { XlaBuilder builder(TestName()); const Shape shape = ShapeUtil::MakeShape(F32, {3}); auto in = Infeed(&builder, shape); diff --git a/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc b/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc index 3f5135438fc59bea98527b1be30ee49339edd455..1fd9cb055c0bebc0f31496eb82f53a7b7a6cbfba 100644 --- a/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc +++ b/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc @@ -208,9 +208,7 @@ XLA_TEST_F(MultiOutputFusionTest, FusionNodeIsRoot) { ROOT fusion = (s32[]) fusion(x), kind=kLoop, calls=fused_computation } )"; - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param = LiteralUtil::MakeTupleOwned( LiteralUtil::MakeTupleOwned( LiteralUtil::MakeTupleOwned(LiteralUtil::CreateR0(42)), @@ -241,9 +239,7 @@ XLA_TEST_F(MultiOutputFusionTest, MultiOutputLoopFusion) { const = f32[4] constant({0, 0, 0, 0}) ROOT select = f32[4] select(gte0, gte1, const) })"; - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param = LiteralUtil::CreateR1({1.0, 2.0, 3.0, -1.0}); Literal result = ExecuteNoHloPasses(std::move(module), {¶m}); LiteralTestUtil::ExpectR1Equal({0.0, 4.0, 9.0, 1.0}, result); @@ -273,9 +269,7 @@ XLA_TEST_F(MultiOutputFusionTest, MultiOutputLoopFeedingMap) { p1 = f32[3] parameter(0) ROOT map = f32[3] map(p1), to_apply=map_computation })"; - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param = LiteralUtil::CreateR1({1.0, 2.0, 3.0}); Literal result = ExecuteNoHloPasses(std::move(module), {¶m}); LiteralTestUtil::ExpectR1Equal({0.0, 4.0, 9.0}, result); @@ -315,9 +309,7 @@ XLA_TEST_F(MultiOutputFusionTest, ROOT fusion = (f32[2,2]{1,0}, f32[2,2]{1,0}) fusion(p), kind=kInput, calls=fused_reduce })"); - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param = LiteralUtil::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}); Literal result = ExecuteNoHloPasses(std::move(module), {¶m}); @@ -346,9 +338,7 @@ XLA_TEST_F(MultiOutputFusionTest, ROOT fusion = (f32[2,2]{1,0}, f32[2,2]{1,0}) fusion(p), kind=kInput, calls=fused_reduce })"); - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param = LiteralUtil::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}); Literal result = ExecuteNoHloPasses(std::move(module), {¶m}); @@ -378,9 +368,7 @@ XLA_TEST_F(MultiOutputFusionTest, ROOT fusion = (f32[2]{0}, f32[2]{0}, f32[2]{0}) fusion(p), kind=kInput, calls=fused_reduce })"); - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param = LiteralUtil::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}); Literal result = ExecuteNoHloPasses(std::move(module), {¶m}); @@ -410,9 +398,7 @@ XLA_TEST_F(MultiOutputFusionTest, ROOT fusion = (f32[2,2,2]{2,1,0}, f32[2,2]{1,0}, f32[2,2]{1,0}) fusion(p), kind=kInput, calls=fused_reduce })"); - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param = LiteralUtil::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}); Literal result = ExecuteNoHloPasses(std::move(module), {¶m}); @@ -443,9 +429,7 @@ XLA_TEST_F(MultiOutputFusionTest, ROOT fusion = (f32[2,2]{1,0}, f32[2,2,2]{2,1,0}, f32[2,2]{1,0}) fusion(p), kind=kInput, calls=fused_reduce })"); - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param = LiteralUtil::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}); Literal result = ExecuteNoHloPasses(std::move(module), {¶m}); @@ -478,9 +462,7 @@ XLA_TEST_F(MultiOutputFusionTest, ROOT fusion = (f32[2]{0}, f32[2,2,2]{2,1,0}, f32[2,2,2]{2,1,0}) fusion(p), kind=kInput, calls=fused_reduce })"); - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param = LiteralUtil::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}); Literal result = ExecuteNoHloPasses(std::move(module), {¶m}); @@ -513,9 +495,7 @@ XLA_TEST_F(MultiOutputFusionTest, ROOT fusion = (f32[2,2]{1,0}, f32[2,2]{1,0}) fusion(p, i, j), kind=kInput, calls=fused_reduce })"); - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param = LiteralUtil::CreateR3({{{0, 2}, {3, 4}}, {{5, 6}, {7, 8}}}); auto init1 = LiteralUtil::CreateR0(5); @@ -549,9 +529,7 @@ XLA_TEST_F(MultiOutputFusionTest, ROOT fusion = (f32[2,2]{1,0}, f32[2,2]{1,0}, f16[2,2,2]{2,1,0}) fusion(p), kind=kInput, calls=fused_reduce })"); - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param = LiteralUtil::CreateR3( {{{Eigen::half(1), Eigen::half(2)}, {Eigen::half(3), Eigen::half(4)}}, {{Eigen::half(5), Eigen::half(6)}, {Eigen::half(7), Eigen::half(8)}}}); diff --git a/tensorflow/compiler/xla/tests/prng_test.cc b/tensorflow/compiler/xla/tests/prng_test.cc index 8f2c26f0eea9c7a3b33cd77e5977924c1659535a..e49bcf26bd6e50f8fb36c86f217907b5d4901eae 100644 --- a/tensorflow/compiler/xla/tests/prng_test.cc +++ b/tensorflow/compiler/xla/tests/prng_test.cc @@ -80,7 +80,9 @@ XLA_TEST_F(PrngTest, LargeU01) { UniformTest(0, 1, {0x100, 0x100}); } XLA_TEST_F(PrngTest, TwelveValuesU524) { UniformTest(5, 24, {12}); } // TODO(b/71543667): Fix Rng ops on LLVM backends. -XLA_TEST_F(PrngTest, DISABLED_ON_GPU(DISABLED_ON_CPU(ScalarBF16Tests))) { +// TODO(b/122047800): Interpreter does not support BF16 for RNG ops. +XLA_TEST_F(PrngTest, DISABLED_ON_INTERPRETER( + DISABLED_ON_GPU(DISABLED_ON_CPU(ScalarBF16Tests)))) { for (int64 seed = 0; seed < 100; ++seed) { // The largest negative number smaller than zero in bf16 that's not // denormalized. @@ -103,7 +105,9 @@ XLA_TEST_F(PrngTest, DISABLED_ON_GPU(DISABLED_ON_CPU(ScalarBF16Tests))) { } // TODO(b/71543667): Fix Rng ops on LLVM backends. -XLA_TEST_F(PrngTest, DISABLED_ON_GPU(DISABLED_ON_CPU(ScalarBF16CountTests))) { +// TODO(b/122047800): Interpreter does not support BF16 for RNG ops. +XLA_TEST_F(PrngTest, DISABLED_ON_INTERPRETER(DISABLED_ON_GPU( + DISABLED_ON_CPU(ScalarBF16CountTests)))) { // There are 3 BF16 values in the range of [32.25, 33): 32.25, 32.5, 32.75, // they should get similar counts. bfloat16 low = static_cast(32.25); @@ -276,6 +280,39 @@ XLA_TEST_F(PrngTest, PassInGlobalRngSeed) { EXPECT_FALSE(LiteralTestUtil::Equal(result5, result6)); } +// This test verifies that the two RNG instructions with the same parameters in +// the same HloComputation produces different values. +XLA_TEST_F(PrngTest, DifferentValuesForIdenticalRngNodesInSameComputation) { + // Build a U[0,1) computation. + auto build_computation = [this]() { + XlaBuilder builder(TestName()); + auto a = RngUniform(ConstantR0(&builder, 0), + ConstantR0(&builder, 100), + ShapeUtil::MakeShape(S32, {10})); + auto b = RngUniform(ConstantR0(&builder, 0), + ConstantR0(&builder, 100), + ShapeUtil::MakeShape(S32, {10})); + Tuple(&builder, {a, b}); + return builder.Build(); + }; + + ExecutionOptions execution_options = execution_options_; + execution_options.set_seed(42); + + Literal result_tuple; + { + TF_ASSERT_OK_AND_ASSIGN(auto computation, build_computation()); + TF_ASSERT_OK_AND_ASSIGN( + result_tuple, client_->ExecuteAndTransfer(computation, /*arguments=*/{}, + &execution_options)); + } + + auto results = result_tuple.DecomposeTuple(); + ASSERT_EQ(results.size(), 2); + + EXPECT_FALSE(LiteralTestUtil::Equal(results[0], results[1])); +} + XLA_TEST_F(PrngTest, TenValuesN01) { XlaBuilder builder(TestName()); RngNormal(ConstantR0(&builder, 0), ConstantR0(&builder, 1), diff --git a/tensorflow/compiler/xla/tests/ptxas_bug_120501638.cc b/tensorflow/compiler/xla/tests/ptxas_bug_120501638.cc new file mode 100644 index 0000000000000000000000000000000000000000..0e5d7db97e88936e7336ed02a5c7a1171254b0cf --- /dev/null +++ b/tensorflow/compiler/xla/tests/ptxas_bug_120501638.cc @@ -0,0 +1,82 @@ +/* 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/debug_options_flags.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/tests/test_macros.h" + +namespace xla { +namespace { + +class PtxasBugTest : public HloTestBase {}; + +// Checks for a bug in ptxas, tracked as Google bug 120501638, and nvidia bug +// 2459377. We never received an explanation of what exactly was going wrong +// here in ptxas. Known-bad in ptxas 10.0.145, known-good in ptxas 10.0.249. +TEST_F(PtxasBugTest, DoIt) { + const char* const kModuleStr = R"( +HloModule test + +add_F32.14 { + lhs.15 = f32[] parameter(0) + rhs.16 = f32[] parameter(1) + ROOT add.17 = f32[] add(lhs.15, rhs.16) +} + +ENTRY testcase { + arg0.1 = f32[2,5,2]{2,1,0} parameter(0) + reshape.2 = f32[2,5,2]{2,1,0} reshape(arg0.1) + constant.3 = f32[] constant(0) + pad.4 = f32[2,6,2]{2,1,0} pad(reshape.2, constant.3), padding=0_0x0_1x0_0 + reshape.5 = f32[2,3,2,2]{3,2,1,0} reshape(pad.4) + transpose.6 = f32[2,2,3,2]{3,0,2,1} transpose(reshape.5), dimensions={2,0,1,3} + reshape.7 = f32[4,3,2]{2,1,0} reshape(transpose.6) + reshape.8 = f32[4,1,3,2]{3,2,1,0} reshape(reshape.7) + transpose.9 = f32[4,2,1,3]{1,3,2,0} transpose(reshape.8), dimensions={0,3,1,2} + convert.10 = f32[4,2,1,3]{1,3,2,0} convert(transpose.9) + constant.12 = f32[] constant(0) + pad.13 = f32[4,2,1,3]{3,2,1,0} pad(convert.10, constant.12), padding=0_0x0_0x0_0x0_0 + constant.11 = f32[] constant(0) + reduce-window.18 = f32[4,2,1,3]{3,2,1,0} reduce-window(pad.13, constant.11), + window={size=1x1x1x1}, to_apply=add_F32.14 + constant.19 = f32[] constant(1) + broadcast.20 = f32[4,2,1,3]{3,2,1,0} broadcast(constant.19), dimensions={} + divide.21 = f32[4,2,1,3]{3,2,1,0} divide(reduce-window.18, broadcast.20) + convert.22 = f32[4,2,1,3]{3,2,1,0} convert(divide.21) + transpose.23 = f32[4,1,3,2]{2,1,3,0} transpose(convert.22), dimensions={0,2,3,1} + reshape.24 = f32[4,3,2]{2,1,0} reshape(transpose.23) + reshape.25 = f32[2,2,3,2]{3,2,1,0} reshape(reshape.24) + transpose.26 = f32[2,3,2,2]{3,1,0,2} transpose(reshape.25), dimensions={1,2,0,3} + reshape.27 = f32[2,6,2]{2,1,0} reshape(transpose.26) + slice.28 = f32[2,5,2]{2,1,0} slice(reshape.27), slice={[0:2], [0:5], [0:2]} + reshape.29 = f32[2,5,2]{2,1,0} reshape(slice.28) + tuple.30 = (f32[2,5,2]{2,1,0}) tuple(reshape.29) + ROOT get-tuple-element.31 = f32[2,5,2]{2,1,0} get-tuple-element(tuple.30), index=0 +})"; + + // Create a module with the true-default flags, not the default-for-testing + // flags. In particular, true-default flags enable unrolling, whereas for + // testing we disable unrolling, and this bug doesn't trigger without + // unrolling. + HloModuleConfig config; + config.set_debug_options(DefaultDebugOptionsIgnoringFlags()); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kModuleStr, config)); + EXPECT_TRUE(RunAndCompare(std::move(module), ErrorSpec{0.01, 0.01})); +} + +} // anonymous namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/reduce_precision_test.cc b/tensorflow/compiler/xla/tests/reduce_precision_test.cc index f80d29b9de440b11c36e8c9bc65d4a93353a6267..e2cf4c0be289b52d5cc581ea07752ed6e98da76f 100644 --- a/tensorflow/compiler/xla/tests/reduce_precision_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_precision_test.cc @@ -34,7 +34,6 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/compiler/xla/tests/test_macros.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/types.h" namespace xla { diff --git a/tensorflow/compiler/xla/tests/reduce_window_test.cc b/tensorflow/compiler/xla/tests/reduce_window_test.cc index 22fe4a2670e2e0e1fedc45036a1ceec19f44e42e..30e2d24184a5d399e5e058a9c4a382f57e82866f 100644 --- a/tensorflow/compiler/xla/tests/reduce_window_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_window_test.cc @@ -607,7 +607,10 @@ class R4ReduceWindowTest : public ReduceWindowTestBase, Array4D input(param.base_bounds[0], param.base_bounds[1], param.base_bounds[2], param.base_bounds[3]); - input.FillRandom(0.1f, 0.1f); + // Choose a prime iota length so that each window sees a unique set of + // values. (Technically, the requirement is that the iota length is + // relatively prime to all of the dimensions involved in the reduce-window.) + input.FillRepeatedIota(0, 137); Literal input_literal = LiteralUtil::CreateR4FromArray4DWithLayout( input, LayoutUtil::MakeLayout(param.layout)); XlaOp parameter; @@ -623,9 +626,9 @@ class R4ReduceWindowTest : public ReduceWindowTestBase, CreateConstantFromLiteral(LiteralUtil::CreateR0(kInitValue), &b); CHECK(param.reducer == kAdd || param.reducer == kMax); auto reducer = param.reducer; - if (use_bfloat16() && Product(param.window_bounds) > 128) { - // To avoid numerical issues, force the reducer to be kMax for large bf16 - // windows. + if (use_bfloat16()) { + // To avoid numerical issues, force the reducer to be kMax for bf16 + // inputs. reducer = kMax; } @@ -949,16 +952,16 @@ struct R3ReduceWindowTestData { /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, {/*base_bounds=*/{95, 202, 251}, /*window_bounds=*/{95, 202, 251}, /*strides=*/{1, 1, 1}, /*layout=*/{2, 1, 0}, - /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kMax}, {/*base_bounds=*/{999, 57, 3}, /*window_bounds=*/{999, 57, 3}, /*strides=*/{1, 1, 1}, /*layout=*/{2, 1, 0}, /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, {/*base_bounds=*/{178, 302, 64}, /*window_bounds=*/{178, 302, 64}, /*strides=*/{1, 1, 1}, /*layout=*/{2, 1, 0}, - /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kMax}, {/*base_bounds=*/{63, 261, 257}, /*window_bounds=*/{63, 261, 257}, /*strides=*/{1, 1, 1}, /*layout=*/{2, 1, 0}, - /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kMax}, {/*base_bounds=*/{10003, 10, 5}, /*window_bounds=*/{9999, 7, 3}, /*strides=*/{1, 1, 1}, /*layout=*/{2, 1, 0}, /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, @@ -1001,17 +1004,19 @@ TEST_P(R3ReduceWindowTest, DoIt) { const float kInitValue = 0.0f; Array3D input(param.base_bounds[0], param.base_bounds[1], param.base_bounds[2]); - input.FillRandom(0.1f, 0.1f); + // Choose a prime iota length so that each window sees a unique set of values. + // (Technically, the requirement is that the iota length is relatively prime + // to all of the dimensions involved in the reduce-window.) + input.FillRepeatedIota(0, 137); Literal input_literal = LiteralUtil::CreateR3FromArray3DWithLayout( input, LayoutUtil::MakeLayout(param.layout)); auto reducer = param.reducer; if (use_bfloat16()) { input_literal = LiteralUtil::ConvertF32ToBF16(input_literal); - if (Product(param.window_bounds) > 128) { - // To avoid numerical issues, force the reducer to be kMax for large bf16 - // windows. - reducer = kMax; - } + + // To avoid numerical issues, force the reducer to be kMax for bf16 + // inputs. + reducer = kMax; } XlaOp parameter = Parameter(&b, 0, input_literal.shape(), "input"); @@ -1527,6 +1532,25 @@ ENTRY %reduce-window (parameter.0: s32[81,8], parameter.1: s32[]) -> s32[82,8] { EXPECT_TRUE(RunAndCompare(hlo_string, absl::nullopt)); } +XLA_TEST_F(HloTestBase, ReduceWindowS64) { + const string hlo_string = R"( +HloModule reduce-window + +%identity.pad_to_reduce_window (param0: s64[], param1: s64[]) -> s64[] { + %param0 = s64[] parameter(0) + ROOT %param1 = s64[] parameter(1) +} + +ENTRY %reduce-window (parameter.0: s64[81,8], parameter.1: s64[]) -> s64[82,8] { + %parameter.0 = s64[81,8]{1,0} parameter(0) + %parameter.1 = s64[] parameter(1) + ROOT %reduce-window = s64[82,8]{1,0} reduce-window(s64[81,8]{1,0} %parameter.0, s64[] %parameter.1), window={size=1x1 pad=0_1x0_0}, to_apply=%identity.pad_to_reduce_window +} + +)"; + EXPECT_TRUE(RunAndCompare(hlo_string, absl::nullopt)); +} + XLA_TEST_F(HloTestBase, ReduceWindowF16) { const string hlo_string = R"( HloModule reduce-window diff --git a/tensorflow/compiler/xla/tests/test_utils.cc b/tensorflow/compiler/xla/tests/test_utils.cc index 18df41c1172d86aa72cd5bfabd8a56a202d4bdbd..95c89b0ba6f29c453abab88e29bca13ee006455a 100644 --- a/tensorflow/compiler/xla/tests/test_utils.cc +++ b/tensorflow/compiler/xla/tests/test_utils.cc @@ -19,6 +19,7 @@ limitations under the License. #include "absl/memory/memory.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/primitive_util.h" +#include "tensorflow/compiler/xla/service/hlo_casting_utils.h" #include "tensorflow/compiler/xla/service/hlo_dataflow_analysis.h" #include "tensorflow/compiler/xla/service/hlo_verifier.h" #include "tensorflow/compiler/xla/service/transfer_manager.h" @@ -274,16 +275,9 @@ bool NeedsInitValue(const HloUse& use) { // Generate random values that are constrained to the input_shape minus the // output_shape so as not to produce wrapping slices, for instance. -Literal MakeRandomIndex(absl::Span index_space, - std::minstd_rand0* engine) { - std::vector start_indices(index_space.size()); - if (engine != nullptr) { - for (int i = 0; i < index_space.size(); ++i) { - std::uniform_int_distribution generator(0, index_space[i]); - start_indices[i] = generator(*engine); - } - } - return LiteralUtil::CreateR1(start_indices); +Literal MakeRandomIndex(int64 index_bound, std::minstd_rand0* engine) { + std::uniform_int_distribution generator(0, index_bound); + return LiteralUtil::CreateR0(generator(*engine)); } // Use dataflow analysis on each parameter to see if there are uses that would @@ -300,8 +294,8 @@ std::vector FindConstrainedUses( HloInstruction* instruction = use.instruction; const HloOpcode opcode = instruction->opcode(); const int64 op_num = use.operand_number; - if ((opcode == HloOpcode::kDynamicSlice && op_num == 1) || - (opcode == HloOpcode::kDynamicUpdateSlice && op_num == 2)) { + if ((opcode == HloOpcode::kDynamicSlice && op_num >= 1) || + (opcode == HloOpcode::kDynamicUpdateSlice && op_num >= 2)) { constrained_uses.push_back(instruction); } else if (opcode == HloOpcode::kFusion) { const HloInstruction* const to_analyze = @@ -336,7 +330,7 @@ std::vector FindConstrainedUses( StatusOr CreateLiteralForConstrainedUses( const absl::Span constrained_uses, const HloInstruction& param, std::minstd_rand0* engine) { - std::vector index_space; + int64 index_bound = INT64_MAX; bool no_duplicates = false; bool needs_constant = false; ConstantType constant_type = ConstantType::kUnknown; @@ -348,19 +342,16 @@ StatusOr CreateLiteralForConstrainedUses( const Shape& slice_shape = use->opcode() == HloOpcode::kDynamicSlice ? use->shape() : use->operand(1)->shape(); - const int64 rank = indexed_shape.rank(); - if (!index_space.empty()) { - TF_RET_CHECK(rank == index_space.size()); - for (int64 i = 0; i < rank; ++i) { - index_space[i] = std::min( - index_space[i], ShapeUtil::GetDimension(indexed_shape, i) - - ShapeUtil::GetDimension(slice_shape, i)); - } - } else { - index_space.resize(rank); - for (int64 i = 0; i < rank; ++i) { - index_space[i] = ShapeUtil::GetDimension(indexed_shape, i) - - ShapeUtil::GetDimension(slice_shape, i); + const int64 first_index = + Cast(use)->first_index_operand_number(); + for (int64 operand = first_index; operand < use->operand_count(); + ++operand) { + if (use->operand(operand) == ¶m) { + index_bound = std::min( + index_bound, + ShapeUtil::GetDimension(indexed_shape, operand - first_index) - + ShapeUtil::GetDimension(slice_shape, + operand - first_index)); } } break; @@ -388,16 +379,13 @@ StatusOr CreateLiteralForConstrainedUses( } int constraint_count = 0; constraint_count += no_duplicates ? 1 : 0; - constraint_count += !index_space.empty() ? 1 : 0; + constraint_count += (index_bound != INT64_MAX) ? 1 : 0; constraint_count += needs_constant ? 1 : 0; if (constraint_count > 1) { return Unimplemented("Conflicting operand generation constraints."); } - if (!index_space.empty()) { - // constrained_uses looks through bitcasts, so param and indexed_space may - // not have the same shape. (For example, param might be an R0 while - // indexed_space might have size 1.) - return MakeRandomIndex(index_space, engine) + if (index_bound != INT64_MAX) { + return MakeRandomIndex(index_bound, engine) .Reshape(param.shape().dimensions()); } else if (needs_constant) { switch (constant_type) { diff --git a/tensorflow/compiler/xla/tests/test_utils_test.cc b/tensorflow/compiler/xla/tests/test_utils_test.cc index 4b4f164d1b43382c695dc8bca7ba21a97162d611..321c3fb2df6f0beccded4617e91eff69c2bce2ea 100644 --- a/tensorflow/compiler/xla/tests/test_utils_test.cc +++ b/tensorflow/compiler/xla/tests/test_utils_test.cc @@ -79,25 +79,26 @@ XLA_TEST_F(TestUtilsTest, MultipleIndexSpacesForDynamicSlices) { R"(HloModule index_space_module ENTRY IndexSpace { - index_param = s32[3]{0} parameter(0) - array_param.1 = f32[123,4,789]{0,1,2} parameter(1) - array_param.2 = f32[3,3000,5]{0,1,2} parameter(2) - dynamic-slice.1 = f32[1,2,3] dynamic-slice(array_param.1, index_param), dynamic_slice_sizes={1,2,3} - ROOT dynamic-slice.2 = f32[3,2,2] dynamic-slice(array_param.2, index_param), dynamic_slice_sizes={3,2,2} + index_param.0 = s32[] parameter(0) + index_param.1 = s32[] parameter(1) + index_param.2 = s32[] parameter(2) + array_param.1 = f32[123,4,789]{0,1,2} parameter(3) + array_param.2 = f32[3,3000,5]{0,1,2} parameter(4) + dynamic-slice.1 = f32[1,2,3] dynamic-slice(array_param.1, index_param.0, index_param.1, index_param.2), dynamic_slice_sizes={1,2,3} + ROOT dynamic-slice.2 = f32[3,2,2] dynamic-slice(array_param.2, index_param.0, index_param.1, index_param.2), dynamic_slice_sizes={3,2,2} })") .ValueOrDie(); TF_ASSERT_OK_AND_ASSIGN(std::vector args, MakeFakeArguments(module.get())); - ASSERT_EQ(args.size(), 3); - const Literal& index_arg = args[0]; + ASSERT_EQ(args.size(), 5); - EXPECT_EQ(index_arg.Get({0}), 0); + EXPECT_EQ(args[0].Get({}), 0); - EXPECT_GE(index_arg.Get({1}), 0); - EXPECT_LE(index_arg.Get({1}), 2); + EXPECT_GE(args[1].Get({}), 0); + EXPECT_LE(args[0].Get({}), 2); - EXPECT_GE(index_arg.Get({2}), 0); - EXPECT_LE(index_arg.Get({2}), 3); + EXPECT_GE(args[2].Get({}), 0); + EXPECT_LE(args[2].Get({}), 3); } XLA_TEST_F(TestUtilsTest, MultipleIndexSpacesForDynamicUpdateSlices) { @@ -105,28 +106,29 @@ XLA_TEST_F(TestUtilsTest, MultipleIndexSpacesForDynamicUpdateSlices) { R"(HloModule index_space_module ENTRY IndexSpace { - index_param = s32[3]{0} parameter(0) - array_param.1 = f32[123,4,789]{0,1,2} parameter(1) - array_param.2 = f32[3,3000,5]{0,1,2} parameter(2) - update_param.1 = f32[1,2,3]{0,1,2} parameter(3) - update_param.2 = f32[3,2,2]{0,1,2} parameter(4) - - dynamic-update-slice.1 = f32[123,4,789] dynamic-update-slice(array_param.1, update_param.1, index_param) - ROOT dynamic-update-slice.2 = f32[3,3000,5] dynamic-update-slice(array_param.2, update_param.2, index_param) + index_param.0 = s32[] parameter(0) + index_param.1 = s32[] parameter(1) + index_param.2 = s32[] parameter(2) + array_param.1 = f32[123,4,789]{0,1,2} parameter(3) + array_param.2 = f32[3,3000,5]{0,1,2} parameter(4) + update_param.1 = f32[1,2,3]{0,1,2} parameter(5) + update_param.2 = f32[3,2,2]{0,1,2} parameter(6) + + dynamic-update-slice.1 = f32[123,4,789] dynamic-update-slice(array_param.1, update_param.1, index_param.0, index_param.1, index_param.2) + ROOT dynamic-update-slice.2 = f32[3,3000,5] dynamic-update-slice(array_param.2, update_param.2, index_param.0, index_param.1, index_param.2) })") .ValueOrDie(); TF_ASSERT_OK_AND_ASSIGN(std::vector args, MakeFakeArguments(module.get())); - ASSERT_EQ(args.size(), 5); - const Literal& index_arg = args[0]; + ASSERT_EQ(args.size(), 7); - EXPECT_EQ(index_arg.Get({0}), 0); + EXPECT_EQ(args[0].Get({}), 0); - EXPECT_GE(index_arg.Get({1}), 0); - EXPECT_LE(index_arg.Get({1}), 2); + EXPECT_GE(args[1].Get({}), 0); + EXPECT_LE(args[0].Get({}), 2); - EXPECT_GE(index_arg.Get({2}), 0); - EXPECT_LE(index_arg.Get({2}), 3); + EXPECT_GE(args[2].Get({}), 0); + EXPECT_LE(args[2].Get({}), 3); } XLA_TEST_F(TestUtilsTest, NoDuplicatesFloats) { @@ -134,10 +136,18 @@ XLA_TEST_F(TestUtilsTest, NoDuplicatesFloats) { auto module = ParseHloString(R"( HloModule sort.148.1589 +compare { + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + ROOT lt = pred[] less-than(p.0.lhs, p.0.rhs) +} + ENTRY %sort.148.1589 (parameter.0: f32[1048576], parameter.1: s32[1048576]) -> (f32[1048576], s32[1048576]) { %parameter.0 = f32[1048576]{0} parameter(0) %parameter.1 = s32[1048576]{0} parameter(1) - ROOT %sort.148.1589 = (f32[1048576]{0}, s32[1048576]{0}) sort(f32[1048576]{0} %parameter.0, s32[1048576]{0} %parameter.1), dimensions={0} + ROOT %sort.148.1589 = (f32[1048576]{0}, s32[1048576]{0}) sort(f32[1048576]{0} %parameter.0, s32[1048576]{0} %parameter.1), dimensions={0}, to_apply=compare } )") .ValueOrDie(); @@ -157,10 +167,18 @@ XLA_TEST_F(TestUtilsTest, NoDuplicatesInt32) { auto module = ParseHloString(R"( HloModule sort.148.1589 +compare { + p.0.lhs = s32[] parameter(0) + p.0.rhs = s32[] parameter(1) + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + ROOT lt = pred[] less-than(p.0.lhs, p.0.rhs) +} + ENTRY %sort.148.1589 (parameter.0: s32[1048576], parameter.1: s32[1048576]) -> (s32[1048576], s32[1048576]) { %parameter.0 = s32[1048576]{0} parameter(0) %parameter.1 = s32[1048576]{0} parameter(1) - ROOT %sort.148.1589 = (s32[1048576]{0}, s32[1048576]{0}) sort(s32[1048576]{0} %parameter.0, s32[1048576]{0} %parameter.1), dimensions={0} + ROOT %sort.148.1589 = (s32[1048576]{0}, s32[1048576]{0}) sort(s32[1048576]{0} %parameter.0, s32[1048576]{0} %parameter.1), dimensions={0}, to_apply=compare } )") .ValueOrDie(); @@ -180,10 +198,18 @@ XLA_TEST_F(TestUtilsTest, NoDuplicatesBfloat16) { auto module = ParseHloString(R"( HloModule sort, is_scheduled=true +compare { + p.0.lhs = bf16[] parameter(0) + p.0.rhs = bf16[] parameter(1) + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + ROOT lt = pred[] less-than(p.0.lhs, p.0.rhs) +} + ENTRY %sort. (parameter.0: bf16[2,1452], parameter.1: s32[2,1452]) -> (bf16[2,1452], s32[2,1452]) { %parameter.0 = bf16[2,1452]{1,0} parameter(0) %parameter.1 = s32[2,1452]{1,0} parameter(1) - ROOT %sort = (bf16[2,1452]{1,0}, s32[2,1452]{1,0}) sort(bf16[2,1452]{1,0} %parameter.0, s32[2,1452]{1,0} %parameter.1), dimensions={1} + ROOT %sort = (bf16[2,1452]{1,0}, s32[2,1452]{1,0}) sort(bf16[2,1452]{1,0} %parameter.0, s32[2,1452]{1,0} %parameter.1), dimensions={1}, to_apply=compare } )") .ValueOrDie(); diff --git a/tensorflow/compiler/xla/client/lib/triangular_solve_test.cc b/tensorflow/compiler/xla/tests/triangular_solve_test.cc similarity index 76% rename from tensorflow/compiler/xla/client/lib/triangular_solve_test.cc rename to tensorflow/compiler/xla/tests/triangular_solve_test.cc index 703227c94944feb6858de9464758e024c55b323d..24ab12136ff396bd9ac37bb058311b0d2d6f2515 100644 --- a/tensorflow/compiler/xla/client/lib/triangular_solve_test.cc +++ b/tensorflow/compiler/xla/tests/triangular_solve_test.cc @@ -13,8 +13,6 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/xla/client/lib/triangular_solve.h" - #include #include #include @@ -54,6 +52,20 @@ Array2D AValsUpper() { {kNan, kNan, kNan, 11}}; } +Array2D AValsLowerUnitDiagonal() { + return {{kNan, kNan, kNan, kNan}, + {3, kNan, kNan, kNan}, + {4, 7, kNan, kNan}, + {5, 8, 10, kNan}}; +} + +Array2D AValsUpperUnitDiagonal() { + return {{kNan, 3, 4, 5}, + {kNan, kNan, 7, 8}, + {kNan, kNan, kNan, 10}, + {kNan, kNan, kNan, kNan}}; +} + Array2D BValsRight() { return {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}; } @@ -96,8 +108,8 @@ XLA_TEST_F(TriangularSolveTest, EmptyArrays) { CreateR2Parameter(Array2D(0, 10), 1, "b", &builder, &b); TriangularSolve(a, b, /*left_side=*/true, /*lower=*/true, - /*transpose_a=*/true, /*conjugate_a=*/false, - /*block_size=*/2); + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::TRANSPOSE); ComputeAndCompareR2(&builder, Array2D(0, 10), {a_data.get(), b_data.get()}); @@ -111,8 +123,8 @@ XLA_TEST_F(TriangularSolveTest, SimpleRightLowerTranspose) { auto b_data = CreateR2Parameter(BValsRight(), 1, "b", &builder, &b); TriangularSolve(a, b, /*left_side=*/false, /*lower=*/true, - /*transpose_a=*/true, /*conjugate_a=*/false, - /*block_size=*/2); + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::TRANSPOSE); Array2D expected({ {0.5, 0.08333334, 0.04629629, 0.03367003}, @@ -132,8 +144,8 @@ XLA_TEST_F(TriangularSolveTest, SimpleRightLowerNotranspose) { auto b_data = CreateR2Parameter(BValsRight(), 1, "b", &builder, &b); TriangularSolve(a, b, /*left_side=*/false, /*lower=*/true, - /*transpose_a=*/false, /*conjugate_a=*/false, - /*block_size=*/2); + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::NO_TRANSPOSE); Array2D expected({ {-0.16414141, -0.06902357, -0.07070707, 0.36363636}, @@ -153,8 +165,8 @@ XLA_TEST_F(TriangularSolveTest, SimpleRightUpperTranspose) { auto b_data = CreateR2Parameter(BValsRight(), 1, "b", &builder, &b); TriangularSolve(a, b, /*left_side=*/false, /*lower=*/false, - /*transpose_a=*/true, /*conjugate_a=*/false, - /*block_size=*/2); + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::TRANSPOSE); Array2D expected({ {-0.16414141, -0.06902357, -0.07070707, 0.36363636}, @@ -174,8 +186,8 @@ XLA_TEST_F(TriangularSolveTest, SimpleRightUpperNotranspose) { auto b_data = CreateR2Parameter(BValsRight(), 1, "b", &builder, &b); TriangularSolve(a, b, /*left_side=*/false, /*lower=*/false, - /*transpose_a=*/false, /*conjugate_a=*/false, - /*block_size=*/2); + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::NO_TRANSPOSE); Array2D expected({ {0.5, 0.08333334, 0.04629629, 0.03367003}, @@ -195,8 +207,8 @@ XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerTranspose) { auto b_data = CreateR2Parameter(BValsLeft(), 1, "b", &builder, &b); TriangularSolve(a, b, /*left_side=*/true, /*lower=*/true, - /*transpose_a=*/true, /*conjugate_a=*/false, - /*block_size=*/2); + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::TRANSPOSE); Array2D expected({ {-0.89646465, -0.69444444, -0.49242424}, @@ -217,8 +229,8 @@ XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerNotranspose) { auto b_data = CreateR2Parameter(BValsLeft(), 1, "b", &builder, &b); TriangularSolve(a, b, /*left_side=*/true, /*lower=*/true, - /*transpose_a=*/false, /*conjugate_a=*/false, - /*block_size=*/2); + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::NO_TRANSPOSE); Array2D expected({ {0.5, 1.0, 1.5}, @@ -231,6 +243,25 @@ XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerNotranspose) { ErrorSpec(1e-2, 1e-2)); } +XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerNoTransposeUnitDiagonal) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = + CreateR2Parameter(AValsLowerUnitDiagonal(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter(BValsLeft(), 1, "b", &builder, &b); + TriangularSolve(a, b, + /*left_side=*/true, /*lower=*/true, + /*unit_diagonal=*/true, + /*transpose_a=*/TriangularSolveOptions::NO_TRANSPOSE); + + Array2D expected( + {{1., 2., 3.}, {1., -1., -3.}, {-4., 7., 18.}, {37., -61., -159.}}); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerNotransposeIrregularblock) { XlaBuilder builder(TestName()); @@ -239,8 +270,8 @@ XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerNotransposeIrregularblock) { auto b_data = CreateR2Parameter(BValsLeft(), 1, "b", &builder, &b); TriangularSolve(a, b, /*left_side=*/true, /*lower=*/true, - /*transpose_a=*/false, /*conjugate_a=*/false, - /*block_size=*/3); + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::NO_TRANSPOSE); Array2D expected({ {0.5, 1.0, 1.5}, @@ -261,8 +292,8 @@ XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperTranspose) { auto b_data = CreateR2Parameter(BValsLeft(), 1, "b", &builder, &b); TriangularSolve(a, b, /*left_side=*/true, /*lower=*/false, - /*transpose_a=*/true, /*conjugate_a=*/false, - /*block_size=*/2); + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::TRANSPOSE); Array2D expected({ {0.5, 1.0, 1.5}, @@ -283,8 +314,8 @@ XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperNotranspose) { auto b_data = CreateR2Parameter(BValsLeft(), 1, "b", &builder, &b); TriangularSolve(a, b, /*left_side=*/true, /*lower=*/false, - /*transpose_a=*/false, /*conjugate_a=*/false, - /*block_size=*/2); + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::NO_TRANSPOSE); Array2D expected({ {-0.89646465, -0.69444444, -0.49242424}, @@ -297,6 +328,27 @@ XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperNotranspose) { ErrorSpec(1e-2, 1e-2)); } +XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperNotransposeUnitDiagonal) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = + CreateR2Parameter(AValsUpperUnitDiagonal(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter(BValsLeft(), 1, "b", &builder, &b); + TriangularSolve(a, b, + /*left_side=*/true, /*lower=*/false, + /*unit_diagonal=*/true, + /*transpose_a=*/TriangularSolveOptions::NO_TRANSPOSE); + + Array2D expected({{-1402., -1538., -1674.}, + {575., 631., 687.}, + {-93., -102., -111.}, + {10., 11., 12.}}); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + XLA_TEST_F(TriangularSolveTest, SimpleRightLowerTransposeConjugate) { XlaBuilder builder(TestName()); @@ -307,8 +359,8 @@ XLA_TEST_F(TriangularSolveTest, SimpleRightLowerTransposeConjugate) { CreateR2Parameter(BValsRightComplex(), 1, "b", &builder, &b); TriangularSolve(a, b, /*left_side=*/false, /*lower=*/true, - /*transpose_a=*/true, /*conjugate_a=*/true, - /*block_size=*/2); + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::ADJOINT); Array2D expected({ {0.5, complex64(0.08333333, 0.08333333), @@ -333,8 +385,8 @@ XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperTransposeNoconjugate) { CreateR2Parameter(BValsLeftComplex(), 1, "b", &builder, &b); TriangularSolve(a, b, /*left_side=*/true, /*lower=*/false, - /*transpose_a=*/true, /*conjugate_a=*/false, - /*block_size=*/2); + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::TRANSPOSE); Array2D expected({ {0.5, 1., 1.5}, @@ -368,11 +420,12 @@ XLA_TEST_F(TriangularSolveTest, BatchedLeftUpper) { XlaOp a, b; auto a_data = CreateR3Parameter(avals, 0, "a", &builder, &a); auto b_data = CreateR3Parameter(bvals, 1, "b", &builder, &b); - BatchDot(ConstantR3FromArray3D(&builder, avals), - TriangularSolve(a, b, - /*left_side=*/true, /*lower=*/false, - /*transpose_a=*/false, /*conjugate_a=*/false, - /*block_size=*/2)); + BatchDot( + ConstantR3FromArray3D(&builder, avals), + TriangularSolve(a, b, + /*left_side=*/true, /*lower=*/false, + /*unit_diagonal=*/false, + /*transpose_a=*/TriangularSolveOptions::NO_TRANSPOSE)); ComputeAndCompareR3(&builder, bvals, {a_data.get(), b_data.get()}, ErrorSpec(1e-2, 1e-2)); @@ -382,7 +435,7 @@ struct TriangularSolveTestSpec { int m, n; // A is mxm, B is mxn bool left_side; bool lower; - bool transpose_a; + TriangularSolveOptions::Transpose transpose_a; }; class TriangularSolveParametricTest @@ -408,11 +461,11 @@ XLA_TEST_P(TriangularSolveParametricTest, Random) { XlaOp a, b; auto a_data = CreateR2Parameter(avals, 0, "a", &builder, &a); auto b_data = CreateR2Parameter(bvals, 1, "b", &builder, &b); - auto x = TriangularSolve(a, b, spec.left_side, spec.lower, spec.transpose_a, - /*conjugate_a=*/false, - /*block_size=*/3); + auto x = TriangularSolve(a, b, spec.left_side, spec.lower, + /*unit_diagonal=*/false, spec.transpose_a); auto a_tri = Triangle(a, spec.lower); - a_tri = MaybeTransposeInMinorDims(a_tri, spec.transpose_a); + a_tri = MaybeTransposeInMinorDims( + a_tri, spec.transpose_a != TriangularSolveOptions::NO_TRANSPOSE); if (spec.left_side) { BatchDot(a_tri, x); } else { @@ -429,7 +482,9 @@ std::vector TriangularSolveTests() { for (int n : {5, 10}) { for (bool left_side : {false, true}) { for (bool lower : {false, true}) { - for (bool transpose_a : {false, true}) { + for (TriangularSolveOptions::Transpose transpose_a : + {TriangularSolveOptions::NO_TRANSPOSE, + TriangularSolveOptions::TRANSPOSE}) { specs.push_back({m, n, left_side, lower, transpose_a}); } } @@ -439,9 +494,9 @@ std::vector TriangularSolveTests() { return specs; } -INSTANTIATE_TEST_CASE_P(TriangularSolveParametricTestInstantiation, - TriangularSolveParametricTest, - ::testing::ValuesIn(TriangularSolveTests())); +INSTANTIATE_TEST_SUITE_P(TriangularSolveParametricTestInstantiation, + TriangularSolveParametricTest, + ::testing::ValuesIn(TriangularSolveTests())); } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/tuple_test.cc b/tensorflow/compiler/xla/tests/tuple_test.cc index 9c586bdeb05afb7378e92caed1f3edc408e051bf..cdf2c34fcc3cc005e84626c39c8ab301a9040529 100644 --- a/tensorflow/compiler/xla/tests/tuple_test.cc +++ b/tensorflow/compiler/xla/tests/tuple_test.cc @@ -176,8 +176,9 @@ XLA_TEST_F(TupleTest, AddTupleElements) { {2.f, 4.f, 6.f}, // row 0 {5.f, 7.f, 9.f}, // row 1 }); - ASSERT_TRUE(ShapeUtil::ShapeIs(vector_shape, F32, {3})); - ASSERT_TRUE(ShapeUtil::ShapeIs(matrix_shape, F32, {/*y=*/2, /*x=*/3})); + ASSERT_TRUE(ShapeUtil::Equal(vector_shape, ShapeUtil::MakeShape(F32, {3}))); + ASSERT_TRUE(ShapeUtil::Equal(matrix_shape, + ShapeUtil::MakeShape(F32, {/*y=*/2, /*x=*/3}))); ComputeAndCompareR2(&builder, expected, {}, error_spec_); } @@ -512,8 +513,7 @@ XLA_TEST_F(TupleTest, ComplexTuples) { class TupleHloTest : public HloTestBase {}; -// Disabled on the interpreter because bitcast doesn't exist on the interpreter. -XLA_TEST_F(TupleHloTest, DISABLED_ON_INTERPRETER(BitcastAfterGTE)) { +XLA_TEST_F(TupleHloTest, BitcastAfterGTE) { const char* testcase = R"( HloModule m, is_scheduled=true @@ -525,9 +525,7 @@ XLA_TEST_F(TupleHloTest, DISABLED_ON_INTERPRETER(BitcastAfterGTE)) { ROOT tuple.4 = (f32[1,3]{1,0}) tuple(copy) } )"; - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param = LiteralUtil::MakeTupleOwned(LiteralUtil::CreateR1({1, 2, 3})); auto result = ExecuteNoHloPasses(std::move(module), {¶m}); @@ -559,9 +557,7 @@ XLA_TEST_F(TupleHloTest, ROOT outfeed = token[] outfeed(tuple, token0) } )"; - auto module = - HloRunner::CreateModuleFromString(testcase, GetDebugOptionsForTest()) - .ValueOrDie(); + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); auto param0 = LiteralUtil::CreateR1({1, 2}); auto param1 = LiteralUtil::CreateR1({2, 3}); auto param4 = LiteralUtil::CreateR0(false); diff --git a/tensorflow/compiler/xla/tests/unary_op_test.cc b/tensorflow/compiler/xla/tests/unary_op_test.cc index 4fbd7f2fb174ac899c1e3b23801986cb52db96a2..c51f30f3b5db95962a719ec226dd03f41142a782 100644 --- a/tensorflow/compiler/xla/tests/unary_op_test.cc +++ b/tensorflow/compiler/xla/tests/unary_op_test.cc @@ -64,7 +64,9 @@ class UnaryOpTest : public ClientLibraryTestBase { &builder, {-2, 25, 0, static_cast(-0.0), -123, inf(), -inf()}); Sign(arg); - ComputeAndCompareR1(&builder, {-1, 1, 0, 0, -1, 1, -1}, {}); + ComputeAndCompareR1( + &builder, + {-1, 1, static_cast(+0.0), static_cast(-0.0), -1, 1, -1}, {}); } template diff --git a/tensorflow/compiler/xla/tests/while_test.cc b/tensorflow/compiler/xla/tests/while_test.cc index 6d5f276e82087cedc356691b0ff08df24cec8d20..85212fa56d71088156d2f3edda17f71cdab56da2 100644 --- a/tensorflow/compiler/xla/tests/while_test.cc +++ b/tensorflow/compiler/xla/tests/while_test.cc @@ -861,7 +861,7 @@ XLA_TEST_F(WhileTest, WhileWithDynamicUpdateSlice) { // Update. auto update = ConvertElementType(Broadcast(out0, {2}), F32); // Starts = iteration * 2; - auto starts = Reshape(Mul(iteration, ConstantR0(&builder, 2)), {1}); + auto starts = Mul(iteration, ConstantR0(&builder, 2)); // UpdateSlice. auto out1 = DynamicUpdateSlice(input, update, starts); @@ -901,7 +901,7 @@ XLA_TEST_F(WhileTest, WhileWithDynamicUpdateSlice) { // Per backend the values generated can be different as the different backends // use different random number generators. // TODO(b/32240857): Extend test to verify outputs. -XLA_TEST_F(WhileTest, DISABLED_ON_INTERPRETER(WhileWithPrngScalarResult)) { +XLA_TEST_F(WhileTest, WhileWithPrngScalarResult) { auto v6s32 = ShapeUtil::MakeShape(S32, {6}); // Create a computation for the condition: repeat for count iterations. @@ -1146,7 +1146,7 @@ XLA_TEST_F(WhileTest, NestedWhileWithScalarResult) { // while (f(result).get<0>()) { // result = result + 1; // } -XLA_TEST_F(WhileTest, DISABLED_ON_INTERPRETER(WhileWithCallInsideCondition)) { +XLA_TEST_F(WhileTest, WhileWithCallInsideCondition) { auto result_shape = ShapeUtil::MakeShape(S32, {}); // Create a computation for the condition: repeat for 5 iterations. @@ -1299,9 +1299,9 @@ void BM_WhileLoop(int num_iters) { auto one = ConstantR0(&builder, 1.0); auto update = Broadcast(one, {1, 1024, 1024}); // Starts = iteration * 2; - auto starts = ConstantR1(&builder, {0, 0, 0}); + auto zero = ConstantR0(&builder, 0); // UpdateSlice. - auto out1 = DynamicUpdateSlice(input, update, starts); + auto out1 = DynamicUpdateSlice(input, update, {zero, zero, zero}); Tuple(&builder, {out0, out1}); body = builder.Build().ConsumeValueOrDie(); } diff --git a/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc b/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc index 1538f2afbafad53fa2f7a3fd1a4705e5ab55f536..7b7b8f5d02dc99607b30f898e18c5b448d421e07 100644 --- a/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc +++ b/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc @@ -40,8 +40,6 @@ limitations under the License. namespace xla { namespace { -namespace gtl = ::tensorflow::gtl; - class HloProfileTest : public ClientLibraryTestBase {}; struct ParsedProfileOutputLine { @@ -174,9 +172,8 @@ void ExecuteAndFetchProfile(string* profile_output, LocalClient* client, exec_run_options.set_allocator(backend->memory_allocator()); exec_run_options.set_intra_op_thread_pool( backend->eigen_intra_op_thread_pool_device()); - ServiceExecutableRunOptions run_options( - exec_run_options, /*borrow_stream=*/nullptr, - backend->eigen_intra_op_thread_pool()); + ServiceExecutableRunOptions run_options(exec_run_options, + /*borrow_stream=*/nullptr); std::vector args = {&lhs_arg, &rhs_arg}; TF_ASSERT_OK_AND_ASSIGN( auto execution_result, diff --git a/tensorflow/compiler/xla/text_literal_writer.cc b/tensorflow/compiler/xla/text_literal_writer.cc index 7289ae7df65e56652eeeb67e536e4c721d97d999..fc7949d889dc8ed9fac425982cc555a6c42a7f1d 100644 --- a/tensorflow/compiler/xla/text_literal_writer.cc +++ b/tensorflow/compiler/xla/text_literal_writer.cc @@ -24,7 +24,6 @@ limitations under the License. #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/types.h" diff --git a/tensorflow/compiler/xla/tools/BUILD b/tensorflow/compiler/xla/tools/BUILD index 99b32c19a52bf2a1f02047a1ceea626947d994fc..52fee4770ab940741723514d742e998b25765f24 100644 --- a/tensorflow/compiler/xla/tools/BUILD +++ b/tensorflow/compiler/xla/tools/BUILD @@ -29,33 +29,6 @@ tf_cc_binary( ], ) -cc_library( - name = "dumped_computation_to_graphviz_library", - srcs = ["dumped_computation_to_graphviz.cc"], - deps = [ - "//tensorflow/compiler/xla:debug_options_flags", - "//tensorflow/compiler/xla:statusor", - "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client", - "//tensorflow/compiler/xla/client:client_library", - "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/client:xla_computation", - "//tensorflow/compiler/xla/service", - "//tensorflow/compiler/xla/service:hlo_proto", - "//tensorflow/core:lib", - "@com_google_absl//absl/types:span", - ], -) - -tf_cc_binary( - name = "dumped_computation_to_graphviz", - deps = [ - ":dumped_computation_to_graphviz_library", - "//tensorflow/compiler/xla/service:interpreter_plugin", - ], -) - tf_cc_binary( name = "show_signature", srcs = ["show_signature.cc"], @@ -95,6 +68,7 @@ cc_library( "//tensorflow/compiler/xla/service:hlo_parser", "//tensorflow/compiler/xla/service:hlo_proto", "//tensorflow/compiler/xla/service/gpu:infeed_manager", + "//tensorflow/compiler/xla/service/gpu:outfeed_manager", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", @@ -281,3 +255,9 @@ tf_cc_binary( "@com_google_absl//absl/strings", ], ) + +sh_test( + name = "interactive_graphviz_build_only_test", + srcs = ["interactive_graphviz_test.sh"], + data = [":interactive_graphviz"], +) diff --git a/tensorflow/compiler/xla/tools/dumped_computation_to_graphviz.cc b/tensorflow/compiler/xla/tools/dumped_computation_to_graphviz.cc deleted file mode 100644 index b623556468fb4a5d96be614b6c067d5a1df51a6f..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/tools/dumped_computation_to_graphviz.cc +++ /dev/null @@ -1,84 +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. -==============================================================================*/ - -// Usage: dumped_computation_to_graphviz some_binary_snapshot_proto* -// -// Dumps a graphviz URL for a snapshot computation to the command line. -// -// some_binary_snapshot_proto is obtained by serializing the HloSnapshot from -// ServiceInterface::SnapshotComputation to disk. -// -// The GraphViz URL is placed into the log stderr, whereas computation -// statistics are printed on stdout (implementation note: getting computation -// statistics is how we trigger compilation to split out a GraphViz URL). - -#include -#include -#include - -#include "absl/types/span.h" -#include "tensorflow/compiler/xla/client/client.h" -#include "tensorflow/compiler/xla/client/client_library.h" -#include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/client/xla_computation.h" -#include "tensorflow/compiler/xla/debug_options_flags.h" -#include "tensorflow/compiler/xla/service/hlo.pb.h" -#include "tensorflow/compiler/xla/service/service.h" -#include "tensorflow/compiler/xla/statusor.h" -#include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" -#include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/init_main.h" -#include "tensorflow/core/platform/logging.h" - -namespace xla { -namespace tools { - -void RealMain(absl::Span args) { - Client* client = ClientLibrary::LocalClientOrDie(); - for (char* arg : args) { - HloSnapshot module; - TF_CHECK_OK( - tensorflow::ReadBinaryProto(tensorflow::Env::Default(), arg, &module)); - XlaComputation computation = - client->LoadSnapshot(module).ConsumeValueOrDie(); - DebugOptions debug_options = GetDebugOptionsFromFlags(); - debug_options.set_xla_generate_hlo_graph(".*"); - ComputationStats stats = - client->GetComputationStats(computation, debug_options) - .ConsumeValueOrDie(); - fprintf(stdout, ">>> %s :: %s\n", arg, stats.DebugString().c_str()); - } -} - -} // namespace tools -} // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::AppendDebugOptionsFlags(&flag_list); - xla::string usage = tensorflow::Flags::Usage(argv[0], flag_list); - const bool parse_result = tensorflow::Flags::Parse(&argc, argv, flag_list); - if (!parse_result) { - LOG(ERROR) << "\n" << usage; - return 2; - } - tensorflow::port::InitMain(argv[0], &argc, &argv); - - absl::Span args(argv, argc); - args.remove_prefix(1); // Pop off the binary name, argv[0] - xla::tools::RealMain(args); - return 0; -} diff --git a/tensorflow/compiler/xla/tools/dumped_computation_to_operation_list.cc b/tensorflow/compiler/xla/tools/dumped_computation_to_operation_list.cc index 4375e7c138c9e8d193feaa7a39d63946c4ea3086..df2d3d18b9ff86c0dd2047c2415527aeb1c1f154 100644 --- a/tensorflow/compiler/xla/tools/dumped_computation_to_operation_list.cc +++ b/tensorflow/compiler/xla/tools/dumped_computation_to_operation_list.cc @@ -31,7 +31,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/service.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/init_main.h" #include "tensorflow/core/platform/logging.h" diff --git a/tensorflow/compiler/xla/tools/dumped_computation_to_text.cc b/tensorflow/compiler/xla/tools/dumped_computation_to_text.cc index 723569862c7550387e95003e3a673743464b67b8..35bb82ca22f46d2cdeaac3b9a87b253efe9a07d9 100644 --- a/tensorflow/compiler/xla/tools/dumped_computation_to_text.cc +++ b/tensorflow/compiler/xla/tools/dumped_computation_to_text.cc @@ -26,7 +26,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/service.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/init_main.h" #include "tensorflow/core/platform/logging.h" diff --git a/tensorflow/compiler/xla/tools/hlo_extractor_test.cc b/tensorflow/compiler/xla/tools/hlo_extractor_test.cc index c187222a11ee721b006194a68620c58749707193..4beb099b330cadf4540944979f38681bae07103c 100644 --- a/tensorflow/compiler/xla/tools/hlo_extractor_test.cc +++ b/tensorflow/compiler/xla/tools/hlo_extractor_test.cc @@ -36,9 +36,8 @@ ENTRY %entry { } )"; - TF_ASSERT_OK_AND_ASSIGN( - std::unique_ptr hlo_module, - HloRunner::CreateModuleFromString(hlo_string, GetDebugOptionsForTest())); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, + ParseAndReturnVerifiedModule(hlo_string)); { auto extracted_module = @@ -75,9 +74,8 @@ ENTRY %entry { } )"; - TF_ASSERT_OK_AND_ASSIGN( - std::unique_ptr hlo_module, - HloRunner::CreateModuleFromString(hlo_string, GetDebugOptionsForTest())); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, + ParseAndReturnVerifiedModule(hlo_string)); { auto extracted_module = @@ -120,9 +118,8 @@ ENTRY %entry { } )"; - TF_ASSERT_OK_AND_ASSIGN( - std::unique_ptr hlo_module, - HloRunner::CreateModuleFromString(hlo_string, GetDebugOptionsForTest())); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, + ParseAndReturnVerifiedModule(hlo_string)); { auto extracted_module = diff --git a/tensorflow/compiler/xla/tools/interactive_graphviz.cc b/tensorflow/compiler/xla/tools/interactive_graphviz.cc index 6c90cde5a75a93837ee149fd9b5a60e6413c2ac4..0c7c078b9b9d30427cb01b8930bd012046d852d3 100644 --- a/tensorflow/compiler/xla/tools/interactive_graphviz.cc +++ b/tensorflow/compiler/xla/tools/interactive_graphviz.cc @@ -1,4 +1,4 @@ -/* 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. @@ -29,8 +29,7 @@ limitations under the License. #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" -#include "absl/strings/string_view_utils.h" -#include "absl/strings/util.h" +#include "absl/strings/str_split.h" #include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/service/compiler.h" @@ -56,7 +55,8 @@ bool ReadLine(const char *prompt, string *line) { return util::ReadLine(prompt, line); #else std::cout << prompt; - return std::getline(std::cin, *line); + std::getline(std::cin, *line); + return std::cin.good(); #endif } @@ -139,9 +139,10 @@ HloComputation* FindComputation(const HloModule& module, // Print a help message describing the various available commands. void DoHelpCommand() { std::cout << R"(Commands: - [] - Renders a neighborhood of nodes around . If - is not provided, the default value is )" + [] [/ +] + Renders a neighborhood of nodes around , without going + beyond the optional boundary instructions. If is not provided, + the default value is )" << kDefaultWidth << R"(. allpaths [] Renders a subset of all paths from one instruction to the other. Either @@ -391,9 +392,9 @@ void DisplayGraphHandle(const Options &opts, const string& handle) { std::cout << handle << std::endl; // If it is a url, try to open it up in the user's browser too. - if (strings::StartsWithIgnoreCase(handle, "http://") || - strings::StartsWithIgnoreCase(handle, "https://") || - strings::StartsWithIgnoreCase(handle, "file://")) { + if (absl::StartsWithIgnoreCase(handle, "http://") || + absl::StartsWithIgnoreCase(handle, "https://") || + absl::StartsWithIgnoreCase(handle, "file://")) { const char* browser_bin = opts.browser.empty() ? "/usr/bin/sensible-browser" : opts.browser.c_str(); tensorflow::SubProcess p; @@ -457,12 +458,6 @@ void DoAllPathsCommand(const Options& opts, const HloModule& module, // Plot a given instruction neighborhood or computation with graphviz. void DoPlotCommand(const Options& opts, const HloModule& module, const std::vector& tokens) { - if (tokens.size() > 2) { - std::cerr << R"(Illegal input. Enter e.g. "%fusion.1 42" or "%fusion.1".)" - << std::endl; - return; - } - string node_name = tokens[0]; // Find the node with the given name. @@ -475,16 +470,43 @@ void DoPlotCommand(const Options& opts, const HloModule& module, } uint64 graph_width = kDefaultWidth; - if (tokens.size() == 2) { + absl::flat_hash_set boundary; + if (tokens.size() >= 2) { if (comp) { std::cerr << "Can only use graph-size parameter with instructions, but " << node_name << " is a computation." << std::endl; return; } - if (!absl::SimpleAtoi(tokens[1], &graph_width)) { - std::cerr << "Can't parse '" << tokens[1] << "' as an integer." - << std::endl; - return; + + int bound_index = 1; + // Get the if present. + if (absl::SimpleAtoi(tokens[bound_index], &graph_width)) { + bound_index++; + } else { + // not found, need to reset graph_width. + graph_width = kDefaultWidth; + } + // Get the '/'. + if (bound_index < tokens.size()) { + // This token must be a '/'. + if (tokens[bound_index] != "/") { + std::cerr << "Expect a /, but get a '" << tokens[bound_index] << "'." + << std::endl; + return; + } + bound_index++; + } + // Get the boundary nodes. + while (bound_index < tokens.size()) { + string bnode_name = tokens[bound_index]; + const HloInstruction* binstr = FindInstruction(module, bnode_name); + if (!binstr) { + std::cerr << "Couldn't find HloInstruction named " << bnode_name << "." + << std::endl; + return; + } + boundary.insert(binstr); + bound_index++; } } @@ -496,7 +518,9 @@ void DoPlotCommand(const Options& opts, const HloModule& module, /*show_backend_config=*/show_backend_config)); } else { DisplayGraphHandle(opts, hlo_graph_dumper::DumpNeighborhoodAround( - *instr, graph_width, /*show_backend_config=*/show_backend_config)); + *instr, graph_width, + /*show_backend_config=*/show_backend_config, + /*boundary=*/boundary)); } } @@ -515,7 +539,7 @@ void InteractiveDumpGraphs(const Options& opts, const HloModule& module) { << std::endl; continue; } - std::vector tokens = strings::Split(line, ' '); + std::vector tokens = absl::StrSplit(line, ' ', absl::SkipEmpty()); if (tokens[0] == "quit" || tokens[0] == "exit") { break; } else if (tokens[0] == "help") { diff --git a/tensorflow/compiler/xla/tools/interactive_graphviz_test.sh b/tensorflow/compiler/xla/tools/interactive_graphviz_test.sh new file mode 100755 index 0000000000000000000000000000000000000000..b3e43aa7da062547fb5f187b885e997fc44bbb65 --- /dev/null +++ b/tensorflow/compiler/xla/tools/interactive_graphviz_test.sh @@ -0,0 +1,19 @@ +#! /bin/bash +# /* 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. +# ==============================================================================*/ + +# This is a placeholder for a compile-only test for intractive_graphviz tool. + +exit 0 diff --git a/tensorflow/compiler/xla/tools/replay_computation.cc b/tensorflow/compiler/xla/tools/replay_computation.cc index 75023f6edf541cdeb04698da1c4bc828b11b9047..d66561315b4ad7a5e3f1f7b1bc1e557b71da6705 100644 --- a/tensorflow/compiler/xla/tools/replay_computation.cc +++ b/tensorflow/compiler/xla/tools/replay_computation.cc @@ -51,6 +51,7 @@ limitations under the License. #include "tensorflow/compiler/xla/execution_options_util.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/service/gpu/infeed_manager.h" +#include "tensorflow/compiler/xla/service/gpu/outfeed_manager.h" #include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_parser.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -73,14 +74,24 @@ namespace { // fields. struct Options { string fake_infeed_shape; - bool generate_fake_infeed = false; + string fake_outfeed_shape; + + // generate_fake_infeed == true is a safe default: If the model has 0 or 1 + // infeeds, then it will work like normal. If the model has more than one + // infeed, it will be an error, but that wouldn't have worked anyway if you + // hadn't passed generate_fake_infeed. + // + // Same for generate_fake_outfeed. + bool generate_fake_infeed = true; + bool generate_fake_outfeed = true; + bool use_fake_data = false; bool print_result = true; int num_runs = 1; }; -std::unique_ptr CompileExecutable(const HloSnapshot& module, - LocalClient* client) { +StatusOr> CompileExecutable( + const HloSnapshot& module, LocalClient* client) { XlaComputation computation(module.hlo().hlo_module()); std::vector argument_layouts; argument_layouts.reserve( @@ -91,9 +102,86 @@ std::unique_ptr CompileExecutable(const HloSnapshot& module, argument_layouts.push_back(Shape(param)); argument_layout_ptrs.push_back(&argument_layouts.back()); } - return client - ->Compile(computation, argument_layout_ptrs, ExecutableBuildOptions()) - .ValueOrDie(); + ExecutableBuildOptions exec_build_options; + *exec_build_options.mutable_debug_options() = GetDebugOptionsFromFlags(); + return client->Compile(computation, argument_layout_ptrs, exec_build_options); +} + +absl::optional GetXfeedShape(bool is_infeed, + const HloModuleProto& module, + const Options& opts) { + std::vector xfeed_instrs; + for (const auto& comp : module.computations()) { + for (const auto& instruction : comp.instructions()) { + if (instruction.opcode() == HloOpcodeString(is_infeed + ? HloOpcode::kInfeed + : HloOpcode::kOutfeed)) { + xfeed_instrs.push_back(instruction); + } + } + } + + auto log_xfeed_instrs = [&] { + for (const auto& infeed : xfeed_instrs) { + LOG(ERROR) << " " << ShapeUtil::HumanString(Shape(infeed.shape())) << " " + << infeed.name(); + } + }; + + auto find_instruction_from_id_or_die = [&](int64 id) { + for (const auto& comp : module.computations()) { + for (const auto& instruction : comp.instructions()) { + if (instruction.id() == id) { + return instruction; + } + } + } + LOG(FATAL) << "No instruction with id " << id; + }; + + absl::optional xfeed_shape; + string xfeed_name = is_infeed ? "infeed" : "outfeed"; + string fake_xfeed_shape = + is_infeed ? opts.fake_infeed_shape : opts.fake_outfeed_shape; + bool generate_fake_xfeed = + is_infeed ? opts.generate_fake_infeed : opts.generate_fake_outfeed; + if (!fake_xfeed_shape.empty()) { + xfeed_shape = std::move(ParseShape(fake_xfeed_shape)).ValueOrDie(); + } else if (generate_fake_xfeed) { + CHECK_LT(xfeed_instrs.size(), 2) + << "--generate_fake_" << xfeed_name + << " only works if the model has 0 or 1 " << xfeed_name << " ops."; + if (xfeed_instrs.empty()) { + LOG(INFO) << "Not generating fake " << xfeed_name + << " shape; model has no " << xfeed_name << "s."; + } else if (xfeed_instrs.size() == 1) { + // kInfeed instructions should have a shape (buffer, token). kOutfeed + // instructions should have operand 0 of shape `buffer`. We want to xfeed + // just `buffer`. + xfeed_shape = is_infeed + ? Shape(xfeed_instrs.front().shape()).tuple_shapes(0) + : Shape(find_instruction_from_id_or_die( + xfeed_instrs.front().operand_ids(0)) + .shape()); + LOG(INFO) << "Generating fake " << xfeed_name << " with inferred shape: " + << ShapeUtil::HumanString(*xfeed_shape); + } else { + LOG(ERROR) << "--generate_fake_" << xfeed_name + << " only works if the model has 0 or 1 " << xfeed_name + << " ops, but this model has " << xfeed_instrs.size() + << " of them:"; + log_xfeed_instrs(); + LOG(FATAL) << "Can't run model with --generate_fake_infeed."; + } + } else if (!xfeed_instrs.empty()) { + LOG(ERROR) << "Model contains " << xfeed_instrs.size() << " " << xfeed_name + << " instruction(s), but neither --generate_fake_" << xfeed_name + << " nor --fake_" << xfeed_name + << "_shape was specified. Execution will likely hang."; + log_xfeed_instrs(); + } + + return xfeed_shape; } // Invokes the given computation passing arbitrary data for every (unbound) @@ -142,56 +230,37 @@ StatusOr ReplayComputation(const HloSnapshot& module, } } - bool provide_infeed = false; - Shape infeed_shape; - if (!opts.fake_infeed_shape.empty()) { - StatusOr shape_status = ParseShape(opts.fake_infeed_shape); - TF_CHECK_OK(shape_status.status()); - infeed_shape = std::move(shape_status).ValueOrDie(); - provide_infeed = true; - } else if (opts.generate_fake_infeed) { - for (const auto& comp : computation.proto().computations()) { - for (const auto& instruction : comp.instructions()) { - if (instruction.opcode() == HloOpcodeString(HloOpcode::kInfeed)) { - CHECK(!provide_infeed) - << "--generate_fake_infeed only works if the model has 0 or 1 " - "infeed ops, but this one has >= 2."; - provide_infeed = true; - // The kInfeed instruction should have a shape (buffer, token). We - // want to infeed just `buffer`. - infeed_shape = Shape(instruction.shape()).tuple_shapes(0); - LOG(INFO) << "Generating fake infeed shape for inferred shape: " - << ShapeUtil::HumanString(infeed_shape); - } - } - } - } - // We only instantiate the thread pool if the user has requested that a - // concurrent infeed occur via the fake_infeed_shape, or when - // --generate_fake_infeed is passed and there exists an infeed operation in - // the HloSnapshot. - absl::optional pool; - Literal data; - if (provide_infeed) { - data = std::move(MakeFakeLiteral(infeed_shape)).ValueOrDie(); + if (absl::optional infeed_shape = GetXfeedShape( + /*is_infeed=*/true, computation.proto(), opts)) { + auto infeed_data = std::make_shared( + std::move(MakeFakeLiteral(*infeed_shape)).ValueOrDie()); + xla::gpu::GetOrCreateInfeedManager() + ->RegisterBeforeGetNextDestinationCallback([infeed_data, client] { + TF_CHECK_OK(client->TransferToInfeed(*infeed_data)); + }); } - auto transfer_infeed = [&data, client]() { - TF_CHECK_OK(client->TransferToInfeed(data)); - }; - if (provide_infeed) { - pool.emplace(tensorflow::Env::Default(), "infeed", - /*num_threads=*/1); - pool->Schedule([transfer_infeed]() { - // There may be several infeed buffers needed, however we don't know how - // many. If we proactively transfer too many infeed buffers, we may run - // out of memory. If we transfer too few infeed buffers, the program will - // hang. Therefore, we register a callback that is called when the infeed - // becomes empty, and in this callback we will transfer another fake - // infeed. - auto infeed_manager = xla::gpu::GetOrCreateInfeedManager(); - infeed_manager->RegisterOnEmptyCallback(transfer_infeed); - transfer_infeed(); - }); + + absl::optional outfeed_thread_pool; + if (absl::optional outfeed_shape = GetXfeedShape( + /*is_infeed=*/false, computation.proto(), opts)) { + // For each an outfeed that runs, enqueue a task that will consume it. We + // need a thread pool because the act of running an outfeed blocks on there + // being a destination available, and the act of making a destination + // available blocks on there being outfeed data available. + outfeed_thread_pool.emplace(tensorflow::Env::Default(), "infeed", + /*num_threads=*/1); + auto consume_outfeed = [client, outfeed_shape] { + TF_CHECK_OK( + client->TransferFromOutfeedLocal(*outfeed_shape, /*device_ordinal=*/0) + .status()); + VLOG(1) << "Received outfeed data of shape " + << ShapeUtil::HumanStringWithLayout(*outfeed_shape); + }; + xla::gpu::GetOrCreateOutfeedManager() + ->RegisterBeforeGetNextDestinationCallback( + [consume_outfeed, &outfeed_thread_pool] { + outfeed_thread_pool->Schedule(consume_outfeed); + }); } // Do not attempt to run the executable if num_runs is less than 1. @@ -260,7 +329,10 @@ StatusOr ParseInputFile(const string& filename, fprintf(stderr, "%s: is not HloProto. Trying HLO text.\n", filename.c_str()); string contents; TF_RETURN_IF_ERROR(tensorflow::ReadFileToString(env, filename, &contents)); - StatusOr> module = ParseHloString(contents); + HloModuleConfig config; + config.set_debug_options(GetDebugOptionsFromFlags()); + StatusOr> module = + ParseHloString(contents, config); if (module.ok()) { *snapshot.mutable_hlo()->mutable_hlo_module() = module.ValueOrDie()->ToProto(); @@ -288,7 +360,7 @@ int RealMain(absl::Span args, const Options& opts) { // Compile all the modules in parallel. LOG(INFO) << "Compiling " << snapshots.size() << " modules in parallel."; - std::vector> executables; + std::vector>> executables; { // ThreadPool CHECK-fails if we give it 0 threads. tensorflow::thread::ThreadPool thread_pool( @@ -305,9 +377,16 @@ int RealMain(absl::Span args, const Options& opts) { LOG(INFO) << "Done compiling; now running the modules."; for (int64 i = 0; i < executables.size(); ++i) { - LocalExecutable* executable = executables[i].get(); + if (!executables[i].ok()) { + LOG(ERROR) << "Compilation failed: " << executables[i].status(); + exit_status = EXIT_FAILURE; + continue; + } + LocalExecutable* executable = executables[i].ValueOrDie().get(); + LOG(ERROR) << "Running iteration " << i; StatusOr result_status = ReplayComputation(snapshots[i], executable, client, opts); + LOG(ERROR) << "iteration complete."; if (!result_status.ok()) { fprintf(stderr, "%s: error: %s\n", args[i], result_status.status().ToString().c_str()); @@ -352,9 +431,14 @@ int main(int argc, char** argv) { "Number of times to run each computation"), tensorflow::Flag("fake_infeed_shape", &opts.fake_infeed_shape, "Shape of fake data to construct for (infinite) infeed"), + tensorflow::Flag("fake_outfeed_shape", &opts.fake_outfeed_shape, + "Shape of fake data to outfeed from computation"), tensorflow::Flag("generate_fake_infeed", &opts.generate_fake_infeed, - "Whether a fake infeed shape should be generated " - "derived from the computation"), + "Whether a fake infeed shape should be derived " + "from the computation"), + tensorflow::Flag("generate_fake_outfeed", &opts.generate_fake_outfeed, + "Whether a fake outfeed shape should be derived " + "from the computation"), }; xla::string usage = tensorflow::Flags::Usage(argv[0], flag_list); bool parse_ok = tensorflow::Flags::Parse(&argc, argv, flag_list); diff --git a/tensorflow/compiler/xla/tools/show_signature.cc b/tensorflow/compiler/xla/tools/show_signature.cc index cdf306dfd1027cf6022c5d8ae844b4308f580e8d..b80d0db8d812380d8144713109d1c05168713c77 100644 --- a/tensorflow/compiler/xla/tools/show_signature.cc +++ b/tensorflow/compiler/xla/tools/show_signature.cc @@ -37,7 +37,6 @@ limitations under the License. #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/init_main.h" #include "tensorflow/core/platform/logging.h" diff --git a/tensorflow/compiler/xla/types.h b/tensorflow/compiler/xla/types.h index b645acb700b0f168112a40c9c72b4669435f717d..daf678f69017b9eb86cbc464a1f33b434021901d 100644 --- a/tensorflow/compiler/xla/types.h +++ b/tensorflow/compiler/xla/types.h @@ -41,6 +41,7 @@ using ::tensorflow::uint32; using ::tensorflow::uint64; using complex64 = std::complex; +using complex128 = std::complex; using ::Eigen::half; diff --git a/tensorflow/compiler/xla/util.cc b/tensorflow/compiler/xla/util.cc index 68cab7387cf1576072f96878b50f07def6862d8b..bb8bbf57c4252b16836553334901a3c896a17f39 100644 --- a/tensorflow/compiler/xla/util.cc +++ b/tensorflow/compiler/xla/util.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "absl/container/inlined_vector.h" #include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" @@ -80,13 +81,9 @@ bool IsPermutation(absl::Span permutation, int64 rank) { if (rank != permutation.size()) { return false; } - std::vector output(permutation.size(), -1); - for (auto index : permutation) { - CHECK_GE(index, 0); - CHECK_LT(index, rank); - output[index] = 0; - } - return std::find(output.begin(), output.end(), -1) == output.end(); + absl::InlinedVector trivial_permutation(rank); + absl::c_iota(trivial_permutation, 0); + return absl::c_is_permutation(permutation, trivial_permutation); } std::vector InversePermutation( diff --git a/tensorflow/compiler/xla/util.h b/tensorflow/compiler/xla/util.h index 6722641e9d2c177440361e6f0d1f6c0804eb7cda..f2fd17dc99455a921bf875aad2a3661b4d456823 100644 --- a/tensorflow/compiler/xla/util.h +++ b/tensorflow/compiler/xla/util.h @@ -324,8 +324,7 @@ bool IsIdentityPermutation(absl::Span permutation); template int64 PositionInContainer(const Container& container, int64 value) { - return std::distance(container.begin(), - std::find(container.begin(), container.end(), value)); + return std::distance(container.begin(), absl::c_find(container, value)); } // Formats the container as a comma-separated string. StrAppend must support diff --git a/tensorflow/compiler/xla/window_util.cc b/tensorflow/compiler/xla/window_util.cc index 51c73b3d17e4c32d9a8a14d3055ab56f02922af3..e001cc35f9fcea2783b3952e825838af6bbece72 100644 --- a/tensorflow/compiler/xla/window_util.cc +++ b/tensorflow/compiler/xla/window_util.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include "absl/algorithm/container.h" #include "absl/strings/str_cat.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -137,25 +138,23 @@ bool HasPadding(const Window& window) { } bool HasSymmetricPadding(const Window& window) { - return std::all_of(window.dimensions().begin(), window.dimensions().end(), - [](const WindowDimension& dim) { - return dim.padding_low() == dim.padding_high(); - }); + return absl::c_all_of(window.dimensions(), [](const WindowDimension& dim) { + return dim.padding_low() == dim.padding_high(); + }); } bool HasSymmetricPadding(const PaddingConfig& padding_config) { - return std::all_of(padding_config.dimensions().begin(), - padding_config.dimensions().end(), - [](const PaddingConfig::PaddingConfigDimension& dim) { - return dim.edge_padding_low() == dim.edge_padding_high(); - }); + return absl::c_all_of(padding_config.dimensions(), + [](const PaddingConfig::PaddingConfigDimension& dim) { + return dim.edge_padding_low() == + dim.edge_padding_high(); + }); } bool HasNegativePadding(const Window& window) { - return std::any_of(window.dimensions().begin(), window.dimensions().end(), - [](const WindowDimension& dim) { - return dim.padding_low() < 0 || dim.padding_high() < 0; - }); + return absl::c_any_of(window.dimensions(), [](const WindowDimension& dim) { + return dim.padding_low() < 0 || dim.padding_high() < 0; + }); } bool HasBaseDilation(const Window& window) { @@ -190,10 +189,9 @@ bool AllOrNoneReversed(const Window& window) { return true; } bool reversed = window.dimensions()[0].window_reversal(); - return std::all_of(window.dimensions().begin(), window.dimensions().end(), - [&](const WindowDimension& dim) { - return dim.window_reversal() == reversed; - }); + return absl::c_all_of(window.dimensions(), [&](const WindowDimension& dim) { + return dim.window_reversal() == reversed; + }); } bool HasDilation(const Window& window) { diff --git a/tensorflow/compiler/xla/xla.bzl b/tensorflow/compiler/xla/xla.bzl index 1439f1bcc5cec39203a7cb4b1f8604e7349382c6..c743dfd32b3a1327af480cf32ae3cdeb08ee814e 100644 --- a/tensorflow/compiler/xla/xla.bzl +++ b/tensorflow/compiler/xla/xla.bzl @@ -1,30 +1,42 @@ """Wrapper around cc_proto_library used inside the XLA codebase.""" -load("//tensorflow/core:platform/default/build_config.bzl", - "cc_proto_library") -load("//tensorflow/core:platform/default/build_config_root.bzl", - "if_static") +load( + "//tensorflow/core:platform/default/build_config.bzl", + "cc_proto_library", +) +load( + "//tensorflow/core:platform/default/build_config_root.bzl", + "if_static", +) +load("//tensorflow:tensorflow.bzl", "if_cuda_is_configured") # xla_proto_library() is a convenience wrapper around cc_proto_library. -def xla_proto_library(name, srcs=[], deps=[], visibility=None, testonly=0, **kwargs): - if kwargs.get('use_grpc_plugin'): - kwargs['use_grpc_namespace'] = True - cc_proto_library(name=name, - srcs=srcs, - deps=deps, - cc_libs = if_static( - ["@protobuf_archive//:protobuf"], - otherwise=["@protobuf_archive//:protobuf_headers"], - ), - protoc="@protobuf_archive//:protoc", - testonly=testonly, - visibility=visibility, - **kwargs) +def xla_proto_library(name, srcs = [], deps = [], visibility = None, testonly = 0, **kwargs): + if kwargs.get("use_grpc_plugin"): + kwargs["use_grpc_namespace"] = True + cc_proto_library( + name = name, + srcs = srcs, + # Append well-known proto dep. As far as I know this is the only way + # for xla_proto_library to access google.protobuf.{Any,Duration,...}. + deps = deps + ["@protobuf_archive//:cc_wkt_protos"], + cc_libs = if_static( + ["@protobuf_archive//:protobuf"], + otherwise = ["@protobuf_archive//:protobuf_headers"], + ), + protoc = "@protobuf_archive//:protoc", + testonly = testonly, + visibility = visibility, + **kwargs + ) def xla_py_grpc_library(**kwargs): - # Note: we don't currently define any special targets for Python GRPC in OSS. - _ignore = kwargs - pass - + # Note: we don't currently define any special targets for Python GRPC in OSS. + _ignore = kwargs + pass ORC_JIT_MEMORY_MAPPER_TARGETS = [] + +# We link the GPU plugin into the XLA Python extension if CUDA is enabled. +def xla_python_default_plugins(): + return if_cuda_is_configured(["//tensorflow/compiler/xla/service:gpu_plugin"]) diff --git a/tensorflow/compiler/xla/xla.proto b/tensorflow/compiler/xla/xla.proto index e2d7b6ef4666c533951960fd3dcf6869ec2b52c5..92834dbb02cdcd6383ceec3ffd079834b163ee6a 100644 --- a/tensorflow/compiler/xla/xla.proto +++ b/tensorflow/compiler/xla/xla.proto @@ -265,6 +265,10 @@ message ExecutionOptions { // computation on. The computation will be partitioned across these devices. // If not provided, the default device will be chosen. repeated DeviceHandle device_handles = 5; + + // Number of replicas of the computation to run. If zero, uses the default + // number of replicas for the XLA service. + int32 num_replicas = 6; } message GetDeviceHandlesRequest { diff --git a/tensorflow/compiler/xla/xla_data.proto b/tensorflow/compiler/xla/xla_data.proto index d029669cbeab346dd01e573dba7a6f915b5364dd..4e127356a9fa7c921386c13c5ecd64af5ab19ed3 100644 --- a/tensorflow/compiler/xla/xla_data.proto +++ b/tensorflow/compiler/xla/xla_data.proto @@ -56,6 +56,7 @@ enum PrimitiveType { // Complex values of fixed width. 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 // sub-shapes. They are used for things like returning multiple values from a @@ -75,7 +76,7 @@ enum PrimitiveType { // primitive type will have empty dimensions and tuple_shapes fields. TOKEN = 17; - // Next = 18 + // Next = 19 } // Describes the padding configuration for Pad operation. The padding amount on @@ -367,6 +368,7 @@ message LiteralProto { repeated float f32s = 8; repeated double f64s = 9; 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 bytes f16s = 11; @@ -374,7 +376,7 @@ message LiteralProto { bytes u16s = 16; bytes s16s = 17; repeated int64 sparse_indices = 14; - // Next = 18 + // Next = 19 } message WindowDimension { @@ -543,6 +545,26 @@ enum RandomDistribution { // Next: 4 } +message TriangularSolveOptions { + // If true, solves ax = b. If false, solves xa = b. + bool left_side = 1; + + // If true, 'a' is lower triangular. If false, 'a' is upper triangular. + bool lower = 2; + + // If true, the diagonal elements of 'a' are assumed to be 1 and not accessed. + bool unit_diagonal = 3; + + // Should we transpose or use the adjoint of 'a'? + enum Transpose { + TRANSPOSE_INVALID = 0; + NO_TRANSPOSE = 1; // Don't transpose 'a'. + TRANSPOSE = 2; // Transpose 'a'. + ADJOINT = 3; // Complex conjugate and transpose 'a'. + }; + Transpose transpose_a = 4; +} + message OpSharding { enum Type { // This sharding is replicated across all devices (implies maximal, diff --git a/tensorflow/compiler/xrt/BUILD b/tensorflow/compiler/xrt/BUILD index 2dae746d034a1bf52e84de74dfb0c6e23aaed4d1..b2718c5c283358d98da175a8d3b21bb1f2b01c75 100644 --- a/tensorflow/compiler/xrt/BUILD +++ b/tensorflow/compiler/xrt/BUILD @@ -11,9 +11,15 @@ package( load( "//tensorflow:tensorflow.bzl", + "tf_custom_op_py_library", "tf_gen_op_libs", + "tf_gen_op_wrapper_py", ) load("//tensorflow/compiler/xla:xla.bzl", "xla_proto_library") +load( + "//tensorflow/core:platform/default/build_config.bzl", + "tf_proto_library_py", +) xla_proto_library( name = "xrt_proto", @@ -27,6 +33,12 @@ xla_proto_library( ], ) +tf_proto_library_py( + name = "xrt_proto", # bzl adds a _py suffix + srcs = ["xrt.proto"], + visibility = ["//visibility:public"], +) + cc_library( name = "xrt_utils", srcs = [ @@ -78,6 +90,25 @@ tf_gen_op_libs( ], ) +tf_gen_op_wrapper_py( + name = "xrt_ops_wrapper_py", + out = "xrt_ops.py", + deps = [ + ":xrt_compile_ops_op_lib", + ":xrt_execute_op_op_lib", + ":xrt_state_ops_op_lib", + ], +) + +tf_custom_op_py_library( + name = "xrt_ops", + kernels = ["//tensorflow/compiler/xrt/kernels:xrt_ops"], + visibility = ["//visibility:public"], + deps = [ + ":xrt_ops_wrapper_py", + ], +) + cc_library( name = "xrt_server", visibility = ["//visibility:public"], diff --git a/tensorflow/compiler/xrt/kernels/BUILD b/tensorflow/compiler/xrt/kernels/BUILD index 67f475846e5f16060c1080759b0acb4216c4e72b..1e325191bba828e3d5e4599f87dcf4f4d0674945 100644 --- a/tensorflow/compiler/xrt/kernels/BUILD +++ b/tensorflow/compiler/xrt/kernels/BUILD @@ -11,20 +11,15 @@ cc_library( name = "xrt_state_ops", hdrs = ["xrt_state_ops.h"], deps = [ + "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_compiler", - "//tensorflow/compiler/xla:debug_options_flags", "//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:client_library", - "//tensorflow/compiler/xla/client:compile_only_client", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/client:xla_computation", - "//tensorflow/compiler/xla/service:compiler", "//tensorflow/compiler/xla/service:computation_placer", - "//tensorflow/compiler/xla/service:hlo_proto", "//tensorflow/compiler/xrt:xrt_proto", "//tensorflow/compiler/xrt:xrt_utils", "//tensorflow/core:core_cpu_internal", @@ -55,14 +50,18 @@ cc_library( "//tensorflow/compiler/xla/client:xla_computation", "//tensorflow/compiler/xla/service:compiler", "//tensorflow/compiler/xla/service:computation_placer", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xrt:xrt_compile_ops_op_lib", + "//tensorflow/compiler/xrt:xrt_execute_op_op_lib", "//tensorflow/compiler/xrt:xrt_proto", + "//tensorflow/compiler/xrt:xrt_state_ops_op_lib", "//tensorflow/compiler/xrt:xrt_utils", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", - "//tensorflow/stream_executor:stream_executor_headers_lib", + "//tensorflow/stream_executor:stream_executor_headers", "@com_google_absl//absl/strings", ], alwayslink = 1, diff --git a/tensorflow/compiler/xrt/kernels/xrt_compile_ops.cc b/tensorflow/compiler/xrt/kernels/xrt_compile_ops.cc index 2ccdf0f02d840600d5e0649c4805e3672d4a1286..b791519c09758a4f4124c95add5351a9433ecb8f 100644 --- a/tensorflow/compiler/xrt/kernels/xrt_compile_ops.cc +++ b/tensorflow/compiler/xrt/kernels/xrt_compile_ops.cc @@ -68,9 +68,11 @@ class XRTCompileOp : public OpKernel { Status CompilationCacheKey(const xrt::XLAComputation& computation, string* key) { - string serialized; - TF_RET_CHECK(SerializeToStringDeterministic(computation, &serialized)); - uint64 fingerprint = Fingerprint64(serialized); + const size_t size = computation.ByteSizeLong(); + auto serialized = absl::make_unique(size); + TF_RET_CHECK( + SerializeToBufferDeterministic(computation, serialized.get(), size)); + uint64 fingerprint = Fingerprint64(absl::string_view(serialized.get(), size)); *key = absl::StrCat(fingerprint); return Status::OK(); } @@ -215,11 +217,6 @@ XRTReleaseCompilationRefOp::~XRTReleaseCompilationRefOp() = default; void XRTReleaseCompilationRefOp::Compute(OpKernelContext* ctx) { VLOG(1) << "XRTReleaseCompilationRefOp::Compute"; - const Tensor& key_tensor = ctx->input(0); - OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(key_tensor.shape()), - errors::Internal("computation key should be a string scalar")); - int64 uid = key_tensor.scalar()(); - ResourceMgr* rm; OP_REQUIRES_OK(ctx, XRTGenericDeviceAccessor::GetResourceManager(ctx, &rm)); @@ -230,9 +227,13 @@ void XRTReleaseCompilationRefOp::Compute(OpKernelContext* ctx) { kXRTCompilationCacheResourceName, &cache)); core::ScopedUnref cache_unref(cache); - OP_REQUIRES_OK(ctx, cache->Release(uid)); - - VLOG(2) << "Released computation handle " << uid; + const Tensor& keys_tensor = ctx->input(0); + auto flat_keys = keys_tensor.flat(); + for (int64 i = 0; i < flat_keys.size(); ++i) { + int64 key = flat_keys(i); + OP_REQUIRES_OK(ctx, cache->Release(key)); + VLOG(2) << "Released computation handle " << key; + } } } // namespace diff --git a/tensorflow/compiler/xrt/kernels/xrt_execute_op.cc b/tensorflow/compiler/xrt/kernels/xrt_execute_op.cc index 7f0ac123a57dc9de4992d50d4b2915ab1d2f1b1d..42ef88168af4b6f391ffc2e69ab4c4000d1cbee1 100644 --- a/tensorflow/compiler/xrt/kernels/xrt_execute_op.cc +++ b/tensorflow/compiler/xrt/kernels/xrt_execute_op.cc @@ -19,10 +19,10 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/service/computation_placer.h" +#include "tensorflow/compiler/xla/service/hlo_input_output_alias_config.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/compiler/xrt/xrt.pb.h" #include "tensorflow/compiler/xrt/xrt_compilation_cache.h" #include "tensorflow/compiler/xrt/xrt_device.h" @@ -228,6 +228,25 @@ Status XRTExecuteOp::DoWork(OpKernelContext* context) { TF_RETURN_IF_ERROR(XRTTupleAllocation::CreateFromBuffer( shaped_buffer, device_ref.backend(), device_ref.device_ordinal(), &output_tuple)); + + // The ScopedShapedBuffer returned by the executable Run() API, in case of + // input/output buffer aliasing, might have holes in it, which need to be + // filled using the proper input tuples buffers which are the source of + // aliasing. + const xla::HloInputOutputAliasConfig& input_output_alias = + executable->executable()->module().input_output_alias_config(); + auto alias_function = + [&](const xla::ShapeIndex& output_index, + const xla::HloInputOutputAliasConfig::Alias& alias) -> Status { + TF_RET_CHECK(alias.parameter_number < input_tuples.size()); + return alias.kind == xla::HloInputOutputAliasConfig::AliasKind::kUserAlias + ? output_tuple->AliasBufferFrom( + *input_tuples[alias.parameter_number], + alias.parameter_index, output_index) + : Status::OK(); + }; + TF_RETURN_IF_ERROR(input_output_alias.ForEachAliasWithStatus(alias_function)); + if (config_proto.return_exploded_tuple() && output_tuple->on_device_shape().IsTuple()) { int64 tuple_element_count = diff --git a/tensorflow/compiler/xrt/kernels/xrt_state_ops.cc b/tensorflow/compiler/xrt/kernels/xrt_state_ops.cc index 1a5bfac337baf773b84b92af5f88ef7a4c8ba81f..6a7f10652533920ba3fa48fba1d5161f7c4d4530 100644 --- a/tensorflow/compiler/xrt/kernels/xrt_state_ops.cc +++ b/tensorflow/compiler/xrt/kernels/xrt_state_ops.cc @@ -37,6 +37,17 @@ REGISTER_KERNEL_BUILDER(Name("XRTAllocate") .HostMemory("handle"), XRTAllocateOp); +REGISTER_KERNEL_BUILDER(Name("XRTAllocateFromTensor") + .Device(DEVICE_XLA_GPU) + .HostMemory("inputs") + .HostMemory("handle"), + XRTAllocateFromTensorOp); +REGISTER_KERNEL_BUILDER(Name("XRTAllocateFromTensor") + .Device(DEVICE_XLA_CPU) + .HostMemory("inputs") + .HostMemory("handle"), + XRTAllocateFromTensorOp); + REGISTER_KERNEL_BUILDER(Name("XRTSubTuple") .Device(DEVICE_XLA_GPU) .HostMemory("base_handle") diff --git a/tensorflow/compiler/xrt/kernels/xrt_state_ops.h b/tensorflow/compiler/xrt/kernels/xrt_state_ops.h index 2e2f3ff116a7b331df8dbd58a9fe40096f524140..e2c223b3dbb2311d0f42e1a36e316fd9d5f66040 100644 --- a/tensorflow/compiler/xrt/kernels/xrt_state_ops.h +++ b/tensorflow/compiler/xrt/kernels/xrt_state_ops.h @@ -19,10 +19,14 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XRT_KERNELS_XRT_STATE_OPS_H_ #define TENSORFLOW_COMPILER_XRT_KERNELS_XRT_STATE_OPS_H_ +#include #include #include +#include "tensorflow/compiler/tf2xla/literal_util.h" +#include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" @@ -30,6 +34,7 @@ limitations under the License. #include "tensorflow/compiler/xrt/xrt.pb.h" #include "tensorflow/compiler/xrt/xrt_device.h" #include "tensorflow/compiler/xrt/xrt_state.h" +#include "tensorflow/core/common_runtime/dma_helper.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/framework/tensor.h" @@ -200,6 +205,109 @@ class XRTAllocateOp : public OpKernel { } }; +// Op that allocates memory for a tensor (with optional layout) and transfers it +// to the device, returning an allocation handle. +template +class XRTAllocateFromTensorOp : public OpKernel { + public: + explicit XRTAllocateFromTensorOp(OpKernelConstruction* ctx) : OpKernel(ctx) { + bool make_tuple = false; + OP_REQUIRES_OK(ctx, ctx->GetAttr("shapes", &tf_shapes_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("dtypes", &dtypes_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("make_tuple", &make_tuple)); + if (ctx->HasAttr("layouts")) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("layouts", &minor_to_major_)); + } + OP_REQUIRES( + ctx, tf_shapes_.size() == dtypes_.size(), + errors::InvalidArgument("shapes and dtypes must be the same length")); + std::vector xla_shapes; + for (int i = 0; i < tf_shapes_.size(); i++) { + xla::Shape xla_shape; + OP_REQUIRES_OK( + ctx, TensorShapeToXLAShape(dtypes_[i], tf_shapes_[i], &xla_shape)); + xla_shapes.push_back(xla_shape); + } + if (xla_shapes.size() > 1 || make_tuple) { + shape_ = xla::ShapeUtil::MakeTupleShape(xla_shapes); + } else { + shape_.Swap(&xla_shapes.front()); + } + if (!minor_to_major_.empty()) { + xla::Shape shape_with_layouts; + OP_REQUIRES_OK(ctx, GetShapeWithLayout(shape_, minor_to_major_, + /*layout_func=*/nullptr, + &shape_with_layouts)); + shape_.Swap(&shape_with_layouts); + } + } + + ~XRTAllocateFromTensorOp() override = default; + XRTAllocateFromTensorOp(const XRTAllocateFromTensorOp&) = delete; + XRTAllocateFromTensorOp& operator=(const XRTAllocateFromTensorOp&) = delete; + + void Compute(OpKernelContext* ctx) override { + VLOG(1) << "XRTAllocateFromTensorOp::Compute"; + + OpInputList values; + OP_REQUIRES_OK(ctx, ctx->input_list("inputs", &values)); + OP_REQUIRES(ctx, values.size() == tf_shapes_.size(), + errors::InvalidArgument( + "Wrong number of inputs to XRTAllocateFromTensor: ", + values.size(), " vs. ", tf_shapes_.size())); + + std::vector tensors_data; + for (size_t i = 0; i < values.size(); ++i) { + const Tensor& input_tensor = values[i]; + OP_REQUIRES(ctx, input_tensor.dtype() == dtypes_[i], + errors::InvalidArgument( + "Input tensor type and input dtype do not match")); + // We allow the requested on-device shape to differ from the shape of the + // input tensor, as long as they have the same number of elements. + OP_REQUIRES( + ctx, + input_tensor.shape().num_elements() == tf_shapes_[i].num_elements(), + errors::InvalidArgument( + "Input tensor must have the number of elements specified " + "in the matching input shape: ", + input_tensor.shape().num_elements(), " vs. ", + tf_shapes_[i].num_elements(), " at index ", i)); + tensors_data.push_back( + static_cast(DMAHelper::base(&input_tensor))); + } + // Use the buffer straight out of the input tensors to create the literal. + xla::BorrowingLiteral literal = + shape_.IsTuple() ? xla::BorrowingLiteral(tensors_data, shape_) + : xla::BorrowingLiteral(tensors_data.front(), shape_); + ResourceMgr* rm; + OP_REQUIRES_OK(ctx, DeviceAccessor::GetResourceManager(ctx, &rm)); + + // We are guaranteed that the underlying device object won't be deleted out + // from under us, while the ScopedRef is live. + class DeviceAccessor::ScopedRef device_ref; + OP_REQUIRES_OK(ctx, DeviceAccessor::InitScopedRef(ctx, &device_ref)); + + XRTTupleAllocation* allocation; + OP_REQUIRES_OK(ctx, XRTTupleAllocation::CreateAndTransfer( + literal, device_ref.backend(), + device_ref.device_ordinal(), &allocation)); + + // Intern takes ownership of our reference to allocation. + int64 key; + OP_REQUIRES_OK(ctx, allocation->Intern(rm, &key)); + + Tensor output(DT_INT64, TensorShape({})); + output.scalar()() = key; + ctx->set_output(0, output); + } + + private: + std::vector tf_shapes_; + DataTypeVector dtypes_; + std::vector minor_to_major_; + xla::Shape shape_; +}; + // Op that takes a tuple handle input and returns a handle to a sub-tuple of the // input. template @@ -453,17 +561,17 @@ class XRTReleaseAllocationOp : public OpKernel { void Compute(OpKernelContext* ctx) override { VLOG(1) << "XRTReleaseAllocationOp::Compute"; - const Tensor& allocation_handle = ctx->input(0); - OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(allocation_handle.shape()), - errors::Internal("handle input should be an int64 scalar")); - int64 key = allocation_handle.scalar()(); - ResourceMgr* rm; OP_REQUIRES_OK(ctx, DeviceAccessor::GetResourceManager(ctx, &rm)); - OP_REQUIRES_OK(ctx, XRTTupleAllocation::DeleteFromResourceManager(rm, key)); - - VLOG(2) << "Released allocation handle " << key; + const Tensor& allocation_handle = ctx->input(0); + auto flat_keys = allocation_handle.flat(); + for (int64 i = 0; i < flat_keys.size(); ++i) { + int64 key = flat_keys(i); + OP_REQUIRES_OK(ctx, + XRTTupleAllocation::DeleteFromResourceManager(rm, key)); + VLOG(2) << "Released allocation handle " << key; + } } }; diff --git a/tensorflow/compiler/xrt/ops/xrt_compile_ops.cc b/tensorflow/compiler/xrt/ops/xrt_compile_ops.cc index 7b3b50c69559f6003a108fdf6a1325dbdbaa80a6..9dd964e5467cd855d67764a512e95a6a18f482e1 100644 --- a/tensorflow/compiler/xrt/ops/xrt_compile_ops.cc +++ b/tensorflow/compiler/xrt/ops/xrt_compile_ops.cc @@ -44,10 +44,10 @@ REGISTER_OP("XRTReleaseCompilationHandle") .SetShapeFn(tensorflow::shape_inference::NoOutputs) .Doc( R"( -Discards a computation from the compilation cache. The handle cannot be -subsequently used. +Discards one or more computation handles from the compilation cache. +The handle(s) cannot be subsequently used. -'handle' is an id returned from a XRTCompile Op. +'handle' is an ID (or vector of IDs) returned from a XRTCompile Op. )"); } // namespace tensorflow diff --git a/tensorflow/compiler/xrt/ops/xrt_state_ops.cc b/tensorflow/compiler/xrt/ops/xrt_state_ops.cc index fe6bee0dacf5dc2050613fc9ad34d3235b5a7b63..2e743fec4963a52ee1abf64525f26e3d89479670 100644 --- a/tensorflow/compiler/xrt/ops/xrt_state_ops.cc +++ b/tensorflow/compiler/xrt/ops/xrt_state_ops.cc @@ -26,12 +26,41 @@ REGISTER_OP("XRTAllocate") .SetShapeFn(tensorflow::shape_inference::ScalarShape) .Doc( R"( -Reads a literal proto and transfers it to TPU device memory. +Reads a literal proto and transfers it to device memory. -'allocation' is a serialized xrt::TPUAllocation proto. +'allocation' is a serialized xrt::XLAAllocation proto. 'handle' is an id that can be used in other ops to refer to the allocation. )"); +REGISTER_OP("XRTAllocateFromTensor") + .Input("inputs: dtypes") + .Output("handle: int64") + .Attr("dtypes: list(type)") + .Attr("shapes: list(shape)") + .Attr("layouts: list(int) = []") + .Attr("make_tuple: bool = false") + .SetShapeFn(tensorflow::shape_inference::ScalarShape) + .Doc( + R"( +Reads a list of tensors with optional layouts, and transfers it to device +memory. + +inputs: The tensors holding the input data. +shapes: The shapes which the tensors should have on device. The i-th shape +corresponds to the i-th input. The shapes, together with the (optional) +layouts, helps creating the fully qualified shape of the data on the device. +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. +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 +subshape can be set to -1, to leave such subshape to the default shape. +handle: An id that can be used in other ops to refer to the allocation. +)"); + REGISTER_OP("XRTSubTuple") .Input("base_handle: int64") .Input("shape_index: int32") @@ -127,10 +156,11 @@ REGISTER_OP("XRTReleaseAllocationHandle") .SetShapeFn(tensorflow::shape_inference::NoOutputs) .Doc( R"( -Discards an allocation from device memory. The handle cannot be subsequently +Discards one or more device memory handles. The handle(s) cannot be subsequently used. -'handle' is the id returned from the Op that produced the on-device allocation. +'handle' is the ID (or a vector of IDs) returned from the Op that produced the +on-device allocation. )"); REGISTER_OP("XRTReleaseAllAllocations") diff --git a/tensorflow/compiler/xrt/tests/BUILD b/tensorflow/compiler/xrt/tests/BUILD index be44a3474acdeb9905c1d21b932fa0dd10b5a212..3a19327e5b5d8072fbecdbe10e9959c8491780eb 100644 --- a/tensorflow/compiler/xrt/tests/BUILD +++ b/tensorflow/compiler/xrt/tests/BUILD @@ -24,6 +24,7 @@ cc_library( "//tensorflow/cc:client_session", "//tensorflow/cc:ops", "//tensorflow/cc:scope", + "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:literal_util", diff --git a/tensorflow/compiler/xrt/tests/raw_api_test.cc b/tensorflow/compiler/xrt/tests/raw_api_test.cc index 5f8121703e108f26b048feb7a0412a282f52892c..1111f8240512e81c10a42a28c09f5b0a94daf1ee 100644 --- a/tensorflow/compiler/xrt/tests/raw_api_test.cc +++ b/tensorflow/compiler/xrt/tests/raw_api_test.cc @@ -22,6 +22,8 @@ limitations under the License. #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/framework/scope.h" #include "tensorflow/cc/ops/standard_ops.h" +#include "tensorflow/compiler/tf2xla/literal_util.h" +#include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/client/xla_builder.h" @@ -53,6 +55,14 @@ string DeviceFromFlag() { return absl::StrCat("/device:", xla_test_device, ":0"); } +std::vector GetAttrLayout(absl::Span minor_to_mayor) { + std::vector layout; + for (auto dim : minor_to_mayor) { + layout.push_back(static_cast(dim)); + } + return layout; +} + xla::LiteralProto TwoElementTuple() { auto array = xla::LiteralUtil::CreateR1({1.0f, 3.0f}); auto matrix = xla::LiteralUtil::CreateR2({{4, 5}, {6, 7}}); @@ -96,14 +106,21 @@ xla::LiteralProto FloatMatrix( return array.ToProto(); } +xla::Literal ReadOutputLiteral(const std::vector& outputs, size_t idx) { + xla::LiteralProto response; + CHECK(response.ParseFromString(outputs[idx].scalar()())); + return xla::Literal::CreateFromProto(response).ValueOrDie(); +} + bool CompareLiteralProtos(const xla::LiteralProto& a, const xla::LiteralProto& b) { auto l_a = xla::Literal::CreateFromProto(a).ValueOrDie(); auto l_b = xla::Literal::CreateFromProto(b).ValueOrDie(); bool equal = l_a == l_b; if (!equal) { - LOG(INFO) << "LiteralProtos don't match: " << a.DebugString() - << " != " << b.DebugString(); + LOG(INFO) << "LiteralProtos don't match:\n" + << a.DebugString() << "\n!=\n" + << b.DebugString(); } return equal; } @@ -113,8 +130,19 @@ bool CompareLiteralToLiteralProto(const xla::Literal& a, auto l_b = xla::Literal::CreateFromProto(b).ValueOrDie(); bool equal = a == l_b; if (!equal) { - LOG(INFO) << "Literal and LiteralProto don't match " - << a.ToProto().DebugString() << " != " << b.DebugString(); + LOG(INFO) << "Literal and LiteralProto don't match:\n" + << a.ToProto().DebugString() << "\n!=\n" + << b.DebugString(); + } + return equal; +} + +bool CompareLiterals(const xla::Literal& a, const xla::Literal& b) { + bool equal = a == b; + if (!equal) { + LOG(INFO) << "Literals don't match:\n" + << a.ToProto().DebugString() << "\n!=\n" + << b.ToProto().DebugString(); } return equal; } @@ -215,6 +243,120 @@ xla::ProgramShape XlaCompiledProgramShape( ->ComputeProgramShape(); } +TEST(RawApiTest, AllocFromTensor) { + xla::Literal literal = + xla::LiteralUtil::CreateR2({{4.0f, 5.0f}, {6.0f, 7.0f}}); + Tensor tensor; + TF_ASSERT_OK(LiteralToHostTensor(literal, DT_FLOAT, &tensor)); + + Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); + std::vector layout = + GetAttrLayout(literal.shape().layout().minor_to_major()); + ops::XRTAllocateFromTensor::Attrs alloc_attrs = + ops::XRTAllocateFromTensor::Layouts(layout); + auto handle = + ops::XRTAllocateFromTensor(root, {tensor}, {tensor.shape()}, alloc_attrs); + auto read_back = ops::XRTReadLiteralAndRelease(root, handle); + TF_ASSERT_OK(root.status()); + + ClientSession session(root); + std::vector outputs; + TF_EXPECT_OK(session.Run({read_back}, &outputs)); + EXPECT_EQ(outputs.size(), 1); + + xla::LiteralProto response; + EXPECT_TRUE(response.ParseFromString(outputs[0].scalar()())); + EXPECT_TRUE(CompareLiteralToLiteralProto(literal, response)); +} + +TEST(RawApiTest, AllocFromTensorTuple) { + xla::Literal literal0 = + xla::LiteralUtil::CreateR2({{4.0f, 5.0f}, {6.0f, 7.0f}}); + xla::Literal literal1 = + xla::LiteralUtil::CreateR2({{14.0f, -5.0f}, {16.0f, 17.0f}}); + xla::Literal literal = xla::LiteralUtil::MakeTuple({&literal0, &literal1}); + Tensor tensor0; + TF_ASSERT_OK(LiteralToHostTensor(literal0, DT_FLOAT, &tensor0)); + Tensor tensor1; + TF_ASSERT_OK(LiteralToHostTensor(literal1, DT_FLOAT, &tensor1)); + + Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); + std::vector layout = GetShapeLayoutVector(literal.shape()).ValueOrDie(); + ops::XRTAllocateFromTensor::Attrs alloc_attrs = + ops::XRTAllocateFromTensor::Layouts(layout); + auto handle = ops::XRTAllocateFromTensor(root, {tensor0, tensor1}, + {tensor0.shape(), tensor1.shape()}, + alloc_attrs); + auto read_back = ops::XRTReadLiteralAndRelease(root, handle); + TF_ASSERT_OK(root.status()); + + ClientSession session(root); + std::vector outputs; + TF_EXPECT_OK(session.Run({read_back}, &outputs)); + EXPECT_EQ(outputs.size(), 1); + + xla::LiteralProto response; + EXPECT_TRUE(response.ParseFromString(outputs[0].scalar()())); + EXPECT_TRUE(CompareLiteralToLiteralProto(literal, response)); +} + +TEST(RawApiTest, AllocFromTensorTupleSingle) { + xla::Literal literal0 = + xla::LiteralUtil::CreateR2({{4.0f, 5.0f}, {6.0f, 7.0f}}); + xla::Literal literal = xla::LiteralUtil::MakeTuple({&literal0}); + Tensor tensor0; + TF_ASSERT_OK(LiteralToHostTensor(literal0, DT_FLOAT, &tensor0)); + + Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); + std::vector layout = GetShapeLayoutVector(literal.shape()).ValueOrDie(); + ops::XRTAllocateFromTensor::Attrs alloc_attrs = + ops::XRTAllocateFromTensor::Layouts(layout).MakeTuple(true); + auto handle = ops::XRTAllocateFromTensor(root, {tensor0}, {tensor0.shape()}, + alloc_attrs); + auto read_back = ops::XRTReadLiteralAndRelease(root, handle); + TF_ASSERT_OK(root.status()); + + ClientSession session(root); + std::vector outputs; + TF_EXPECT_OK(session.Run({read_back}, &outputs)); + EXPECT_EQ(outputs.size(), 1); + + xla::LiteralProto response; + EXPECT_TRUE(response.ParseFromString(outputs[0].scalar()())); + EXPECT_TRUE(CompareLiteralToLiteralProto(literal, response)); +} + +TEST(RawApiTest, AllocFromTensorRelayout) { + xla::Literal literal = + xla::LiteralUtil::CreateR2({{4.0f, 5.0f}, {6.0f, 7.0f}}); + Tensor tensor; + TF_ASSERT_OK(LiteralToHostTensor(literal, DT_FLOAT, &tensor)); + + Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); + // Use inverse array layout with the tensor data above. + std::vector layout({0, 1}); + ops::XRTAllocateFromTensor::Attrs alloc_attrs = + ops::XRTAllocateFromTensor::Layouts(layout); + auto handle = + ops::XRTAllocateFromTensor(root, {tensor}, {tensor.shape()}, alloc_attrs); + auto read_back = ops::XRTReadLiteralAndRelease(root, handle); + TF_ASSERT_OK(root.status()); + + ClientSession session(root); + std::vector outputs; + TF_EXPECT_OK(session.Run({read_back}, &outputs)); + EXPECT_EQ(outputs.size(), 1); + + xla::LiteralProto response; + EXPECT_TRUE(response.ParseFromString(outputs[0].scalar()())); + // We have sent literal's data (in array layout) with a attribute layout + // {0,1}, so the expected literal read from device needs to be changed + // accordingly. + xla::Literal expected_literal = + xla::LiteralUtil::CreateR2({{4.0f, 6.0f}, {5.0f, 7.0f}}); + EXPECT_TRUE(CompareLiteralToLiteralProto(expected_literal, response)); +} + TEST(RawApiTest, AllocAndRewrite) { xrt::XLAAllocation alloc; *alloc.mutable_value() = @@ -258,8 +400,102 @@ TEST(RawApiTest, AllocAndRewrite) { EXPECT_TRUE(new_response.ParseFromString(outputs[0].scalar()())); EXPECT_TRUE(CompareLiteralProtos(new_literal, new_response)); - auto release = - ops::XRTReleaseAllocationHandle(root, Input(allocation_handle)); + Tensor release_tensor(DT_INT64, TensorShape({1})); + release_tensor.flat()(0) = allocation_handle; + + auto release = ops::XRTReleaseAllocationHandle(root, release_tensor); + TF_EXPECT_OK(session.Run(tensorflow::ClientSession::FeedType(), {}, {release}, + &outputs)); +} + +TEST(RawApiTest, AllocReleaseMany) { + xrt::XLAAllocation alloc1; + *alloc1.mutable_value() = + xla::LiteralUtil::CreateR2({{4, 5}, {6, 7}}).ToProto(); + xrt::XLAAllocation alloc2; + *alloc2.mutable_value() = + xla::LiteralUtil::CreateR2({{6, 7}, {4, 5}}).ToProto(); + + Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); + auto value1 = + ops::Const(root.WithDevice("/device:CPU:0"), alloc1.SerializeAsString()); + auto value2 = + ops::Const(root.WithDevice("/device:CPU:0"), alloc2.SerializeAsString()); + auto handle1 = ops::XRTAllocate(root, value1); + auto handle2 = ops::XRTAllocate(root, value2); + TF_ASSERT_OK(root.status()); + + tensorflow::ClientSession session(root); + std::vector outputs; + TF_EXPECT_OK(session.Run({handle1, handle2}, &outputs)); + EXPECT_EQ(outputs.size(), 2); + + int64 allocation_handle1 = outputs[0].scalar()(); + int64 allocation_handle2 = outputs[1].scalar()(); + + Tensor release_tensor(DT_INT64, TensorShape({2})); + release_tensor.flat()(0) = allocation_handle1; + release_tensor.flat()(1) = allocation_handle2; + + auto release = ops::XRTReleaseAllocationHandle(root, release_tensor); + outputs.clear(); + TF_EXPECT_OK(session.Run(tensorflow::ClientSession::FeedType(), {}, {release}, + &outputs)); +} + +TEST(RawApiTest, CompileAndReleaseMany) { + xrt::XLAComputation c1; + auto config1 = c1.mutable_config(); + auto shapes1 = config1->mutable_program_shape(); + *shapes1->add_parameters() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); + *shapes1->add_parameters() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); + *shapes1->mutable_result() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); + StoreComputationSnapshot(AddAndScale(), c1.mutable_hlo_snapshot()); + + xrt::XLAComputation c2; + auto config2 = c2.mutable_config(); + auto shapes2 = config2->mutable_program_shape(); + *shapes2->add_parameters() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); + *shapes2->add_parameters() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); + *shapes2->mutable_result() = + xla::ShapeUtil::MakeTupleShape({xla::ShapeUtil::MakeShape(xla::F32, {2})}) + .ToProto(); + StoreComputationSnapshot(AddAndTuple(), c2.mutable_hlo_snapshot()); + + xrt::XRTExecutionConfig e; + e.set_release_input_handles(true); + e.set_release_compilation_handle(false); + + Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); + auto e_config = + ops::Const(root.WithDevice("/device:CPU:0"), e.SerializeAsString()); + auto computation1 = + ops::Const(root.WithDevice("/device:CPU:0"), c1.SerializeAsString()); + auto c_handle1 = ops::XRTCompile(root, computation1); + auto computation2 = + ops::Const(root.WithDevice("/device:CPU:0"), c2.SerializeAsString()); + auto c_handle2 = ops::XRTCompile(root, computation2); + TF_ASSERT_OK(root.status()); + + ClientSession session(root); + std::vector outputs; + TF_EXPECT_OK(session.Run({c_handle1.handle, c_handle2.handle}, &outputs)); + EXPECT_EQ(outputs.size(), 2); + + int64 compilation_handle1 = outputs[0].scalar()(); + int64 compilation_handle2 = outputs[1].scalar()(); + + Tensor release_tensor(DT_INT64, TensorShape({2})); + release_tensor.flat()(0) = compilation_handle1; + release_tensor.flat()(1) = compilation_handle2; + + auto release = ops::XRTReleaseCompilationHandle(root, release_tensor); + outputs.clear(); TF_EXPECT_OK(session.Run(tensorflow::ClientSession::FeedType(), {}, {release}, &outputs)); } @@ -845,6 +1081,107 @@ TEST(RawApiTest, LeakCompilationReference) { TF_EXPECT_OK(session.Run({c_handle.handle}, &outputs)); } +TEST(RawApiTest, CompileAndExecuteWithReusedBuffers) { + xla::Shape element_shape = xla::ShapeUtil::MakeShape(xla::F32, {2}); + xla::Shape shape = + xla::ShapeUtil::MakeTupleShape({element_shape, element_shape}); + xla::Shape return_shape = xla::ShapeUtil::MakeTupleShape( + {element_shape, element_shape, element_shape, element_shape}); + xla::XlaBuilder builder("ReuseBuffer"); + auto param = xla::Parameter(&builder, 0, shape, "param"); + auto p0 = xla::GetTupleElement(param, 0); + auto p1 = xla::GetTupleElement(param, 1); + auto add = xla::Add(p0, p1); + auto sub = xla::Sub(p0, p1); + xla::Tuple(&builder, {add, sub, p0, p1}); + + // Flip the tuple literals in the input handle. + builder.SetUpAlias({1}, 0, {0}); + builder.SetUpAlias({0}, 0, {1}); + + auto computation = builder.Build().ValueOrDie(); + + auto literal0 = xla::LiteralUtil::CreateR1({1.0f, 2.0f}); + auto literal1 = xla::LiteralUtil::CreateR1({5.0f, 9.0f}); + auto literal = xla::LiteralUtil::MakeTuple({&literal0, &literal1}); + + xrt::XLAAllocation param_alloc; + *param_alloc.mutable_value() = literal.ToProto(); + + xrt::XLAComputation c; + auto config = c.mutable_config(); + auto shapes = config->mutable_program_shape(); + *shapes->add_parameters() = shape.ToProto(); + *shapes->mutable_result() = return_shape.ToProto(); + StoreComputationSnapshot(computation, c.mutable_hlo_snapshot()); + + xrt::XRTExecutionConfig e; + e.set_release_input_handles(false); + e.set_release_compilation_handle(true); + + Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); + ClientSession session(root); + auto e_config = + ops::Const(root.WithDevice("/device:CPU:0"), e.SerializeAsString()); + auto c_data = + ops::Const(root.WithDevice("/device:CPU:0"), c.SerializeAsString()); + auto c_handle = ops::XRTCompile(root, c_data); + auto param_value = ops::Const(root.WithDevice("/device:CPU:0"), + param_alloc.SerializeAsString()); + auto param_handle = ops::XRTAllocate(root, param_value); + TF_ASSERT_OK(root.status()); + + std::vector outputs; + TF_EXPECT_OK(session.Run({param_handle}, &outputs)); + + int64 alloc_handle = outputs[0].scalar()(); + + // Note that we release the result handle immediately, but since we aliased + // the output buffers onto the input allocation ones (held in alloc_handle), + // we can fetch the result from there. + auto result = + ops::XRTExecute(root, c_handle.handle, e_config, {Input(alloc_handle)}); + auto read_back = ops::XRTReadLiteral(root, result); + auto release = ops::XRTReleaseAllocationHandle( + root.WithControlDependencies(read_back), result); + TF_ASSERT_OK(root.status()); + + outputs.clear(); + TF_EXPECT_OK(session.Run(tensorflow::ClientSession::FeedType(), {read_back}, + {release}, &outputs)); + + xla::Literal exec_literal = ReadOutputLiteral(outputs, 0); + auto exec_literal_parts = exec_literal.DecomposeTuple(); + ASSERT_EQ(exec_literal_parts.size(), 4); + + EXPECT_TRUE(CompareLiterals(exec_literal_parts[2], literal0)); + EXPECT_TRUE(CompareLiterals(exec_literal_parts[3], literal1)); + + // Now we read back the original input handle values, which at this point + // should contain the result of the XLA computation. + auto read_handle = ops::XRTReadLiteral(root, Input(alloc_handle)); + TF_ASSERT_OK(root.status()); + auto release_handle = ops::XRTReleaseAllocationHandle( + root.WithControlDependencies(read_handle), Input(alloc_handle)); + TF_ASSERT_OK(root.status()); + + outputs.clear(); + TF_EXPECT_OK(session.Run(tensorflow::ClientSession::FeedType(), {read_handle}, + {release_handle}, &outputs)); + + xla::Literal return_literal = ReadOutputLiteral(outputs, 0); + + auto expected_literal0 = xla::LiteralUtil::CreateR1({6.0f, 11.0f}); + auto expected_literal1 = xla::LiteralUtil::CreateR1({-4.0f, -7.0f}); + // The first element of the computation returned tuple would be the add + // (expected_literal0), but since we flipped the buffers, the sub + // (expected_literal1) should come first. + auto expected_literal = + xla::LiteralUtil::MakeTuple({&expected_literal1, &expected_literal0}); + + EXPECT_TRUE(CompareLiterals(return_literal, expected_literal)); +} + TEST(RawApiTest, CompileAndExecuteWithS64Argument) { xrt::XLAAllocation p0; *p0.mutable_value() = xla::LiteralUtil::CreateR0(11031965).ToProto(); @@ -862,6 +1199,7 @@ TEST(RawApiTest, CompileAndExecuteWithS64Argument) { xrt::XRTExecutionConfig e; e.set_release_input_handles(true); e.set_release_compilation_handle(true); + e.set_return_exploded_tuple(true); Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); auto e_config = diff --git a/tensorflow/compiler/xrt/xrt_compilation_cache.cc b/tensorflow/compiler/xrt/xrt_compilation_cache.cc index d1405eae468492748ae88d842334a922dce272c6..8bf0f28d2233d9e7593365bc42187e327a1c4ac4 100644 --- a/tensorflow/compiler/xrt/xrt_compilation_cache.cc +++ b/tensorflow/compiler/xrt/xrt_compilation_cache.cc @@ -273,6 +273,8 @@ Status XRTCompilationCache::Lookup( return Status::OK(); } -string XRTCompilationCache::DebugString() { return "XRTCompilationCache"; } +string XRTCompilationCache::DebugString() const { + return "XRTCompilationCache"; +} } // namespace tensorflow diff --git a/tensorflow/compiler/xrt/xrt_compilation_cache.h b/tensorflow/compiler/xrt/xrt_compilation_cache.h index c43d0fc47873abdc82ee937c155bebc346a05f17..7398e847d8b744f947adb03e1bcfd5c0a5b2cc55 100644 --- a/tensorflow/compiler/xrt/xrt_compilation_cache.h +++ b/tensorflow/compiler/xrt/xrt_compilation_cache.h @@ -118,7 +118,7 @@ class XRTCompilationCache : public ResourceBase { // EntryRef holding the program is returned in entry. Status Lookup(int64 uid, std::unique_ptr* entry); - string DebugString() override; + string DebugString() const override; private: // An entry in the compilation cache. The entry is deleted once it has been diff --git a/tensorflow/compiler/xrt/xrt_state.cc b/tensorflow/compiler/xrt/xrt_state.cc index 343460ff107fa81be127950837f786fe4eeadf26..78a1b6afc05cc49af3be51e97e2d72fdbc82f334 100644 --- a/tensorflow/compiler/xrt/xrt_state.cc +++ b/tensorflow/compiler/xrt/xrt_state.cc @@ -31,7 +31,6 @@ 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/xla_data.pb.h" #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/random/random.h" @@ -133,7 +132,8 @@ Status AllocateScopedShapedBuffer( XRTBufferAllocation::XRTBufferAllocation(const se::DeviceMemoryBase& allocation, int device_ordinal, xla::DeviceMemoryAllocator* allocator) - : allocation_(allocation), + : size_(allocation.size()), + allocation_(allocation), device_ordinal_(device_ordinal), allocator_(allocator) { if (VLOG_IS_ON(2)) { @@ -181,7 +181,7 @@ XRTTupleAllocation::~XRTTupleAllocation() { } /*static*/ Status XRTTupleAllocation::CreateAndTransfer( - const xla::Literal& literal, xla::Backend* backend, int device_ordinal, + const xla::LiteralBase& literal, xla::Backend* backend, int device_ordinal, XRTTupleAllocation** allocation) { auto transfer_manager = backend->transfer_manager(); auto allocator = backend->memory_allocator(); @@ -223,8 +223,19 @@ Status XRTTupleAllocation::ToLiteral(xla::Backend* backend, int device_ordinal, xla::Literal* literal) { auto transfer_manager = backend->transfer_manager(); TF_ASSIGN_OR_RETURN(auto stream, backend->BorrowStream(device_ordinal)); + + // Validate the allocation buffers as if nulls gets to + // TransferLiteralFromDevice() a CHECK is issued. + xla::ShapedBuffer shaped_buffer = ToShapedBuffer(); + for (auto& index_buffer : shaped_buffer.buffers()) { + if (index_buffer.second.is_null()) { + return errors::InvalidArgument("Literal buffer at index ", + index_buffer.first.ToString(), + " has been released"); + } + } TF_ASSIGN_OR_RETURN(*literal, transfer_manager->TransferLiteralFromDevice( - stream.get(), ToShapedBuffer())); + stream.get(), shaped_buffer)); return Status::OK(); } @@ -505,11 +516,34 @@ xla::ShapedBuffer XRTTupleAllocation::ToShapedBuffer() { return shaped_buffer; } +Status XRTTupleAllocation::AliasBufferFrom(const XRTTupleAllocation& source, + const xla::ShapeIndex& source_index, + const xla::ShapeIndex& dest_index) { + XRTBufferAllocation* source_buffer = source.buffers_.element(source_index); + XRTBufferAllocation* dest_buffer = buffers_.element(dest_index); + // We allow the destination size being zero, because there are cases where we + // are coming in later filling in null/uninitialized device buffers. + // In all other cases, the size of the new buffer must match. + if (source_buffer->size() != dest_buffer->size() && + dest_buffer->size() != 0) { + return errors::InvalidArgument( + "Source buffer at index ", source_index.ToString(), + " does not match the size of destination buffer at index ", + dest_index.ToString(), ": ", source_buffer->size(), " vs ", + dest_buffer->size()); + } + *buffers_.mutable_element(dest_index) = source_buffer; + source_buffer->Ref(); + dest_buffer->Unref(); + return Status::OK(); +} + xla::ShapeTree -XRTTupleAllocation::ToDeviceMemoryTree(bool release) { +XRTTupleAllocation::ToDeviceMemoryTree( + const std::function& release_checker) { xla::ShapeTree shaped_tree(on_device_shape()); for (const auto& buffer : buffers_) { - if (!release) { + if (!release_checker(buffer.first)) { *shaped_tree.mutable_element(buffer.first) = buffer.second->allocation(); } else { *shaped_tree.mutable_element(buffer.first) = xla::OwningDeviceMemory( diff --git a/tensorflow/compiler/xrt/xrt_state.h b/tensorflow/compiler/xrt/xrt_state.h index 3e3d5024124e13b87eed6f79596d50cd64325914..ddf2656e6f51775024a6d1cd0d7a387605faae6f 100644 --- a/tensorflow/compiler/xrt/xrt_state.h +++ b/tensorflow/compiler/xrt/xrt_state.h @@ -18,6 +18,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XRT_XRT_STATE_H_ #define TENSORFLOW_COMPILER_XRT_XRT_STATE_H_ +#include #include #include #include @@ -58,7 +59,14 @@ class XRTBufferAllocation : public core::RefCounted { // freed when the reference count drops to zero. void DiscardAllocation(); + // Returns the expected size of the allocation. Since DiscardAllocation() will + // set allocation_ to {null,0}, and since later we might want to replace the + // discarded buffer with a new one, we need to be able to verify the size + // compatibility. + uint64 size() const { return size_; } + private: + uint64 size_ = 0; se::DeviceMemoryBase allocation_; int device_ordinal_; xla::DeviceMemoryAllocator* allocator_; @@ -80,7 +88,7 @@ class XRTTupleAllocation : public ResourceBase { // Allocates new device memory buffers sufficient to store literal, transfers // literal to that memory, and returns a XRTTupleAllocation handle to the // allocated buffers. - static Status CreateAndTransfer(const xla::Literal& literal, + static Status CreateAndTransfer(const xla::LiteralBase& literal, xla::Backend* backend, int device_ordinal, XRTTupleAllocation** allocation); @@ -168,11 +176,20 @@ class XRTTupleAllocation : public ResourceBase { // the same shape as on_host_shape. xla::ShapedBuffer ToShapedBuffer(); - // Returns the device memory tree of this allocation. If 'release' is set, the - // ownership of the device memory is transferred to the result. - xla::ShapeTree ToDeviceMemoryTree(bool release); + // Aliases the source buffer at source_index into the current tuple allocation + // dest_index. + Status AliasBufferFrom(const XRTTupleAllocation& source, + const xla::ShapeIndex& source_index, + const xla::ShapeIndex& dest_index); + + // Returns the device memory tree of this allocation. If the release_checker + // function returns true for a given index, the ownership of the device memory + // at that index is transferred to the result. Every attempt to read the value + // at that index will fail. + xla::ShapeTree ToDeviceMemoryTree( + const std::function& release_checker); - string DebugString() override { return "XLA allocation handle"; } + string DebugString() const override { return "XLA allocation handle"; } private: // Creates a new handle with (tuple) shape. diff --git a/tensorflow/contrib/BUILD b/tensorflow/contrib/BUILD index 832db0f4ab46911e067d17b4a125706c276cf798..25f2640e35af5f65eab25dc60c44e3ed7ce4e512 100644 --- a/tensorflow/contrib/BUILD +++ b/tensorflow/contrib/BUILD @@ -63,7 +63,6 @@ py_library( "//tensorflow/contrib/libsvm", "//tensorflow/contrib/linear_optimizer:sdca_estimator_py", "//tensorflow/contrib/linear_optimizer:sdca_ops_py", - "//tensorflow/contrib/lite/python:lite", "//tensorflow/contrib/lookup:lookup_py", "//tensorflow/contrib/losses:losses_py", "//tensorflow/contrib/losses:metric_learning_py", @@ -197,7 +196,7 @@ cc_library( "//tensorflow/contrib/kinesis:dataset_kernels", ], }) + if_not_windows([ - "//tensorflow/contrib/tensorrt:trt_engine_op_kernel", + "//tensorflow/contrib/tensorrt:trt_op_kernels", ]), ) @@ -239,7 +238,7 @@ cc_library( "//tensorflow/contrib/kinesis:dataset_ops_op_lib", ], }) + if_not_windows([ - "//tensorflow/contrib/tensorrt:trt_engine_op_op_lib", + "//tensorflow/compiler/tf2tensorrt:trt_op_libs", ]) + select({ "//tensorflow:android": [], "//tensorflow:ios": [], diff --git a/tensorflow/contrib/__init__.py b/tensorflow/contrib/__init__.py index 4f1a2a5693235183c8f486817b82c8c81fa389ec..48d5296c71cbdb470fa405b30547a32b7022f29b 100644 --- a/tensorflow/contrib/__init__.py +++ b/tensorflow/contrib/__init__.py @@ -20,13 +20,14 @@ from __future__ import division from __future__ import print_function import os +import platform # Add projects here, they will show up under tf.contrib. from tensorflow.contrib import autograph from tensorflow.contrib import batching from tensorflow.contrib import bayesflow from tensorflow.contrib import checkpoint -if os.name != "nt": +if os.name != "nt" and platform.machine() != "s390x": from tensorflow.contrib import cloud from tensorflow.contrib import cluster_resolver from tensorflow.contrib import coder @@ -91,7 +92,6 @@ from tensorflow.contrib import tpu from tensorflow.contrib import training from tensorflow.contrib import util from tensorflow.contrib.eager.python import tfe as eager -from tensorflow.contrib.lite.python import lite from tensorflow.contrib.optimizer_v2 import optimizer_v2_symbols as optimizer_v2 from tensorflow.contrib.receptive_field import receptive_field_api as receptive_field from tensorflow.contrib.recurrent.python import recurrent_api as recurrent @@ -103,6 +103,8 @@ from tensorflow.python.util.lazy_loader import LazyLoader ffmpeg = LazyLoader("ffmpeg", globals(), "tensorflow.contrib.ffmpeg") del os +del platform + del LazyLoader del absolute_import diff --git a/tensorflow/contrib/autograph/examples/notebooks/rnn_keras_estimator.ipynb b/tensorflow/contrib/autograph/examples/notebooks/rnn_keras_estimator.ipynb index 44532cb078f9bd1578172f8a7d8a4b55cd21a7cb..831c613f2c8c9a4fcc2cb9d313077fe79ee96fd7 100644 --- a/tensorflow/contrib/autograph/examples/notebooks/rnn_keras_estimator.ipynb +++ b/tensorflow/contrib/autograph/examples/notebooks/rnn_keras_estimator.ipynb @@ -186,8 +186,8 @@ "\n", " def __init__(self):\n", " super(RnnColorbot, self).__init__()\n", - " self.lower_cell = tf.contrib.rnn.LSTMBlockCell(256)\n", - " self.upper_cell = tf.contrib.rnn.LSTMBlockCell(128)\n", + " self.lower_cell = tf.contrib.rnn.LSTMBlockCell(256, dtype=tf.float32)\n", + " self.upper_cell = tf.contrib.rnn.LSTMBlockCell(128, dtype=tf.float32)\n", " self.relu_layer = tf.layers.Dense(3, activation=tf.nn.relu)\n", "\n", " def _rnn_layer(self, chars, cell, batch_size, training):\n", @@ -241,7 +241,7 @@ " seq = self._rnn_layer(seq, self.upper_cell, batch_size, training)\n", "\n", " # Grab just the end-of-sequence from each output.\n", - " indices = (length - 1, range(batch_size))\n", + " indices = (length - 1, list(range(batch_size)))\n", " indices = tf.stack(indices, 1)\n", " sequence_ends = tf.gather_nd(seq, indices)\n", " return self.relu_layer(sequence_ends)\n", diff --git a/tensorflow/contrib/batching/BUILD b/tensorflow/contrib/batching/BUILD index 648f3ebb05646a66144bcb118347cbc391909409..5174afe0a63d37e3ea3e19ac9bab644d1d83ecf1 100644 --- a/tensorflow/contrib/batching/BUILD +++ b/tensorflow/contrib/batching/BUILD @@ -37,6 +37,7 @@ py_library( cc_library( name = "batch_ops_kernels", deps = [ + "//tensorflow/core:batch_ops_op_lib", "//tensorflow/core/kernels:batch_kernels", ], alwayslink = 1, diff --git a/tensorflow/contrib/bigtable/kernels/bigtable_kernels.cc b/tensorflow/contrib/bigtable/kernels/bigtable_kernels.cc index 6138d7912601344ef7422fd50fb35c8401fd2e63..f0637595db08cbeb3b3ee0c94c5399df4c8c83e6 100644 --- a/tensorflow/contrib/bigtable/kernels/bigtable_kernels.cc +++ b/tensorflow/contrib/bigtable/kernels/bigtable_kernels.cc @@ -19,7 +19,6 @@ limitations under the License. #include "tensorflow/core/lib/core/threadpool.h" namespace tensorflow { - namespace { class BigtableClientOp : public OpKernel { @@ -341,8 +340,8 @@ class ToBigtableOp : public AsyncOpKernel { } template - Status ParseScalarArgument(OpKernelContext* ctx, - const StringPiece& argument_name, T* output) { + Status ParseScalarArgument(OpKernelContext* ctx, StringPiece argument_name, + T* output) { const Tensor* argument_t; TF_RETURN_IF_ERROR(ctx->input(argument_name, &argument_t)); if (!TensorShapeUtils::IsScalar(argument_t->shape())) { @@ -360,5 +359,4 @@ REGISTER_KERNEL_BUILDER(Name("DatasetToBigtable").Device(DEVICE_CPU), } // namespace } // namespace data - } // namespace tensorflow diff --git a/tensorflow/contrib/bigtable/kernels/bigtable_lib.h b/tensorflow/contrib/bigtable/kernels/bigtable_lib.h index 4652021fecabfa11fa6a8754dc884d89e151b590..e3b4535bac4a01a1277290e0d1ea6d3c7613731c 100644 --- a/tensorflow/contrib/bigtable/kernels/bigtable_lib.h +++ b/tensorflow/contrib/bigtable/kernels/bigtable_lib.h @@ -42,7 +42,7 @@ class BigtableClientResource : public ResourceBase { return client_; } - string DebugString() override { + string DebugString() const override { return strings::StrCat("BigtableClientResource(project_id: ", project_id_, ", instance_id: ", instance_id_, ")"); } @@ -67,7 +67,7 @@ class BigtableTableResource : public ResourceBase { ::google::cloud::bigtable::noex::Table& table() { return table_; } - string DebugString() override { + string DebugString() const override { return strings::StrCat( "BigtableTableResource(client: ", client_->DebugString(), ", table: ", table_name_, ")"); diff --git a/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc b/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc index c04d5578eef0aadfc918c2de734723fe05c5cfee..e6fda9e61757f1441b3691c2a3d57c6f1a5a0d42 100644 --- a/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc +++ b/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc @@ -410,6 +410,17 @@ BigtableTestClient::AsyncCheckAndMutateRow( return nullptr; } +std::unique_ptr< + grpc::ClientAsyncReaderInterface> +BigtableTestClient::AsyncReadRows( + grpc::ClientContext* context, + const google::bigtable::v2::ReadRowsRequest& request, + grpc::CompletionQueue* cq, void* tag) { + LOG(WARNING) << "Call to InMemoryDataClient::" << __func__ + << "(); this will likely cause a crash!"; + return nullptr; +} + std::shared_ptr BigtableTestClient::Channel() { LOG(WARNING) << "Call to InMemoryDataClient::Channel(); this will likely " "cause a crash!"; diff --git a/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.h b/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.h index 85705904573e9e7710912e3f4ff30dd8fed5bf85..8e1326f2ce841368ea81fc7194a0588e5d6cd637 100644 --- a/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.h +++ b/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.h @@ -87,6 +87,12 @@ class BigtableTestClient : public ::google::cloud::bigtable::DataClient { const google::bigtable::v2::CheckAndMutateRowRequest& request, grpc::CompletionQueue* cq) override; + std::unique_ptr< + grpc::ClientAsyncReaderInterface> + AsyncReadRows(grpc::ClientContext* context, + const google::bigtable::v2::ReadRowsRequest& request, + grpc::CompletionQueue* cq, void* tag) override; + std::shared_ptr Channel() override; private: diff --git a/tensorflow/contrib/bigtable/ops/bigtable_ops.cc b/tensorflow/contrib/bigtable/ops/bigtable_ops.cc index 416b719e30aa5f2504449d151a48e95c9105c68b..39c2a2e775d5d5287b137bf33eef66251738e6d3 100644 --- a/tensorflow/contrib/bigtable/ops/bigtable_ops.cc +++ b/tensorflow/contrib/bigtable/ops/bigtable_ops.cc @@ -59,7 +59,7 @@ REGISTER_OP("BigtablePrefixKeyDataset") .Input("table: resource") .Input("prefix: string") .Output("handle: variant") - .SetIsStateful() // TODO(b/65524810): Source dataset ops must be marked + .SetIsStateful() // TODO(b/123753214): Source dataset ops must be marked // stateful to inhibit constant folding. .SetShapeFn(shape_inference::ScalarShape); @@ -68,14 +68,14 @@ REGISTER_OP("BigtableRangeKeyDataset") .Input("start_key: string") .Input("end_key: string") .Output("handle: variant") - .SetIsStateful() // TODO(b/65524810): Source dataset ops must be marked + .SetIsStateful() // TODO(b/123753214): Source dataset ops must be marked // stateful to inhibit constant folding. .SetShapeFn(shape_inference::ScalarShape); REGISTER_OP("BigtableSampleKeysDataset") .Input("table: resource") .Output("handle: variant") - .SetIsStateful() // TODO(b/65524810): Source dataset ops must be marked + .SetIsStateful() // TODO(b/123753214): Source dataset ops must be marked // stateful to inhibit constant folding. .SetShapeFn(shape_inference::ScalarShape); @@ -85,7 +85,7 @@ REGISTER_OP("BigtableSampleKeyPairsDataset") .Input("start_key: string") .Input("end_key: string") .Output("handle: variant") - .SetIsStateful() // TODO(b/65524810): Source dataset ops must be marked + .SetIsStateful() // TODO(b/123753214): Source dataset ops must be marked // stateful to inhibit constant folding. .SetShapeFn(shape_inference::ScalarShape); @@ -100,7 +100,7 @@ REGISTER_OP("BigtableScanDataset") .Input("columns: string") .Input("probability: float") .Output("handle: variant") - .SetIsStateful() // TODO(b/65524810): Source dataset ops must be marked + .SetIsStateful() // TODO(b/123753214): Source dataset ops must be marked // stateful to inhibit constant folding. .SetShapeFn(shape_inference::ScalarShape); diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/BUILD b/tensorflow/contrib/boosted_trees/estimator_batch/BUILD index d3b23d949ee2c7674c3918d39e8b71d76eefcfec..64e4c4560ba3a1b177db12a09997ff7afe8775a3 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/BUILD +++ b/tensorflow/contrib/boosted_trees/estimator_batch/BUILD @@ -193,8 +193,9 @@ py_test( py_test( name = "estimator_test", - size = "large", + size = "medium", srcs = ["estimator_test.py"], + shard_count = 4, srcs_version = "PY2AND3", tags = [ "no_gpu", diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/estimator_test.py b/tensorflow/contrib/boosted_trees/estimator_batch/estimator_test.py index ee052ac60387d8f993e4942dd7dff39e191dd3a4..47d910d42a27db4b857eeb12209dfbb429dd1be2 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/estimator_test.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/estimator_test.py @@ -487,8 +487,8 @@ class BoostedTreeEstimatorTest(test_util.TensorFlowTestCase): self.assertTrue(frac_below_upper_0 <= 0.98) self.assertTrue(frac_below_upper_1 >= 0.92) self.assertTrue(frac_below_upper_1 <= 0.98) - self.assertTrue(frac_both_below_upper >= 0.92) - self.assertTrue(frac_both_below_upper <= 0.98) + self.assertTrue(frac_both_below_upper >= 0.91) + self.assertTrue(frac_both_below_upper <= 0.99) train_input_fn, test_input_fn, _ = _quantile_regression_input_fns( two_dimension=True) @@ -516,8 +516,8 @@ class BoostedTreeEstimatorTest(test_util.TensorFlowTestCase): self.assertTrue(frac_above_lower_0 <= 0.98) self.assertTrue(frac_above_lower_1 >= 0.92) self.assertTrue(frac_above_lower_1 <= 0.98) - self.assertTrue(frac_both_above_lower >= 0.92) - self.assertTrue(frac_both_above_lower <= 0.98) + self.assertTrue(frac_both_above_lower >= 0.91) + self.assertTrue(frac_both_above_lower <= 0.99) class CoreGradientBoostedDecisionTreeEstimators(test_util.TensorFlowTestCase): @@ -806,8 +806,8 @@ class CoreGradientBoostedDecisionTreeEstimators(test_util.TensorFlowTestCase): self.assertTrue(frac_below_upper_0 <= 0.98) self.assertTrue(frac_below_upper_1 >= 0.92) self.assertTrue(frac_below_upper_1 <= 0.98) - self.assertTrue(frac_both_below_upper >= 0.92) - self.assertTrue(frac_both_below_upper <= 0.98) + self.assertTrue(frac_both_below_upper >= 0.91) + self.assertTrue(frac_both_below_upper <= 0.99) train_input_fn, test_input_fn, _ = _quantile_regression_input_fns( two_dimension=True) @@ -835,8 +835,8 @@ class CoreGradientBoostedDecisionTreeEstimators(test_util.TensorFlowTestCase): self.assertTrue(frac_above_lower_0 <= 0.98) self.assertTrue(frac_above_lower_1 >= 0.92) self.assertTrue(frac_above_lower_1 <= 0.98) - self.assertTrue(frac_both_above_lower >= 0.92) - self.assertTrue(frac_both_above_lower <= 0.98) + self.assertTrue(frac_both_above_lower >= 0.91) + self.assertTrue(frac_both_above_lower <= 0.99) if __name__ == "__main__": diff --git a/tensorflow/contrib/boosted_trees/kernels/stats_accumulator_ops.cc b/tensorflow/contrib/boosted_trees/kernels/stats_accumulator_ops.cc index e446c411a8d5075563b8f8b912b29df310e16c8c..6faf6963011b698a3b233329d87471da7608e44a 100644 --- a/tensorflow/contrib/boosted_trees/kernels/stats_accumulator_ops.cc +++ b/tensorflow/contrib/boosted_trees/kernels/stats_accumulator_ops.cc @@ -96,7 +96,7 @@ class StatsAccumulatorResource : public boosted_trees::StampedResource { TensorShapeUtils::IsScalar(hessian_shape)); } - string DebugString() override { + string DebugString() const override { return strings::StrCat("StatsAccumulatorResource[size=", values_.size(), "]"); } diff --git a/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py b/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py index 42d69645acaae063fcd46bd1f6c819ccb68f48bd..aa3f24f08a0f762507df83def72e7d595265221f 100644 --- a/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py +++ b/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py @@ -227,7 +227,7 @@ class ModelOpsTest(test_util.TensorFlowTestCase): tree_ensemble_config=tree_ensemble_config.SerializeToString(), name="restore_tree") resources.initialize_resources(resources.shared_resources()).run() - variables.initialize_all_variables().run() + variables.global_variables_initializer().run() my_saver = saver.Saver() # Add the second tree and replace the ensemble of the handle. diff --git a/tensorflow/contrib/boosted_trees/python/ops/model_ops.py b/tensorflow/contrib/boosted_trees/python/ops/model_ops.py index fca22c71a83459cb290eaebcf107cf1c14c222b7..c3685b54e201f73039f6623443c67ba2b217a51e 100644 --- a/tensorflow/contrib/boosted_trees/python/ops/model_ops.py +++ b/tensorflow/contrib/boosted_trees/python/ops/model_ops.py @@ -62,8 +62,8 @@ class TreeEnsembleVariableSavable(saver.BaseSaverBuilder.SaveableObject): saver.BaseSaverBuilder.SaveSpec(ensemble_config, slice_spec, name + "_config"), ] - super(TreeEnsembleVariableSavable, - self).__init__(tree_ensemble_handle, specs, name) + super(TreeEnsembleVariableSavable, self).__init__(tree_ensemble_handle, + specs, name) self._tree_ensemble_handle = tree_ensemble_handle self._create_op = create_op @@ -115,7 +115,7 @@ class TreeEnsembleVariable(tracking.TrackableResource): def _gather_saveables_for_checkpoint(self): return { - "tree_ensemble_variable": + self.resource_handle.op.name + "/tree_ensemble_variable": functools.partial( TreeEnsembleVariableSavable, tree_ensemble_handle=self.resource_handle, @@ -131,8 +131,8 @@ def tree_ensemble_variable(stamp_token, Args: stamp_token: The initial stamp token value for the ensemble resource. - tree_ensemble_config: A `Tensor` of type `string`. - Serialized proto of the tree ensemble. + tree_ensemble_config: A `Tensor` of type `string`. Serialized proto of the + tree ensemble. name: A name for the ensemble variable. container: An optional `string`. Defaults to `""`. diff --git a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py index a5951fb7377d48748f5eb578c034176517df7749..e78ec476ab3b43e5eb56a2502008bb8020ae97e0 100644 --- a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py +++ b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py @@ -566,9 +566,10 @@ class GradientBoostedDecisionTreeModel(object): # Determine if ensemble is colocated with the inputs. if self._ensemble_handle.device != input_deps[0].device: # Create a local ensemble and get its local stamp. - with ops.name_scope("local_ensemble", "TreeEnsembleVariable") as name: + with ops.name_scope("local_ensemble", "TreeEnsembleVariable"): local_ensemble_handle = ( - gen_model_ops.decision_tree_ensemble_resource_handle_op(name=name)) + gen_model_ops.decision_tree_ensemble_resource_handle_op( + self._ensemble_handle.op.name + "/local_ensemble")) create_op = gen_model_ops.create_tree_ensemble_variable( local_ensemble_handle, stamp_token=-1, tree_ensemble_config="") with ops.control_dependencies([create_op]): diff --git a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch_test.py b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch_test.py index 92068e88a76cb8bfdd394c1093347a8fb8a63449..7e45d0b2cecefa4bdec77d6cf7cfca7dba04db9c 100644 --- a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch_test.py +++ b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch_test.py @@ -43,7 +43,7 @@ from tensorflow.python.platform import googletest def _squared_loss(label, unused_weights, predictions): """Unweighted loss implementation.""" loss = math_ops.reduce_sum( - math_ops.square(predictions - label), 1, keepdims=True) + math_ops.squared_difference(predictions, label), 1, keepdims=True) return loss diff --git a/tensorflow/contrib/boosted_trees/python/utils/losses.py b/tensorflow/contrib/boosted_trees/python/utils/losses.py index 220e981618b7c0bfb1e4e98c087d83b451b9b3cf..1ad40aca2880940c78d746674c7378ff0427c057 100644 --- a/tensorflow/contrib/boosted_trees/python/utils/losses.py +++ b/tensorflow/contrib/boosted_trees/python/utils/losses.py @@ -166,7 +166,7 @@ def per_example_squared_loss(labels, weights, predictions): update_op: An update operation to update the loss's internal state. """ unweighted_loss = math_ops.reduce_sum( - math_ops.square(predictions - labels), 1, keepdims=True) + math_ops.squared_difference(predictions, labels), 1, keepdims=True) return unweighted_loss * weights, control_flow_ops.no_op() diff --git a/tensorflow/contrib/boosted_trees/resources/decision_tree_ensemble_resource.h b/tensorflow/contrib/boosted_trees/resources/decision_tree_ensemble_resource.h index 94aeb2c7bb48c6eddb6c7894f8bf6f1567470113..0fe57c0a4e8375cc7ec7aca9553bded87e238b33 100644 --- a/tensorflow/contrib/boosted_trees/resources/decision_tree_ensemble_resource.h +++ b/tensorflow/contrib/boosted_trees/resources/decision_tree_ensemble_resource.h @@ -34,7 +34,7 @@ class DecisionTreeEnsembleResource : public StampedResource { protobuf::Arena::CreateMessage< boosted_trees::trees::DecisionTreeEnsembleConfig>(&arena_)) {} - string DebugString() override { + string DebugString() const override { return strings::StrCat("GTFlowDecisionTreeEnsemble[size=", decision_tree_ensemble_->trees_size(), "]"); } diff --git a/tensorflow/contrib/boosted_trees/resources/quantile_stream_resource.h b/tensorflow/contrib/boosted_trees/resources/quantile_stream_resource.h index fdaaae7f472c8f564ab45a8366d3746cbf1158ee..574e3065e7f46049815897ef73e44d33f0d23f0f 100644 --- a/tensorflow/contrib/boosted_trees/resources/quantile_stream_resource.h +++ b/tensorflow/contrib/boosted_trees/resources/quantile_stream_resource.h @@ -43,7 +43,7 @@ class QuantileStreamResource : public StampedResource { set_stamp(stamp_token); } - string DebugString() override { return "QuantileStreamResource"; } + string DebugString() const override { return "QuantileStreamResource"; } tensorflow::mutex* mutex() { return &mu_; } diff --git a/tensorflow/contrib/checkpoint/__init__.py b/tensorflow/contrib/checkpoint/__init__.py index 94b7f4f867655bf7fdf94e8488eeae7088c41622..99ed4959fad9699f265183d71a1f3b609d7e6d30 100644 --- a/tensorflow/contrib/checkpoint/__init__.py +++ b/tensorflow/contrib/checkpoint/__init__.py @@ -51,11 +51,11 @@ from tensorflow.contrib.checkpoint.python.split_dependency import split_dependen from tensorflow.contrib.checkpoint.python.visualize import dot_graph_from_checkpoint from tensorflow.core.protobuf.checkpointable_object_graph_pb2 import CheckpointableObjectGraph from tensorflow.python.training.checkpoint_management import CheckpointManager -from tensorflow.python.training.checkpointable.base import CheckpointableBase +from tensorflow.python.training.checkpointable.base import Checkpointable as CheckpointableBase from tensorflow.python.training.checkpointable.data_structures import List from tensorflow.python.training.checkpointable.data_structures import Mapping from tensorflow.python.training.checkpointable.data_structures import NoDependency -from tensorflow.python.training.checkpointable.tracking import Checkpointable +from tensorflow.python.training.checkpointable.tracking import AutoCheckpointable as Checkpointable from tensorflow.python.training.checkpointable.util import capture_dependencies from tensorflow.python.training.checkpointable.util import list_objects from tensorflow.python.training.checkpointable.util import object_metadata diff --git a/tensorflow/contrib/checkpoint/python/containers.py b/tensorflow/contrib/checkpoint/python/containers.py index 5418e2605b724edb60878e250d2c50fcc6ff5633..97936d9e9dfd5d6e62fdf8312707a276b63e1267 100644 --- a/tensorflow/contrib/checkpoint/python/containers.py +++ b/tensorflow/contrib/checkpoint/python/containers.py @@ -63,7 +63,7 @@ class UniqueNameTracker(data_structures.CheckpointableDataStructure): ValueError: If `checkpointable` is not a checkpointable object. """ - if not isinstance(checkpointable, checkpointable_lib.CheckpointableBase): + if not isinstance(checkpointable, checkpointable_lib.Checkpointable): raise ValueError( ("Expected a checkpointable value, got %s which does not inherit " "from CheckpointableBase.") % (checkpointable,)) diff --git a/tensorflow/contrib/checkpoint/python/containers_test.py b/tensorflow/contrib/checkpoint/python/containers_test.py index ac85c7be803cd4c2f8ba19d3ef887a3c65a15933..a2d453ec6eb3dcf9aba4c52fe866756a92673c63 100644 --- a/tensorflow/contrib/checkpoint/python/containers_test.py +++ b/tensorflow/contrib/checkpoint/python/containers_test.py @@ -52,7 +52,7 @@ class UniqueNameTrackerTests(test.TestCase): save_root = util.Checkpoint(slots=slots) save_path = save_root.save(checkpoint_prefix) - restore_slots = tracking.Checkpointable() + restore_slots = tracking.AutoCheckpointable() restore_root = util.Checkpoint( slots=restore_slots) status = restore_root.restore(save_path) @@ -68,7 +68,7 @@ class UniqueNameTrackerTests(test.TestCase): @test_util.run_in_graph_and_eager_modes def testExample(self): - class SlotManager(tracking.Checkpointable): + class SlotManager(tracking.AutoCheckpointable): def __init__(self): self.slotdeps = containers.UniqueNameTracker() diff --git a/tensorflow/contrib/checkpoint/python/python_state.py b/tensorflow/contrib/checkpoint/python/python_state.py index 302d5cfb79a08b6adf52ebd44533152c5454eadc..969c90c78871ebff02b360f8f09623df56c9c077 100644 --- a/tensorflow/contrib/checkpoint/python/python_state.py +++ b/tensorflow/contrib/checkpoint/python/python_state.py @@ -34,7 +34,7 @@ except ImportError: # pylint: enable=g-import-not-at-top -class NumpyState(base.CheckpointableBase): +class NumpyState(base.Checkpointable): """A checkpointable object whose NumPy array attributes are saved/restored. Example usage: @@ -130,7 +130,7 @@ class NumpyState(base.CheckpointableBase): @six.add_metaclass(abc.ABCMeta) -class PythonStateWrapper(base.CheckpointableBase): +class PythonStateWrapper(base.Checkpointable): """Wraps a Python object for storage in an object-based checkpoint.""" @abc.abstractmethod diff --git a/tensorflow/contrib/checkpoint/python/split_dependency.py b/tensorflow/contrib/checkpoint/python/split_dependency.py index 7e77453f3d848c2e321ed2ba66917a742d95459a..3e9700ad74618e24843181d169f3fb39ac96bff6 100644 --- a/tensorflow/contrib/checkpoint/python/split_dependency.py +++ b/tensorflow/contrib/checkpoint/python/split_dependency.py @@ -43,7 +43,7 @@ class _CallbackSaveable(saver_lib.BaseSaverBuilder.SaveableObject): return self._restore_callback(tensor) -class _SplitDependency(checkpointable.CheckpointableBase): +class _SplitDependency(checkpointable.Checkpointable): """Looks like a regular variable while synchronizing save/restores.""" def __init__(self, save_buffer, restore_buffer, name, dtype, num_components, diff --git a/tensorflow/contrib/checkpoint/python/split_dependency_test.py b/tensorflow/contrib/checkpoint/python/split_dependency_test.py index 00a805af25d5d0ea723db5d015fb12bf45c53857..664a4e76ab31bf31c7a57924e4af866f2d746804 100644 --- a/tensorflow/contrib/checkpoint/python/split_dependency_test.py +++ b/tensorflow/contrib/checkpoint/python/split_dependency_test.py @@ -44,7 +44,7 @@ def _combine_variable_closure(variable): return _consume_restore_buffer_fn -class SaveTensorSlicesAsDeps(base.CheckpointableBase): +class SaveTensorSlicesAsDeps(base.Checkpointable): def __init__(self): self.combined = resource_variable_ops.ResourceVariable([0., 0., 0., 0.]) @@ -59,14 +59,14 @@ class SaveTensorSlicesAsDeps(base.CheckpointableBase): self._track_checkpointable(dep, name=name) -class HasRegularDeps(tracking.Checkpointable): +class HasRegularDeps(tracking.AutoCheckpointable): def __init__(self): self.first_half = resource_variable_ops.ResourceVariable([0., 0.]) self.second_half = resource_variable_ops.ResourceVariable([0., 0.]) -class OnlyOneDep(tracking.Checkpointable): +class OnlyOneDep(tracking.AutoCheckpointable): def __init__(self): self.first_half = resource_variable_ops.ResourceVariable([0., 0.]) diff --git a/tensorflow/contrib/cluster_resolver/__init__.py b/tensorflow/contrib/cluster_resolver/__init__.py index 390b3e7550b3d991269bb84707c3500f2fa33290..a4dea85efd98893c881abbd3f7ebda78755b8189 100644 --- a/tensorflow/contrib/cluster_resolver/__init__.py +++ b/tensorflow/contrib/cluster_resolver/__init__.py @@ -23,7 +23,7 @@ from __future__ import print_function from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver from tensorflow.python.distribute.cluster_resolver.cluster_resolver import SimpleClusterResolver from tensorflow.python.distribute.cluster_resolver.cluster_resolver import UnionClusterResolver -from tensorflow.python.distribute.cluster_resolver.gce_cluster_resolver import GceClusterResolver +from tensorflow.python.distribute.cluster_resolver.gce_cluster_resolver import GCEClusterResolver from tensorflow.python.distribute.cluster_resolver.kubernetes_cluster_resolver import KubernetesClusterResolver from tensorflow.python.distribute.cluster_resolver.slurm_cluster_resolver import SlurmClusterResolver from tensorflow.python.distribute.cluster_resolver.tfconfig_cluster_resolver import TFConfigClusterResolver @@ -36,7 +36,7 @@ _allowed_symbols = [ 'ClusterResolver', 'SimpleClusterResolver', 'UnionClusterResolver', - 'GceClusterResolver', + 'GCEClusterResolver', 'KubernetesClusterResolver', 'TFConfigClusterResolver', 'TPUClusterResolver', diff --git a/tensorflow/contrib/cluster_resolver/python/training/__init__.py b/tensorflow/contrib/cluster_resolver/python/training/__init__.py index 10d93549ebbd4f7e900796d0516b0af1744224af..ef1e9f11a07a5be6c0b181f5e0b80e0e2214f972 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/__init__.py +++ b/tensorflow/contrib/cluster_resolver/python/training/__init__.py @@ -25,7 +25,7 @@ from __future__ import print_function from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver from tensorflow.python.distribute.cluster_resolver.cluster_resolver import SimpleClusterResolver from tensorflow.python.distribute.cluster_resolver.cluster_resolver import UnionClusterResolver -from tensorflow.python.distribute.cluster_resolver.gce_cluster_resolver import GceClusterResolver +from tensorflow.python.distribute.cluster_resolver.gce_cluster_resolver import GCEClusterResolver from tensorflow.python.distribute.cluster_resolver.kubernetes_cluster_resolver import KubernetesClusterResolver from tensorflow.python.distribute.cluster_resolver.slurm_cluster_resolver import SlurmClusterResolver from tensorflow.python.distribute.cluster_resolver.tfconfig_cluster_resolver import TFConfigClusterResolver @@ -43,7 +43,7 @@ _allowed_symbols = [ 'ClusterResolver', 'SimpleClusterResolver', 'UnionClusterResolver', - 'GceClusterResolver', + 'GCEClusterResolver', 'KubernetesClusterResolver', 'TFConfigClusterResolver', 'TPUClusterResolver', diff --git a/tensorflow/contrib/cluster_resolver/python/training/gce_cluster_resolver.py b/tensorflow/contrib/cluster_resolver/python/training/gce_cluster_resolver.py index 55e61155c683c928efab9bb018868faec3e3df8c..5b49116ff6a4e17a774ea79b33ae1b948ba9f187 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/gce_cluster_resolver.py +++ b/tensorflow/contrib/cluster_resolver/python/training/gce_cluster_resolver.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Stub file for GceClusterResolver to maintain backwards compatibility.""" +"""Stub file for GCEClusterResolver to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division @@ -23,13 +23,14 @@ from __future__ import print_function # existing OSS code will not be broken. # pylint: disable=unused-import -from tensorflow.python.distribute.cluster_resolver.gce_cluster_resolver import GceClusterResolver +from tensorflow.python.distribute.cluster_resolver.gce_cluster_resolver import GCEClusterResolver # pylint: enable=unused-import from tensorflow.python.util.all_util import remove_undocumented + _allowed_symbols = [ - 'GceClusterResolver', + 'GCEClusterResolver', ] remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/cmake/README.md b/tensorflow/contrib/cmake/README.md index b2badc5785bdb1ea90c7f07e544ea9047146eebd..60ee1b4b3fd7d0b6afaefcc05effd3bbae00cf2c 100644 --- a/tensorflow/contrib/cmake/README.md +++ b/tensorflow/contrib/cmake/README.md @@ -147,19 +147,19 @@ suitable interface for project configuration and dependency setting. * Go (required if you need ssl support, optional) * NASM/YASM (required by grpc for ssl support, optional) 2. Start CMake GUI -3. Click on `Browse Source` and direct to the the folder +3. Click on `Browse Source` and direct to the folder `/tensorflow/contrib/cmake` 4. Click on `Browse Build` and spectify a location that you want tensorflow to be build 5. Click on `Configure`, a new window will be prompted out, specify the generator mode for the project generation. For Windows, choose `Visual Studio Win64`, for Linux, choose `Unix Makefiles`, then - press `Finish`. Wait for a moment, the default project dependecy would + press `Finish`. Wait for a moment, the default project dependency would automatically generate. 6. There are a few options that you can customize your own build. **The setting here is crucial for a successful build, please check all items carefully.** - * `tensorflow_BUILD_ALL_KERNELS` should alway be `on` + * `tensorflow_BUILD_ALL_KERNELS` should always be `on` * `tensorflow_BUILD_CC_EXAMPLE` is default to be `on`. This can help you to test build (optional) * `tensorflow_BUILD_CONTRIB_KERNELS` is default to be `on`, but it won't @@ -278,7 +278,7 @@ suitable interface for project configuration and dependency setting. `make -sj install` Where `` is the threads used for the compilation, change - to any integer less or equal to your computer's maxiumum thread number. + to any integer less or equal to your computer's maximum thread number. Headers are discretely located in the build folders. Tensorflow library can be found at ``, namely `tensorflow.so` (Linux) or diff --git a/tensorflow/contrib/cmake/external/abseil_cpp.cmake b/tensorflow/contrib/cmake/external/abseil_cpp.cmake index 46a193971c5084523d432065f265fa7a9909f595..6c6a5df7f76723800740a81ccdcb137a0ec33846 100644 --- a/tensorflow/contrib/cmake/external/abseil_cpp.cmake +++ b/tensorflow/contrib/cmake/external/abseil_cpp.cmake @@ -31,17 +31,17 @@ if (systemlib_ABSEIL_CPP) message(STATUS " abseil_cpp includes: ${ABSEIL_CPP_INCLUDE_DIR}") message(STATUS " abseil_cpp libraries: ${ABSEIL_CPP_LIBRARIES}") - add_custom_target(abseil_cpp) - list(APPEND tensorflow_EXTERNAL_DEPENDENCIES abseil_cpp) + add_custom_target(abseil_cpp_build) + list(APPEND tensorflow_EXTERNAL_DEPENDENCIES abseil_cpp_build) else (systemlib_ABSEIL_CPP) include (ExternalProject) - set(abseil_cpp_INCLUDE_DIR ${CMAKE_BINARY_DIR}/abseil_cpp/src/abseil_cpp) - set(abseil_cpp_URL https://github.com/abseil/abseil-cpp/archive/e01d95528ea2137a4a27a88d1f57c6cb260aafed.tar.gz) - set(abseil_cpp_HASH SHA256=84043ed402d2a2a6ba4cdddb7e85118b1158fd81fe4ac3a14adc343d054c1e2e) - set(abseil_cpp_BUILD ${CMAKE_BINARY_DIR}/abseil_cpp/src/abseil_cpp-build) + set(abseil_cpp_INCLUDE_DIR ${CMAKE_BINARY_DIR}/abseil_cpp/src/abseil_cpp_build) + set(abseil_cpp_URL https://github.com/abseil/abseil-cpp.git) + set(abseil_cpp_TAG master) + set(abseil_cpp_BUILD ${CMAKE_BINARY_DIR}/abseil_cpp/src/abseil_cpp_build) if(WIN32) if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") @@ -49,8 +49,11 @@ else (systemlib_ABSEIL_CPP) ${abseil_cpp_BUILD}/absl/base/Release/absl_base.lib ${abseil_cpp_BUILD}/absl/base/Release/absl_dynamic_annotations.lib ${abseil_cpp_BUILD}/absl/base/Release/absl_internal_malloc_internal.lib + ${abseil_cpp_BUILD}/absl/base/Release/absl_internal_throw_delegate.lib + ${abseil_cpp_BUILD}/absl/numeric/Release/absl_int128.lib ${abseil_cpp_BUILD}/absl/strings/Release/absl_strings.lib ${abseil_cpp_BUILD}/absl/strings/Release/str_format_internal.lib + ${abseil_cpp_BUILD}/absl/time/Release/absl_time.lib ${abseil_cpp_BUILD}/absl/types/Release/absl_bad_optional_access.lib) else() set(abseil_cpp_STATIC_LIBRARIES @@ -62,6 +65,7 @@ else (systemlib_ABSEIL_CPP) ${abseil_cpp_BUILD}/absl/numeric/absl_int128.lib ${abseil_cpp_BUILD}/absl/strings/absl_strings.lib ${abseil_cpp_BUILD}/absl/strings/str_format_internal.lib + ${abseil_cpp_BUILD}/absl/time/absl_time.lib ${abseil_cpp_BUILD}/absl/types/absl_bad_optional_access.lib) endif() else() @@ -74,15 +78,18 @@ else (systemlib_ABSEIL_CPP) ${abseil_cpp_BUILD}/absl/numeric/libabsl_int128.a ${abseil_cpp_BUILD}/absl/strings/libabsl_strings.a ${abseil_cpp_BUILD}/absl/strings/libstr_format_internal.a + ${abseil_cpp_BUILD}/absl/time/libabsl_time.a ${abseil_cpp_BUILD}/absl/types/libabsl_bad_optional_access.a) endif() - ExternalProject_Add(abseil_cpp + ExternalProject_Add(abseil_cpp_build PREFIX abseil_cpp - URL ${abseil_cpp_URL} - URL_HASH ${abseil_cpp_HASH} + GIT_REPOSITORY ${abseil_cpp_URL} DOWNLOAD_DIR "${DOWNLOAD_LOCATION}" + BUILD_IN_SOURCE 1 BUILD_BYPRODUCTS ${abseil_cpp_STATIC_LIBRARIES} + BUILD_COMMAND ${CMAKE_COMMAND} --build . --config Release + COMMAND ${CMAKE_COMMAND} --build . --config Release INSTALL_COMMAND "" CMAKE_CACHE_ARGS -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=${tensorflow_ENABLE_POSITION_INDEPENDENT_CODE} @@ -91,8 +98,10 @@ else (systemlib_ABSEIL_CPP) ) include_directories(${abseil_cpp_INCLUDE_DIR}) + message(STATUS ${abseil_cpp_INCLUDE_DIR}) + list(APPEND tensorflow_EXTERNAL_LIBRARIES ${abseil_cpp_STATIC_LIBRARIES}) - list(APPEND tensorflow_EXTERNAL_DEPENDENCIES abseil_cpp) + list(APPEND tensorflow_EXTERNAL_DEPENDENCIES abseil_cpp_build) endif (systemlib_ABSEIL_CPP) \ No newline at end of file diff --git a/tensorflow/contrib/cmake/external/grpc.cmake b/tensorflow/contrib/cmake/external/grpc.cmake index e570c09ecb5e64130ed6f3375a51d74850cc3989..379b530361f42279a8d489282bf1b35f08ba74cf 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 69b6c047bc767b4d80e7af4d00ccb7c45b683dae) +set(GRPC_TAG d0d93bdab84f2befb425e9a991d17dc78c195c6d) if(WIN32) # We use unsecure gRPC because boringssl does not build on windows diff --git a/tensorflow/contrib/cmake/external/nsync.cmake b/tensorflow/contrib/cmake/external/nsync.cmake index 479609458c64f7c7bd7b3ce6b23aceaa3db17f21..b15143bfc1cd787b156c9d6dd724a17730f0f8fb 100644 --- a/tensorflow/contrib/cmake/external/nsync.cmake +++ b/tensorflow/contrib/cmake/external/nsync.cmake @@ -16,7 +16,7 @@ include (ExternalProject) set(nsync_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR}/external/nsync/public) set(nsync_URL https://github.com/google/nsync) -set(nsync_TAG 1.20.1) +set(nsync_TAG 1.20.2) set(nsync_BUILD ${CMAKE_CURRENT_BINARY_DIR}/nsync/src/nsync) set(nsync_INSTALL ${CMAKE_CURRENT_BINARY_DIR}/nsync/install) diff --git a/tensorflow/contrib/cmake/python_modules.txt b/tensorflow/contrib/cmake/python_modules.txt index 96160568fa79291a7b391761373e1eaf0f70974e..8b6395304bb81476775e5a2d8f2ec7876035778c 100644 --- a/tensorflow/contrib/cmake/python_modules.txt +++ b/tensorflow/contrib/cmake/python_modules.txt @@ -1,6 +1,9 @@ # python_sanity_test.py will complain about invalid or missing entries # problematic entries can be commented for temporary whitelisting tensorflow +tensorflow/compiler +tensorflow/compiler/xla +tensorflow/compiler/xla/service tensorflow/core tensorflow/core/example tensorflow/core/framework @@ -10,6 +13,7 @@ tensorflow/core/lib tensorflow/core/lib/core tensorflow/core/profiler tensorflow/core/protobuf +tensorflow/core/protobuf/tpu tensorflow/core/util tensorflow/examples tensorflow/examples/tutorials @@ -434,7 +438,6 @@ tensorflow/contrib/timeseries/python/timeseries/state_space_models tensorflow/contrib/tpu tensorflow/contrib/tpu/ops tensorflow/contrib/tpu/profiler -tensorflow/contrib/tpu/proto tensorflow/contrib/tpu/python tensorflow/contrib/tpu/python/ops tensorflow/contrib/tpu/python/profiler diff --git a/tensorflow/contrib/cmake/python_protos.txt b/tensorflow/contrib/cmake/python_protos.txt index 013180c89083748b240ad061b342300e886d3568..b4603206da419f44af0857b9b933eb7df1b255ff 100644 --- a/tensorflow/contrib/cmake/python_protos.txt +++ b/tensorflow/contrib/cmake/python_protos.txt @@ -1,6 +1,7 @@ tensorflow/core tensorflow/core/kernels/boosted_trees tensorflow/core/profiler +tensorflow/core/protobuf/tpu tensorflow/python tensorflow/contrib/boosted_trees/proto tensorflow/contrib/cloud/kernels @@ -12,7 +13,6 @@ tensorflow/contrib/mpi_collectives tensorflow/contrib/session_bundle tensorflow/contrib/tensor_forest/proto tensorflow/contrib/tensorboard/plugins/projector -tensorflow/contrib/tpu/proto tensorflow/contrib/tpu/profiler tensorflow/contrib/training/python/training tensorflow/contrib/verbs diff --git a/tensorflow/contrib/cmake/tf_core_framework.cmake b/tensorflow/contrib/cmake/tf_core_framework.cmake index d7b2a1339e047aba0a9424a53a63726805e89721..cc263d7995c01100f1c51436bcb584b600c8c161 100644 --- a/tensorflow/contrib/cmake/tf_core_framework.cmake +++ b/tensorflow/contrib/cmake/tf_core_framework.cmake @@ -125,9 +125,9 @@ endfunction() file(GLOB_RECURSE tf_protos_cc_srcs RELATIVE ${tensorflow_source_dir} "${tensorflow_source_dir}/tensorflow/core/*.proto" + "${tensorflow_source_dir}/tensorflow/core/protobuf/tpu/*.proto" "${tensorflow_source_dir}/tensorflow/compiler/xla/*.proto" "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/proto/*.proto" - "${tensorflow_source_dir}/tensorflow/contrib/tpu/proto/*.proto" ) RELATIVE_PROTOBUF_GENERATE_CPP(PROTO_SRCS PROTO_HDRS @@ -147,7 +147,6 @@ set(tf_proto_text_srcs "tensorflow/core/framework/function.proto" "tensorflow/core/framework/graph.proto" "tensorflow/core/framework/graph_transfer_info.proto" - "tensorflow/core/framework/iterator.proto" "tensorflow/core/framework/kernel_def.proto" "tensorflow/core/framework/log_memory.proto" "tensorflow/core/framework/node_def.proto" @@ -302,8 +301,8 @@ file(GLOB_RECURSE tf_core_framework_srcs "${tensorflow_source_dir}/tensorflow/core/common_runtime/session.cc" "${tensorflow_source_dir}/tensorflow/core/common_runtime/session_factory.cc" "${tensorflow_source_dir}/tensorflow/core/common_runtime/session_options.cc" - "${tensorflow_source_dir}/tensorflow/contrib/tensorboard/db/*.cc" - "${tensorflow_source_dir}/tensorflow/contrib/tensorboard/db/*.h" + "${tensorflow_source_dir}/tensorflow/core/summary/*.cc" + "${tensorflow_source_dir}/tensorflow/core/summary/*.h" "${tensorflow_source_dir}/public/*.h" ) @@ -317,14 +316,14 @@ file(GLOB_RECURSE tf_core_framework_exclude_srcs "${tensorflow_source_dir}/tensorflow/core/util/*test*.h" "${tensorflow_source_dir}/tensorflow/core/util/*test*.cc" "${tensorflow_source_dir}/tensorflow/core/util/*main.cc" - "${tensorflow_source_dir}/tensorflow/contrib/tensorboard/db/*test*.cc" - "${tensorflow_source_dir}/tensorflow/contrib/tensorboard/db/loader.cc" - "${tensorflow_source_dir}/tensorflow/contrib/tensorboard/db/vacuum.cc" + "${tensorflow_source_dir}/tensorflow/core/summary/*test*.cc" + "${tensorflow_source_dir}/tensorflow/core/summary/loader.cc" + "${tensorflow_source_dir}/tensorflow/core/summary/vacuum.cc" ) # TODO(jart): Why doesn't this work? # set_source_files_properties( -# ${tensorflow_source_dir}/tensorflow/contrib/tensorboard/db/snapfn.cc +# ${tensorflow_source_dir}/tensorflow/core/lib/db/snapfn.cc # PROPERTIES COMPILE_FLAGS -DSQLITE_OMIT_LOAD_EXTENSION) list(REMOVE_ITEM tf_core_framework_srcs ${tf_core_framework_exclude_srcs}) diff --git a/tensorflow/contrib/cmake/tf_python.cmake b/tensorflow/contrib/cmake/tf_python.cmake index 8faccf8d55902e6701ebb4ce534b84705304fd5f..1fe8795ddf00232eba5a60a130e0845a6f6a8e17 100755 --- a/tensorflow/contrib/cmake/tf_python.cmake +++ b/tensorflow/contrib/cmake/tf_python.cmake @@ -802,6 +802,7 @@ add_custom_command( # tensorflow/__init__.py depends on files generated in this step. So, remove it while # this step is running since the files aren't there yet. COMMAND ${CMAKE_COMMAND} -E remove -f ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/__init__.py + COMMAND ${CMAKE_COMMAND} -E touch ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/__init__.py # Run create_python_api.py to generate API init files. COMMAND ${CMAKE_COMMAND} -E env PYTHONPATH=${CMAKE_CURRENT_BINARY_DIR}/tf_python "${PY_RUNTIME_ENV}" ${PYTHON_EXECUTABLE} diff --git a/tensorflow/contrib/compiler/xla.py b/tensorflow/contrib/compiler/xla.py index 2aa5d77b108d0f000550bd4d57fc4a922e7594e6..238c6ab1366a50710efabea2f33eb1bd06fe9423 100644 --- a/tensorflow/contrib/compiler/xla.py +++ b/tensorflow/contrib/compiler/xla.py @@ -85,7 +85,14 @@ def compile(computation, inputs=None): # pylint: disable=redefined-builtin with `tf.convert_to_tensor`. Returns: - A list of output tensors. + 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. """ # pylint: disable=protected-access return _compile_internal(computation, inputs) @@ -137,6 +144,30 @@ class XLACompileContext(control_flow_ops.XLAControlFlowContext): logging.warning('... and %d more', len(self._unsupported_ops) - _MAX_WARNING_LINES) + 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): """Create op in XLACompileContext and notifies outer context recursively.""" # pylint: disable=protected-access @@ -186,11 +217,14 @@ class XLACompileContext(control_flow_ops.XLAControlFlowContext): 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. - external_control_inputs = [ - array_ops.identity(x.outputs[0]).op - for x in external_control_inputs - if x.outputs - ] + 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 @@ -251,13 +285,21 @@ def _compile_internal(computation, inputs=None): Args: computation: A Python function that builds the computation to compile and execute. - inputs: A list of input tensors or `None` (equivalent to `[]`). Its order - should match ordering of computation arguments. + inputs: A list of inputs 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-dimension 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`. + Returns: - A list of output tensors from computation. + Same data structure as if computation(*inputs) is called directly with some + exceptions for correctness. Exceptions include: 1) None output 2) Single + value output 3) Operation-only outputs Raises: ValueError: If any element in computation outputs is neither an operations or a value that can be converted to tensor. + ValueError: If computation outputs is non-flat and contains any Operations. TypeError: If `inputs` is not a list or tuple. """ if inputs is None: @@ -300,66 +342,166 @@ def _compile_internal(computation, inputs=None): # Restore variable scope after computation. vscope.set_use_resource(saved_use_resource) - # If the computation returns `None`, make it an empty tuple. - if outputs is None: - outputs = tuple() - # If the computation only returned one value, make it a tuple. - if not isinstance(outputs, collections.Sequence): - outputs = (outputs,) - - # Append `no_op` here so that return value of this function always contains - # at least one op that can trigger XlaLaunch node. - outputs += (control_flow_ops.no_op(),) - try: - outputs = [ - o if isinstance(o, ops.Operation) else ops.convert_to_tensor(o) - for o in outputs - ] - except Exception as e: - raise ValueError( - 'XLA computation function return values must all either be Operations' - ' or convertible to Tensors. Got error: "%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( - 'XLA computation function must return zero or more Tensor values ' - 'followed by zero or more Operations.') - output_arity = len(output_tensors) - - new_output_tensors = [] - for t in output_tensors: - with ops.device(t.device if t.device else ''): - new_output_tensors.append(array_ops.identity(t)) + outputs_is_flat = 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) - output_tensors = new_output_tensors context.ExitResult(output_tensors) finally: context.report_unsupported_operations() context.Exit() - outputs = [ - xla_ops.xla_cluster_output(output_tensors[i], name='output{}'.format(i)) - for i in xrange(output_arity) + # When XLA computation returns only operations and no tensors, a NoOp + # dependent on the operations in outputs is returned. Otherwise final + # outputs would be empty and there is no way to trigger returned + # operations. + if not output_tensors: + return control_flow_ops.group(control_deps, name='output_0') + + output_tensors = [ + xla_ops.xla_cluster_output(o, name='output{}'.format(i)) + for i, o in enumerate(output_tensors) ] - with ops.control_dependencies(output_operations): - if output_arity == 0: - # When XLA computation returns only operations and no tensors, a NoOp - # dependent on the operations in outputs is returned. Otherwise final - # outputs would be empty and there is no way to trigger returned - # operations. - return control_flow_ops.no_op(name='output_0') - else: - # Wraps the outputs in identity operators that carries control - # dependencies. - return [ - array_ops.identity(outputs[i], name='output_%d' % i) - for i in xrange(output_arity) - ] + with ops.control_dependencies(control_deps): + # Wraps the outputs in identity operators that carries control + # dependencies. + output_tensors = [ + array_ops.identity(o, name='output_%d' % i) + for i, o in enumerate(output_tensors) + ] + + # If `computation` returned non-flat output structure, pack output tensors + # back into same structure. + if not outputs_is_flat: + output_tensors = nest.pack_sequence_as( + structure=outputs, flat_sequence=output_tensors) + + return output_tensors + + +def is_flat(outputs): + """Checks if outputs is a flat structure. + + Following structures and values are considered flat: + 1) None + 2) A single object + 3) A list or tuple of Tensors/Operations + + The only structures that this function understands are sequences and + dictionaries. E.g. this means that if outputs contains a single + user-defined Object, it is considered to be flat. Errors are raised later on + if that Object cannot be converted to a Tensor. + + Args: + outputs: Output from `computation` inside `xla.compile`. + + Returns: + A boolean indicates whether outputs is flat. + """ + # If outputs is a list or tuple, check if it has any nested structure. If + # there is, then outputs is non-flat. + if isinstance(outputs, collections.Sequence): + for o in outputs: + if isinstance(o, collections.Sequence) or isinstance(o, dict): + return False + + # If outputs is a dict, it is non-flat. + if isinstance(outputs, dict): + return False + + # Getting here means either outputs itself is a single non-structured value + # or it is a flat list of single non-structured values. + return True + + +def _postprocess_flat_outputs(outputs): + """Validates flat outputs and adds back device assignments. + + Args: + outputs: Output from `computation` inside `xla.compile`. + + 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, make it a tuple. + if not isinstance(outputs, collections.Sequence): + outputs = (outputs,) + + # Append `no_op` here so that return value of this function always contains + # at least one op that can trigger XlaLaunch node. + outputs += (control_flow_ops.no_op(),) + try: + outputs = [ + o if isinstance(o, ops.Operation) else ops.convert_to_tensor(o) + for o in outputs + ] + except Exception as e: + raise ValueError( + 'XLA computation function return values must all either be Operations' + ' or convertible to Tensors. Got error: "%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( + 'XLA computation function must return zero or more Tensor values ' + 'followed by zero or more Operations.') + + new_output_tensors = [] + for t in output_tensors: + with ops.device(t.device if t.device else ''): + new_output_tensors.append(array_ops.identity(t)) + + return new_output_tensors, output_operations + + +def _postprocess_non_flat_outputs(outputs): + """Validates non-flat outputs and adds back device assignments. + + Args: + outputs: Output from `computation` inside `xla.compile`. + + Returns: + Tensors extracted from outputs and an empty list because Operations are not + allowed in non-flat outputs.. + """ + # Convert all non-Operation outputs to Tensors. + new_output_tensors = [] + for o in nest.flatten(outputs): + if isinstance(o, ops.Operation): + raise ValueError( + 'xla.compile 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( + 'XLA computation function return values must all either be ' + 'Operations or convertible to Tensors. Got error: "%s"' % str(e)) + + # Makes sure even pass-through inputs/outputs are touched in compile + # context by creating an Identity node inside compile context. + with ops.device(o.device if o.device else ''): + new_output_tensors.append(array_ops.identity(o)) + + return new_output_tensors, [] @contextlib.contextmanager diff --git a/tensorflow/contrib/constrained_optimization/BUILD b/tensorflow/contrib/constrained_optimization/BUILD index eee4329acbeb38c9f37f79227aeb3acd46dce5e7..619153df67c90cea5a5082a411972948bac5fe90 100644 --- a/tensorflow/contrib/constrained_optimization/BUILD +++ b/tensorflow/contrib/constrained_optimization/BUILD @@ -42,11 +42,6 @@ py_test( name = "candidates_test", srcs = ["python/candidates_test.py"], srcs_version = "PY2AND3", - tags = [ - # TODO(b/121223093): Re-enable this test after fixing "Distribution - # should match known solution" errors. - "no_mac", - ], deps = [ ":constrained_optimization", "//tensorflow/python:client_testlib", diff --git a/tensorflow/contrib/constrained_optimization/README.md b/tensorflow/contrib/constrained_optimization/README.md index cb1dd7d836ae11700b2ffaaff4fda5b7f943f87d..7ffb6894d37444fd78015b6c124c46f2855c1cde 100644 --- a/tensorflow/contrib/constrained_optimization/README.md +++ b/tensorflow/contrib/constrained_optimization/README.md @@ -1,5 +1,10 @@ +**NOTE: As tensorflow.contrib is being +[deprecated](https://github.com/tensorflow/community/pull/18), TFCO is moving to +its own repository on +[github](https://github.com/google-research/tensorflow_constrained_optimization).** + # ConstrainedOptimization (TFCO) TFCO is a library for optimizing inequality-constrained problems in TensorFlow. diff --git a/tensorflow/contrib/constrained_optimization/python/candidates_test.py b/tensorflow/contrib/constrained_optimization/python/candidates_test.py index a4c49d48bc5c763489215261a909573af0f19055..280e9acd88638a9385bfd9128ba6d3739879aab2 100644 --- a/tensorflow/contrib/constrained_optimization/python/candidates_test.py +++ b/tensorflow/contrib/constrained_optimization/python/candidates_test.py @@ -52,12 +52,12 @@ class CandidatesTest(test.TestCase): distribution = candidates.find_best_candidate_distribution( objective_vector, constraints_matrix) # Verify that the solution is a probability distribution. - self.assertTrue(np.all(distribution >= 0)) + self.assertTrue(np.all(distribution >= -1e-6)) self.assertAlmostEqual(np.sum(distribution), 1.0) # Verify that the solution satisfies the constraints. maximum_constraint_violation = np.amax( np.dot(constraints_matrix, distribution)) - self.assertLessEqual(maximum_constraint_violation, 0) + self.assertLessEqual(maximum_constraint_violation, 1e-6) # Verify that the solution matches that which we expect. expected_distribution = np.array([0.37872711, 0.62127289, 0, 0]) self.assertAllClose(expected_distribution, distribution, rtol=0, atol=1e-6) diff --git a/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py b/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py index 7e1b4062ce435f3ab4216e90b4f5fcbab984c1dc..ca92c31236a7a3882415834eb32a994a120b6d2d 100644 --- a/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py +++ b/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_test.py @@ -1023,7 +1023,7 @@ class CudnnRNNTestCompatibleRNNCells(test_util.TensorFlowTestCase): outputs_v, output_state_v = sess.run( [outputs, output_state], feed_dict={cell_inputs: inference_input}) - self.assertAllClose(cudnn_outputs_v, outputs_v, atol=2e-5, rtol=2e-5) + self.assertAllClose(cudnn_outputs_v, outputs_v, atol=1e-4, rtol=2e-4) (cudnn_output_h_v,) = cudnn_output_states_v self.assertAllClose(cudnn_output_h_v, output_state_v, atol=2e-5, rtol=2e-5) diff --git a/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py b/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py index 86b48435d11e17bb3554f1d929a1418a01eb91d5..aa99ebfd6b4f72e6739c677224af5f93da923b7c 100644 --- a/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py +++ b/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py @@ -838,7 +838,7 @@ class CudnnLSTMSaveable(CudnnOpaqueParamsSaveable): checkpointable._track_checkpointable(bias, name="bias") # pylint: disable=protected-access assert len(biases) == len(weights) for cell_index, (bias, kernel) in enumerate(zip(biases, weights)): - cell = checkpointable_lib.Checkpointable() + cell = checkpointable_lib.AutoCheckpointable() checkpointable._track_checkpointable(cell, name="cell-%d" % cell_index) # pylint: disable=protected-access cell.bias = bias cell.kernel = kernel 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 6c5f8c6b00975b3fba041271309a93cecd9f5057..4db711c1f3f2815e7b8cf275af315c062ce4c02e 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 @@ -25,11 +25,13 @@ from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import tensor_shape +from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import script_ops from tensorflow.python.platform import test +@test_util.run_v1_only("deprecated API, no eager or V2 test coverage") class AssertElementShapeTest(test_base.DatasetTestBase): def test_assert_element_shape(self): diff --git a/tensorflow/contrib/data/python/kernel_tests/lmdb_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/lmdb_dataset_op_test.py index b9840b1ff1a3df5a05db0e64f436637220f49f80..220f9934b67d1d2a97f6c0fd4ba7779f011e1b09 100644 --- a/tensorflow/contrib/data/python/kernel_tests/lmdb_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/lmdb_dataset_op_test.py @@ -27,12 +27,14 @@ from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors +from tensorflow.python.framework import test_util from tensorflow.python.platform import test from tensorflow.python.util import compat prefix_path = "tensorflow/core/lib" +@test_util.run_v1_only("deprecated API, no eager or V2 test coverage") class LMDBDatasetTest(test_base.DatasetTestBase): def setUp(self): diff --git a/tensorflow/contrib/data/python/kernel_tests/reduce_dataset_test.py b/tensorflow/contrib/data/python/kernel_tests/reduce_dataset_test.py index e7281d531870c75c638b5c48fa3fc6dc606a3623..78019fcc7d810da444f1407f3885d54e76a741c6 100644 --- a/tensorflow/contrib/data/python/kernel_tests/reduce_dataset_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/reduce_dataset_test.py @@ -25,10 +25,12 @@ from tensorflow.contrib.data.python.ops import grouping from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import dtypes +from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.platform import test +@test_util.run_v1_only("deprecated API, no eager or V2 test coverage") class ReduceDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): @parameterized.named_parameters( diff --git a/tensorflow/contrib/data/python/kernel_tests/slide_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/slide_dataset_op_test.py index 2527706709fae8e459aca3489324d4db3c784be6..9275a36582a8c82b936659041129b71e100f883e 100644 --- a/tensorflow/contrib/data/python/kernel_tests/slide_dataset_op_test.py +++ b/tensorflow/contrib/data/python/kernel_tests/slide_dataset_op_test.py @@ -26,11 +26,13 @@ from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import sparse_tensor +from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.platform import test +@test_util.run_v1_only("deprecated API, no eager or V2 test coverage") class SlideDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): @parameterized.named_parameters( diff --git a/tensorflow/contrib/distribute/README.md b/tensorflow/contrib/distribute/README.md index 8a8dc159ade6f2a4a9b5ec29055ea4848492b29f..dbcaf8185fb7a9d2bcf22376439c0ebd49accb1a 100644 --- a/tensorflow/contrib/distribute/README.md +++ b/tensorflow/contrib/distribute/README.md @@ -43,28 +43,19 @@ the workers. Let's see how to scale to multiple GPUs on one machine using `MirroredStrategy` with [tf.keras] (https://www.tensorflow.org/guide/keras). -Take a very simple model consisting of a single layer: +Let's define a simple input dataset for training this model. Note that currently we require using +[`tf.data.Dataset`](https://www.tensorflow.org/api_docs/python/tf/data/Dataset) +with `DistributionStrategy`. ```python import tensorflow as tf from tensorflow import keras -inputs = tf.keras.layers.Input(shape=(1,)) -predictions = tf.keras.layers.Dense(1)(inputs) -model = tf.keras.models.Model(inputs=inputs, outputs=predictions) -``` - -Let's also define a simple input dataset for training this model. Note that currently we require using -[`tf.data.Dataset`](https://www.tensorflow.org/api_docs/python/tf/data/Dataset) -with `DistributionStrategy`. - -```python features = tf.data.Dataset.from_tensors([1.]).repeat(10000).batch(10) labels = tf.data.Dataset.from_tensors([1.]).repeat(10000).batch(10) train_dataset = tf.data.Dataset.zip((features, labels)) ``` - To distribute this Keras model on multiple GPUs using `MirroredStrategy` we first instantiate a `MirroredStrategy` object. @@ -72,14 +63,17 @@ first instantiate a `MirroredStrategy` object. distribution = tf.contrib.distribute.MirroredStrategy() ``` -We then compile the Keras model and pass the `MirroredStrategy` object in the -`distribute` argument (apart from other usual arguments like `loss` and -`optimizer`). +Take a very simple model consisting of a single layer. We need to create and compile +the model under the distribution strategy scope. ```python -model.compile(loss='mean_squared_error', - optimizer=tf.train.GradientDescentOptimizer(learning_rate=0.2), - distribute=distribution) +with distribution.scope(): + inputs = tf.keras.layers.Input(shape=(1,)) + predictions = tf.keras.layers.Dense(1)(inputs) + model = tf.keras.models.Model(inputs=inputs, outputs=predictions) + + model.compile(loss='mean_squared_error', + optimizer=tf.train.GradientDescentOptimizer(learning_rate=0.2)) ``` To train the model we call Keras `fit` API using the input dataset that we diff --git a/tensorflow/contrib/distribute/__init__.py b/tensorflow/contrib/distribute/__init__.py index 4c3f9b8f02087bbc442f8082945f9ef7da961a6a..59d76f5d1c817d7f2cc8ad285b9fb517fe994a81 100644 --- a/tensorflow/contrib/distribute/__init__.py +++ b/tensorflow/contrib/distribute/__init__.py @@ -35,8 +35,8 @@ from tensorflow.contrib.distribute.python.tpu_strategy import TPUStrategy from tensorflow.python.distribute.cross_device_ops import * from tensorflow.python.distribute.distribute_config import DistributeConfig from tensorflow.python.distribute.distribute_coordinator import run_standard_tensorflow_server -from tensorflow.python.training.distribute import * -from tensorflow.python.training.distribution_strategy_context import * +from tensorflow.python.distribute.distribute_lib import * +from tensorflow.python.distribute.distribution_strategy_context import * from tensorflow.python.util.all_util import remove_undocumented @@ -64,7 +64,9 @@ _allowed_symbols = [ 'get_distribution_strategy', 'get_loss_reduction', 'get_replica_context', + 'get_strategy', 'has_distribution_strategy', + 'has_strategy', 'in_cross_replica_context', 'require_replica_context', 'run_standard_tensorflow_server', diff --git a/tensorflow/contrib/distribute/python/BUILD b/tensorflow/contrib/distribute/python/BUILD index 78ad3214fb68d8779f3b8535d100d6a6f917a2a2..44ecc8c4286d594e40378e6811a085ade73cea84 100644 --- a/tensorflow/contrib/distribute/python/BUILD +++ b/tensorflow/contrib/distribute/python/BUILD @@ -1,6 +1,7 @@ # Implementation of a prototype TF distributed computation library. load("//tensorflow/compiler/tests:build_defs.bzl", "tf_xla_py_test") +load("//tensorflow/core:platform/default/distribute.bzl", "distribute_py_test") load("//tensorflow:tensorflow.bzl", "py_test") load("//tensorflow:tensorflow.bzl", "cuda_py_test") @@ -14,8 +15,18 @@ licenses(["notice"]) # Apache 2.0 exports_files(["LICENSE"]) -# TODO(priyag): Figure out testonly issues that are preventing us from -# including our tests in pip for now. +py_library( + name = "distribute_test_lib_pip", + visibility = ["//tensorflow:internal"], + deps = [ + ":combinations", + ":keras_correctness_test_lib", + ":keras_test_lib", + ":multi_worker_test_base", + ":single_loss_example", + ":strategy_test_lib", + ], +) cuda_py_test( name = "values_test", @@ -37,9 +48,6 @@ cuda_py_test( "//tensorflow/python/eager:test", "//tensorflow/python/estimator:estimator_py", ], - tags = [ - "no_pip", - ], ) cuda_py_test( @@ -50,18 +58,13 @@ cuda_py_test( ":mirrored_strategy", ":multi_worker_test_base", "@absl_py//absl/testing:parameterized", - "//tensorflow/core:protos_all_py", "//tensorflow/python:errors", - "//tensorflow/python:framework_test_lib", "//tensorflow/python/data/ops:dataset_ops", "//tensorflow/python/distribute:input_lib", "//tensorflow/python/distribute:values", "//tensorflow/python/eager:context", "//tensorflow/python/eager:test", ], - tags = [ - "no_pip", - ], ) py_library( @@ -80,19 +83,10 @@ py_library( srcs = ["parameter_server_strategy.py"], visibility = ["//tensorflow:internal"], deps = [ - ":mirrored_strategy", - "//tensorflow/core:protos_all_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:framework_ops", - "//tensorflow/python:resource_variable_ops", - "//tensorflow/python:training", - "//tensorflow/python:util", - "//tensorflow/python/distribute:cross_device_ops", - "//tensorflow/python/distribute:input_lib", - "//tensorflow/python/distribute:multi_worker_util", - "//tensorflow/python/distribute:reduce_util", + "//tensorflow/python/distribute:distribute_lib", + "//tensorflow/python/distribute:parameter_server_strategy", "//tensorflow/python/distribute:values", - "//tensorflow/python/eager:context", + "//tensorflow/python/distribute/cluster_resolver:cluster_resolver_lib", ], ) @@ -125,7 +119,6 @@ cuda_py_test( ], tags = [ "multi_and_single_gpu", - "no_pip", ], ) @@ -134,16 +127,17 @@ py_library( srcs = ["one_device_strategy.py"], visibility = ["//tensorflow:internal"], deps = [ - "//tensorflow/python:array_ops", - "//tensorflow/python:dtypes", - "//tensorflow/python:framework_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python/distribute:distribute_lib", - "//tensorflow/python/distribute:input_lib", - "//tensorflow/python/distribute:reduce_util", - "//tensorflow/python/distribute:values", - "//tensorflow/python/eager:context", - "@six_archive//:six", + "//tensorflow/python/distribute:one_device_strategy", + ], +) + +cuda_py_test( + name = "one_device_strategy_test", + srcs = ["one_device_strategy_test.py"], + additional_deps = [ + ":strategy_test_lib", + ":combinations", + "//tensorflow/python/eager:test", ], ) @@ -152,29 +146,16 @@ py_library( srcs = ["collective_all_reduce_strategy.py"], visibility = ["//tensorflow:internal"], deps = [ - ":mirrored_strategy", - "//tensorflow/core:protos_all_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:collective_ops", - "//tensorflow/python:framework_ops", - "//tensorflow/python:training", - "//tensorflow/python/distribute:cross_device_ops", - "//tensorflow/python/distribute:cross_device_utils", - "//tensorflow/python/distribute:input_lib", - "//tensorflow/python/distribute:multi_worker_util", - "//tensorflow/python/distribute:values", - "//tensorflow/python/eager:context", + "//tensorflow/python/distribute:collective_all_reduce_strategy", + "//tensorflow/python/distribute:distribute_lib", + "//tensorflow/python/distribute/cluster_resolver:cluster_resolver_lib", ], ) py_library( name = "strategy_test_lib", - testonly = 1, srcs = ["strategy_test_lib.py"], srcs_version = "PY2AND3", - tags = [ - "no_pip", - ], deps = [ "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", @@ -187,20 +168,18 @@ py_library( "//tensorflow/python/eager:backprop", "//tensorflow/python/eager:context", "//tensorflow/python/eager:test", + "//third_party/py/numpy", ], ) py_library( name = "combinations", - testonly = 1, srcs = ["combinations.py"], srcs_version = "PY2AND3", - tags = [ - "no_pip", - ], deps = [ ":mirrored_strategy", ":one_device_strategy", + ":parameter_server_strategy", ":tpu_strategy", "//tensorflow/contrib/cluster_resolver:cluster_resolver_pip", "//tensorflow/contrib/optimizer_v2:training", @@ -209,6 +188,7 @@ py_library( "//tensorflow/python:util", "//tensorflow/python/distribute:distribute_lib", "//tensorflow/python/eager:context", + "//tensorflow/python/keras/optimizer_v2", "@absl_py//absl/testing:parameterized", ], ) @@ -216,30 +196,12 @@ py_library( py_test( name = "combinations_test", srcs = ["combinations_test.py"], - tags = [ - "no_pip", - ], deps = [ ":combinations", "//tensorflow/python/eager:test", ], ) -py_test( - name = "one_device_strategy_test", - srcs = ["one_device_strategy_test.py"], - srcs_version = "PY2AND3", - tags = [ - "no_pip", - ], - deps = [ - ":one_device_strategy", - ":strategy_test_lib", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python/eager:test", - ], -) - # TODO(priyag): Rename this test to mirrored_strategy_test cuda_py_test( name = "mirrored_strategy_multigpu_test", @@ -265,18 +227,13 @@ cuda_py_test( tags = [ "guitar", "multi_and_single_gpu", - "no_pip", ], ) py_library( name = "multi_worker_test_base", - testonly = 1, srcs = ["multi_worker_test_base.py"], srcs_version = "PY2AND3", - tags = [ - "no_pip", - ], deps = [ "//tensorflow/core:protos_all_py", "//tensorflow/python:client_testlib", @@ -312,6 +269,7 @@ py_library( "//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", ], @@ -344,14 +302,16 @@ cuda_py_test( ], tags = [ "multi_and_single_gpu", - "no_pip", ], ) -py_library( - name = "minimize_loss_test_lib", - testonly = 1, +distribute_py_test( + name = "minimize_loss_test", srcs = ["minimize_loss_test.py"], + main = "minimize_loss_test.py", + tags = [ + "multi_and_single_gpu", + ], deps = [ ":combinations", ":mirrored_strategy", @@ -371,18 +331,6 @@ py_library( ], ) -cuda_py_test( - name = "minimize_loss_test", - srcs = ["minimize_loss_test.py"], - additional_deps = [ - ":minimize_loss_test_lib", - ], - tags = [ - "multi_and_single_gpu", - "no_pip", - ], -) - cuda_py_test( name = "moving_averages_test", srcs = ["moving_averages_test.py"], @@ -396,9 +344,6 @@ cuda_py_test( "//tensorflow/python:training", "//tensorflow/python:variables", ], - tags = [ - "no_pip", - ], ) cuda_py_test( @@ -416,7 +361,6 @@ cuda_py_test( ], tags = [ "multi_and_single_gpu", - "no_pip", ], ) @@ -439,7 +383,6 @@ cuda_py_test( tags = [ "multi_and_single_gpu", "no_oss", # http://b/119349471 - "no_pip", "tf_integration_test", ], ) @@ -450,10 +393,10 @@ cuda_py_test( additional_deps = [ ":keras_test_lib", ], + shard_count = 4, tags = [ "multi_and_single_gpu", "no_oss", # http://b/119349471 - "no_pip", "tf_integration_test", ], ) @@ -483,7 +426,6 @@ cuda_py_test( shard_count = 48, tags = [ "multi_and_single_gpu", - "no_pip", # TODO(b/118768923): Re-enable {a,m,t}san test. "noasan", "nomsan", @@ -505,10 +447,13 @@ py_library( ], ) -py_library( - name = "step_fn_test_lib", - testonly = 1, +distribute_py_test( + name = "step_fn_test", srcs = ["step_fn_test.py"], + main = "step_fn_test.py", + tags = [ + "multi_and_single_gpu", + ], deps = [ ":combinations", ":single_loss_example", @@ -521,18 +466,6 @@ py_library( ], ) -cuda_py_test( - name = "step_fn_test", - srcs = ["step_fn_test.py"], - additional_deps = [ - ":step_fn_test_lib", - ], - tags = [ - "multi_and_single_gpu", - "no_pip", - ], -) - py_library( name = "monitor", srcs = ["monitor.py"], @@ -549,10 +482,10 @@ cuda_py_test( additional_deps = [ ":combinations", ":monitor", - ":one_device_strategy", ":single_loss_example", "@absl_py//absl/testing:parameterized", "//third_party/py/numpy", + "//tensorflow/python/distribute:one_device_strategy", "//tensorflow/python/eager:context", "//tensorflow/python/eager:test", "//tensorflow/python:framework_ops", @@ -560,7 +493,6 @@ cuda_py_test( ], tags = [ "multi_and_single_gpu", - "no_pip", ], ) @@ -577,9 +509,6 @@ cuda_py_test( "//tensorflow/python/eager:context", "//tensorflow/python/eager:test", ], - tags = [ - "no_pip", - ], ) cuda_py_test( @@ -601,16 +530,15 @@ cuda_py_test( ], tags = [ "multi_and_single_gpu", - "no_pip", ], ) py_library( name = "keras_test_lib", - testonly = 1, srcs = [ "keras_backward_compat_test.py", "keras_test.py", + "keras_utils_test.py", ], deps = [ ":combinations", @@ -626,43 +554,69 @@ py_library( ], ) -cuda_py_test( +distribute_py_test( name = "keras_test", srcs = ["keras_test.py"], - additional_deps = [ + full_precision = True, + main = "keras_test.py", + shard_count = 32, + tags = [ + "multi_and_single_gpu", + "no_oss", # TODO(b/117919883): Fix python error. + "no_windows_gpu", + "notsan", + ], + deps = [ ":keras_test_lib", ], - shard_count = 16, +) + +distribute_py_test( + name = "keras_utils_test", + srcs = ["keras_utils_test.py"], + full_precision = True, + main = "keras_utils_test.py", + shard_count = 32, tags = [ "multi_and_single_gpu", "no_oss", # TODO(b/117919883): Fix python error. - "no_pip", "no_windows_gpu", "notsan", ], + deps = [ + ":keras_test", + ":keras_test_lib", + ], ) # TODO(b/121200287): Remove this in 2.0 -cuda_py_test( +distribute_py_test( name = "keras_backward_compat_test", srcs = ["keras_backward_compat_test.py"], - additional_deps = [ - ":keras_test_lib", - ], - shard_count = 16, + full_precision = True, + main = "keras_backward_compat_test.py", + shard_count = 31, tags = [ "multi_and_single_gpu", "no_oss", # TODO(b/117919883): Fix python error. - "no_pip", "no_windows_gpu", "notsan", ], + deps = [ + ":keras_test_lib", + ], ) py_library( name = "keras_correctness_test_lib", - testonly = 1, - srcs = ["keras_correctness_test.py"], + srcs = [ + "keras_correctness_test_base.py", + "keras_dnn_correctness_test.py", + "keras_embedding_model_correctness_test.py", + "keras_image_model_correctness_test.py", + "keras_lstm_model_correctness_test.py", + "keras_stateful_lstm_model_correctness_test.py", + ], deps = [ ":combinations", "//tensorflow/contrib/distribute/python:mirrored_strategy", @@ -677,13 +631,95 @@ py_library( ], ) -cuda_py_test( - name = "keras_correctness_test", - srcs = ["keras_correctness_test.py"], - additional_deps = [ +distribute_py_test( + name = "keras_dnn_correctness_test", + size = "medium", + srcs = ["keras_dnn_correctness_test.py"], + full_precision = True, + main = "keras_dnn_correctness_test.py", + # Shard count is set to an odd number to distribute tasks across + # shards more evenly. + shard_count = 19, + tags = [ + "multi_and_single_gpu", + "no_oss", # TODO(b/117919883): Fix python error. + "no_windows_gpu", + "notsan", + ], + deps = [ ":keras_correctness_test_lib", ], - shard_count = 16, +) + +distribute_py_test( + name = "keras_image_model_correctness_test", + size = "medium", + srcs = ["keras_image_model_correctness_test.py"], + full_precision = True, + main = "keras_image_model_correctness_test.py", + # Shard count is set to an odd number to distribute tasks across + # shards more evenly. + shard_count = 31, + tags = [ + "multi_and_single_gpu", + "no_oss", # TODO(b/117919883): Fix python error. + "no_windows_gpu", + "notsan", + ], + deps = [ + ":keras_correctness_test_lib", + ], +) + +distribute_py_test( + name = "keras_embedding_model_correctness_test", + size = "medium", + srcs = ["keras_embedding_model_correctness_test.py"], + full_precision = True, + main = "keras_embedding_model_correctness_test.py", + # Shard count is set to an odd number to distribute tasks across + # shards more evenly. + shard_count = 31, + tags = [ + "multi_and_single_gpu", + "no_oss", # TODO(b/117919883): Fix python error. + "no_windows_gpu", + "notsan", + ], + deps = [ + ":keras_correctness_test_lib", + ], +) + +distribute_py_test( + name = "keras_lstm_model_correctness_test", + size = "medium", + srcs = ["keras_lstm_model_correctness_test.py"], + full_precision = True, + main = "keras_lstm_model_correctness_test.py", + # Shard count is set to an odd number to distribute tasks across + # shards more evenly. + shard_count = 31, + tags = [ + "multi_and_single_gpu", + "no_oss", # TODO(b/117919883): Fix python error. + "no_windows_gpu", + "notsan", + ], + deps = [ + ":keras_correctness_test_lib", + ], +) + +distribute_py_test( + name = "keras_stateful_lstm_model_correctness_test", + size = "medium", + srcs = ["keras_stateful_lstm_model_correctness_test.py"], + full_precision = True, + main = "keras_stateful_lstm_model_correctness_test.py", + # Shard count is set to an odd number to distribute tasks across + # shards more evenly. + shard_count = 31, tags = [ "multi_and_single_gpu", "no_oss", # TODO(b/117919883): Fix python error. @@ -691,12 +727,18 @@ cuda_py_test( "no_windows_gpu", "notsan", ], + deps = [ + ":keras_correctness_test_lib", + ], ) -py_library( - name = "metrics_v1_test_lib", - testonly = 1, +distribute_py_test( + name = "metrics_v1_test", srcs = ["metrics_v1_test.py"], + main = "metrics_v1_test.py", + tags = [ + "multi_and_single_gpu", + ], deps = [ ":combinations", "//tensorflow/python:math_ops", @@ -708,18 +750,6 @@ py_library( ], ) -cuda_py_test( - name = "metrics_v1_test", - srcs = ["metrics_v1_test.py"], - additional_deps = [ - ":metrics_v1_test_lib", - ], - tags = [ - "multi_and_single_gpu", - "no_pip", - ], -) - cuda_py_test( name = "warm_starting_util_test", size = "medium", @@ -734,7 +764,6 @@ cuda_py_test( ], tags = [ "multi_and_single_gpu", - "no_pip", ], ) @@ -745,7 +774,6 @@ cuda_py_test( additional_deps = [ ":combinations", "//tensorflow/python:client_testlib", - "//tensorflow/python:checkpoint_utils_test", "//tensorflow/python:framework_ops", "//tensorflow/python:training", "//tensorflow/python:variable_scope", @@ -753,7 +781,6 @@ cuda_py_test( ], tags = [ "multi_and_single_gpu", - "no_pip", ], ) @@ -768,7 +795,6 @@ tf_xla_py_test( ], tags = [ "no_oss", - "no_pip", ], deps = [ ":tpu_strategy", diff --git a/tensorflow/contrib/distribute/python/checkpoint_utils_test.py b/tensorflow/contrib/distribute/python/checkpoint_utils_test.py index 31bd0e996a247a2fc01405fb3b8172a40853d698..7ee50f03155636a487020d0a9178107a06775588 100644 --- a/tensorflow/contrib/distribute/python/checkpoint_utils_test.py +++ b/tensorflow/contrib/distribute/python/checkpoint_utils_test.py @@ -25,6 +25,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import os from absl.testing import parameterized from tensorflow.contrib.distribute.python import combinations @@ -33,7 +34,23 @@ from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables from tensorflow.python.platform import test from tensorflow.python.training import checkpoint_utils -from tensorflow.python.training import checkpoint_utils_test +from tensorflow.python.training import saver as saver_lib + + +def _create_checkpoints(sess, checkpoint_dir): + checkpoint_prefix = os.path.join(checkpoint_dir, "model") + checkpoint_state_name = "checkpoint" + v1 = variable_scope.get_variable("var1", [1, 10]) + v2 = variable_scope.get_variable("var2", [10, 10]) + sess.run(variables.global_variables_initializer()) + v1_value, v2_value = sess.run([v1, v2]) + saver = saver_lib.Saver() + saver.save( + sess, + checkpoint_prefix, + global_step=0, + latest_filename=checkpoint_state_name) + return v1_value, v2_value class CheckpointUtilsWithDistributionStrategyTest( @@ -51,8 +68,7 @@ class CheckpointUtilsWithDistributionStrategyTest( def testInitFromCheckpoint(self, distribution, in_replica_mode): checkpoint_dir = self.get_temp_dir() with self.cached_session() as session: - v1_value, v2_value, _, _ = checkpoint_utils_test._create_checkpoints( - session, checkpoint_dir) + v1_value, v2_value = _create_checkpoints(session, checkpoint_dir) def init_and_verify(g): v1 = variable_scope.get_variable("new_var1", [1, 10]) @@ -71,7 +87,7 @@ class CheckpointUtilsWithDistributionStrategyTest( with ops.Graph().as_default() as g, distribution.scope(): if in_replica_mode: - distribution.call_for_each_replica(init_and_verify, args=[g]) + distribution.extended.call_for_each_replica(init_and_verify, args=[g]) else: init_and_verify(g) diff --git a/tensorflow/contrib/distribute/python/checkpointing_test.py b/tensorflow/contrib/distribute/python/checkpointing_test.py index f37026bc95e02f58be20219f48bf37918024cdb2..aa5b9f57b8a5bc12ee94399ec1fc5a55177a5b5d 100644 --- a/tensorflow/contrib/distribute/python/checkpointing_test.py +++ b/tensorflow/contrib/distribute/python/checkpointing_test.py @@ -34,7 +34,7 @@ from tensorflow.python.training.checkpointable import tracking from tensorflow.python.training.checkpointable import util as checkpointable_utils -class NonLayerCheckpointable(tracking.Checkpointable): +class NonLayerCheckpointable(tracking.AutoCheckpointable): def __init__(self): super(NonLayerCheckpointable, self).__init__() diff --git a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py index f6361cb6e89373526dc7c554fd9e5b980d5f0382..19741627980c34d8c281f7aed6f1464d4a03393e 100644 --- a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py +++ b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py @@ -18,28 +18,18 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import copy - -from tensorflow.contrib.distribute.python import mirrored_strategy -from tensorflow.core.protobuf import rewriter_config_pb2 -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 device_util +from tensorflow.python.distribute import collective_all_reduce_strategy from tensorflow.python.distribute import distribute_lib -from tensorflow.python.distribute import input_lib -from tensorflow.python.distribute import multi_worker_util -from tensorflow.python.distribute import values -from tensorflow.python.eager import context -from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import collective_ops -from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.distribute.cluster_resolver import SimpleClusterResolver +from tensorflow.python.distribute.cluster_resolver import TFConfigClusterResolver # TODO(yuefengz): support in-graph replication. class CollectiveAllReduceStrategy(distribute_lib.DistributionStrategy): """Distribution strategy that uses collective ops for all-reduce. + *** contrib version *** + It is similar to the MirroredStrategy but it uses collective ops for reduction. @@ -62,289 +52,19 @@ class CollectiveAllReduceStrategy(distribute_lib.DistributionStrategy): CollectiveAllReduceExtended(self, num_gpus_per_worker)) -class CollectiveAllReduceExtended(mirrored_strategy.MirroredExtended): +class CollectiveAllReduceExtended( + collective_all_reduce_strategy.CollectiveAllReduceExtended): """Implementation of CollectiveAllReduceStrategy.""" def __init__(self, container_strategy, num_gpus_per_worker): - distribute_lib.DistributionStrategyExtended.__init__( - self, container_strategy) - self._cross_device_ops = None - self._num_gpus_per_worker = num_gpus_per_worker - self._initialize_local_worker(num_gpus_per_worker) - assert isinstance(self._get_cross_device_ops(), - cross_device_ops_lib.CollectiveAllReduce) - - def _initialize_local_worker(self, num_gpus_per_worker): - """Initializes the object for local training.""" - self._is_chief = True - self._num_workers = 1 - - if num_gpus_per_worker: - local_devices = tuple( - "/device:GPU:%d" % i for i in range(num_gpus_per_worker) - ) - else: - local_devices = ("/device:CPU:0",) - self._worker_device = device_util.canonicalize("/device:CPU:0") - - self._collective_keys = cross_device_utils.CollectiveKeys() - self._initialize_local(local_devices) - self._cross_device_ops = cross_device_ops_lib.CollectiveAllReduce( - num_workers=self._num_workers, - num_gpus_per_worker=num_gpus_per_worker, - collective_keys=self._collective_keys) - - self._cluster_spec = None - self._task_type = None - self._task_id = None - - logging.info("CollectiveAllReduceStrategy with local_devices = %r", - local_devices) - - def _initialize_multi_worker(self, num_gpus_per_worker, cluster_spec, - task_type, task_id): - """Initializes the object for multi-worker training.""" - if task_type is None or task_id is None: - raise ValueError("When `cluster_spec` is given, you must also specify " - "`task_type` and `task_id`") - if task_type not in ("chief", "worker"): - raise ValueError( - "Unrecognized task_type: %r, valid task types are: \"chief\", " - "\"worker\"." % task_type) - cluster_spec = multi_worker_util.normalize_cluster_spec(cluster_spec) - self._num_workers = multi_worker_util.worker_count(cluster_spec, task_type) - if not self._num_workers: - raise ValueError("No `worker` or `chief` tasks can be found in " - "`cluster_spec`.") - - self._is_chief = multi_worker_util.is_chief(cluster_spec, task_type, - task_id) - - self._worker_device = "/job:%s/task:%d" % (task_type, task_id) - if num_gpus_per_worker: - local_devices = tuple( - "%s/device:GPU:%d" % (self._worker_device, i) - for i in range(num_gpus_per_worker) - ) - else: - local_devices = (self._worker_device,) - - self._collective_keys = cross_device_utils.CollectiveKeys() - self._initialize_local(local_devices) - self._input_workers = input_lib.InputWorkers( - self._device_map, [(self._worker_device, self.worker_devices)]) - self._cross_device_ops = cross_device_ops_lib.CollectiveAllReduce( - num_workers=self._num_workers, - num_gpus_per_worker=num_gpus_per_worker, - collective_keys=self._collective_keys) - - # Add a default device so that ops without specified devices will not end up - # on other workers. - self._default_device = "/job:%s/task:%d" % (task_type, task_id) - - self._cluster_spec = multi_worker_util.normalize_cluster_spec(cluster_spec) - self._task_type = task_type - self._task_id = task_id - - logging.info( - "Multi-worker CollectiveAllReduceStrategy with " - "cluster_spec = %r, task_type = %r, task_id = %r, " - "num_workers = %r, local_devices = %r", cluster_spec.as_dict(), - task_type, task_id, self._num_workers, local_devices) - - def _create_variable(self, next_creator, *args, **kwargs): - 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. - else: - device_map = colocate_with.device_map - logical_device = colocate_with.logical_device - group_size = device_map.num_replicas_in_graph * self._num_workers - group_key = self._collective_keys.get_group_key(self.worker_devices) - - def _real_mirrored_creator(devices, *args, **kwargs): - """Creates one MirroredVariable on the current worker.""" - value_list = [] - unique_var_name = ops.get_default_graph().unique_name( - kwargs["name"], mark_as_used=False).rstrip("/") - collective_instance_key = self._collective_keys.get_instance_key( - key_id=unique_var_name) - if "initial_value" not in kwargs: - raise ValueError("Initial value must be specified.") - initial_value = kwargs["initial_value"] - if callable(initial_value): - initial_value_fn = initial_value - else: - initial_value_fn = lambda: initial_value - - 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) - - # The initial value fn makes sure variables all initialized to - # same values. The first device of the chief worker will send their - # variable values to other devices and other workers. - def _overridden_initial_value_fn(device=d, index=i): # pylint: disable=g-missing-docstring - with ops.device(device): - initial_value = initial_value_fn() - assert not callable(initial_value) - initial_value = ops.convert_to_tensor(initial_value) - - if self._is_chief and index == 0: - bcast_send = collective_ops.broadcast_send( - initial_value, initial_value.shape, initial_value.dtype, - group_size, group_key, collective_instance_key) - with ops.control_dependencies([bcast_send]): - return array_ops.identity(initial_value) - else: - return collective_ops.broadcast_recv( - initial_value.shape, initial_value.dtype, group_size, - group_key, collective_instance_key) - - kwargs["initial_value"] = _overridden_initial_value_fn - - with context.context().device_policy(context.DEVICE_PLACEMENT_SILENT): - v = next_creator(*args, **kwargs) - - if i == 0: - actual_var_name = v.name.split(":")[0] - assert unique_var_name == actual_var_name, "%r vs %r" % ( - unique_var_name, actual_var_name) - assert not isinstance(v, values.DistributedVariable) - value_list.append(v) - return value_list - - # pylint: disable=protected-access - return mirrored_strategy._create_mirrored_variable( - self._container_strategy(), device_map, logical_device, - _real_mirrored_creator, *args, **kwargs) - - def _distribute_dataset(self, dataset_fn): - """Distributes the dataset to each local GPU.""" - # TODO(yuefengz): shard the dataset. - worker_index = 0 - return input_lib.PerReplicaDataset( - self._call_dataset_fn(dataset_fn), self._input_workers, worker_index, - prefetch_on_device=True) - - def _make_dataset_iterator(self, dataset): - 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): - """Distributes the dataset to each local GPU.""" - if self._cluster_spec is None: - input_pipeline_id = 0 - else: - input_pipeline_id = multi_worker_util.id_in_cluster( - self._cluster_spec, self._task_type, self._task_id) - input_context = distribute_lib.InputContext( - num_input_pipelines=self._num_workers, - input_pipeline_id=input_pipeline_id, - num_replicas_in_sync=self._num_replicas_in_sync) - - return input_lib.InputFunctionIterator( - input_fn, self._input_workers, [input_context]) - - def _configure(self, - session_config=None, - cluster_spec=None, - task_type=None, - task_id=None): - """Configures the object. - - Args: - session_config: a `tf.ConfigProto` - cluster_spec: a dict, ClusterDef or ClusterSpec object specifying the - cluster configurations. - task_type: the current task type, such as "worker". - task_id: the current task id. - - Raises: - ValueError: if `task_type` is not in the `cluster_spec`. - """ - if not self._cluster_spec and cluster_spec: - # If a `cluster_spec` is already passed in, do nothing here. - # TODO(yuefengz): check `cluster_spec` is the same if this object has - # already been initialized with a `cluster_spec`. - self._initialize_multi_worker(self._num_gpus_per_worker, cluster_spec, - task_type, task_id) - assert isinstance(self._get_cross_device_ops(), - cross_device_ops_lib.CollectiveAllReduce) - - 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) - # Enable the scoped allocator optimization for CollectiveOps. This - # optimization converts many small all-reduces into fewer larger - # all-reduces. - rewrite_options = updated_config.graph_options.rewrite_options - rewrite_options.scoped_allocator_optimization = ( - rewriter_config_pb2.RewriterConfig.ON) - # We turn on ScopedAllocator only for CollectiveReduce op, i.e. enable_op = - # ["CollectiveReduce"]. Since we can't assign to a repeated proto field, we - # clear and then append. - del rewrite_options.scoped_allocator_opts.enable_op[:] - rewrite_options.scoped_allocator_opts.enable_op.append("CollectiveReduce") - - if not self._cluster_spec: - return updated_config - - assert self._task_type - assert self._task_id is not None - - # Collective group leader is needed for collective ops to coordinate - # workers. - if "chief" in self._cluster_spec.jobs: - updated_config.experimental.collective_group_leader = ( - "/job:chief/replica:0/task:0") - else: - if "worker" not in self._cluster_spec.jobs: - raise ValueError( - "You must have `chief` or `worker` jobs in the `cluster_spec`.") - updated_config.experimental.collective_group_leader = ( - "/job:worker/replica:0/task:0") - - # The device filters prevent communication between workers. - del updated_config.device_filters[:] - updated_config.device_filters.append( - "/job:%s/task:%d" % (self._task_type, self._task_id)) - - return updated_config - - @property - def experimental_between_graph(self): - return True - - @property - def experimental_should_init(self): - return True - - @property - def should_checkpoint(self): - return self._is_chief - - @property - def should_save_summary(self): - return self._is_chief - - @property - def _num_replicas_in_sync(self): - return len(self.worker_devices) * self._num_workers - - # TODO(priyag): Delete this once all strategies use global batch size. - @property - def _global_batch_size(self): - return False + # 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. + tfconfig = TFConfigClusterResolver() + cluster_resolver = SimpleClusterResolver( + cluster_spec=tfconfig.cluster_spec(), + task_type=tfconfig.task_type, + task_id=tfconfig.task_id, + num_accelerators=num_gpus_per_worker) + super(CollectiveAllReduceExtended, self).__init__( + container_strategy, 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 0fb672dded7624e798592d2f5c01945aa830021e..acbe4677b401cbea4fd0ec415415f25c920e68e4 100644 --- a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py +++ b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py @@ -29,9 +29,13 @@ from tensorflow.core.protobuf import config_pb2 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_utils +from tensorflow.python.distribute import distribute_lib +from tensorflow.python.distribute import multi_worker_util from tensorflow.python.distribute import reduce_util from tensorflow.python.distribute import values +from tensorflow.python.distribute.cluster_resolver import SimpleClusterResolver from tensorflow.python.eager import context from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes @@ -49,6 +53,55 @@ from tensorflow.python.ops.losses import losses from tensorflow.python.platform import test from tensorflow.python.training import adam from tensorflow.python.training import training_util +from tensorflow.python.training.server_lib import ClusterSpec + + +class MockCollectiveAllReduceStrategy(distribute_lib.DistributionStrategy): + """Mock the strategy to allow cluster resolver as an argument.""" + + def __init__(self, cluster_resolver): + super(MockCollectiveAllReduceStrategy, self).__init__( + core_collective_all_reduce_strategy.CollectiveAllReduceExtended( + self, cluster_resolver=cluster_resolver)) + + +def create_test_objects(cluster_spec=None, + task_type=None, + task_id=None, + num_gpus=None, + use_core_strategy=False): + sess_config = config_pb2.ConfigProto() + if num_gpus is None: + num_gpus = context.num_gpus() + if use_core_strategy: + if cluster_spec and task_type and task_id is not None: + cluster_resolver = SimpleClusterResolver( + cluster_spec=multi_worker_util.normalize_cluster_spec(cluster_spec), + task_type=task_type, + task_id=task_id, + num_accelerators=num_gpus) + target = 'grpc://' + cluster_spec[task_type][task_id] + else: + cluster_resolver = SimpleClusterResolver( + ClusterSpec({}), num_accelerators=num_gpus) + target = '' + + strategy = MockCollectiveAllReduceStrategy(cluster_resolver) + sess_config = strategy.update_config_proto(sess_config) + else: + strategy = collective_all_reduce_strategy.CollectiveAllReduceStrategy( + num_gpus_per_worker=num_gpus) + if task_type and task_id is not None: + strategy.configure( + session_config=sess_config, + cluster_spec=cluster_spec, + task_type=task_type, + task_id=task_id) + target = 'grpc://' + cluster_spec[task_type][task_id] + else: + target = '' + + return strategy, target, sess_config class CollectiveAllReduceStrategyTestBase( @@ -64,16 +117,18 @@ class CollectiveAllReduceStrategyTestBase( CollectiveAllReduceStrategyTestBase.collective_key_base += 100000 super(CollectiveAllReduceStrategyTestBase, self).setUp() - def _get_test_object(self, task_type, task_id, num_gpus=0): - distribution = collective_all_reduce_strategy.CollectiveAllReduceStrategy( - num_gpus_per_worker=num_gpus) - session_config = config_pb2.ConfigProto() - if task_type and task_id is not None: - distribution.configure( - session_config=session_config, - cluster_spec=self._cluster_spec, - task_type=task_type, - task_id=task_id) + def _get_test_object(self, + task_type, + task_id, + num_gpus=0, + use_core_strategy=False): + strategy, target, session_config = create_test_objects( + cluster_spec=self._cluster_spec, + task_type=task_type, + task_id=task_id, + num_gpus=num_gpus, + use_core_strategy=use_core_strategy) + collective_keys = cross_device_utils.CollectiveKeys( group_key_start=10 * num_gpus + CollectiveAllReduceStrategyTestBase.collective_key_base, @@ -81,16 +136,16 @@ class CollectiveAllReduceStrategyTestBase( CollectiveAllReduceStrategyTestBase.collective_key_base, instance_key_with_id_start=num_gpus * 10000 + CollectiveAllReduceStrategyTestBase.collective_key_base) - distribution.extended._collective_keys = collective_keys - distribution.extended._cross_device_ops._collective_keys = ( - collective_keys) - if task_type and task_id is not None: - return distribution, 'grpc://' + self._cluster_spec[task_type][ - task_id], session_config - else: - return distribution, '', session_config + strategy.extended._collective_keys = collective_keys + strategy.extended._cross_device_ops._collective_keys = (collective_keys) - def _test_minimize_loss_graph(self, task_type, task_id, num_gpus): + return strategy, target, session_config + + def _test_minimize_loss_graph(self, + task_type, + task_id, + num_gpus, + use_core_strategy=False): d, master_target, config = self._get_test_object(task_type, task_id, num_gpus) with ops.Graph().as_default(), \ @@ -123,7 +178,7 @@ class CollectiveAllReduceStrategyTestBase( def step(): """Perform one optimization step.""" # Run forward & backward to get gradients, variables list. - g_v = d.call_for_each_replica(grad_fn, args=[one]) + g_v = d.extended.call_for_each_replica(grad_fn, args=[one]) # Update the variables using the gradients and the update() function. before_list = [] after_list = [] @@ -135,7 +190,7 @@ class CollectiveAllReduceStrategyTestBase( g = d.extended.reduce_to( reduce_util.ReduceOp.SUM, g, destinations=v) with ops.control_dependencies( - d.update(v, update, g, grouped=False)): + d.extended.update(v, update, args=(g,), group=False)): after_list.append(d.extended.read_var(v)) return before_list, after_list @@ -158,7 +213,11 @@ class CollectiveAllReduceStrategyTestBase( self.assertLess(error_after, error_before) return error_after < error_before - def _test_complex_model(self, task_type, task_id, num_gpus): + def _test_complex_model(self, + task_type, + task_id, + num_gpus, + use_core_strategy=False): d, master_target, config = self._get_test_object(task_type, task_id, num_gpus) @@ -192,6 +251,7 @@ class CollectiveAllReduceStrategyTestBase( image = random_ops.random_uniform([2, 28, 28]) label = random_ops.random_uniform([2, 1], maxval=10, dtype=dtypes.int32) logits = model(image, training=True) + # TODO(yuefengz): make loss a callable for eager mode. loss = losses.sparse_softmax_cross_entropy(labels=label, logits=logits) optimizer = adam.AdamOptimizer(learning_rate=1e-4) train_op = optimizer.minimize(loss, @@ -202,14 +262,18 @@ class CollectiveAllReduceStrategyTestBase( self.cached_session(config=config, target=master_target) as sess: with d.scope(): - train_op = d.call_for_each_replica(model_fn) + train_op = d.extended.call_for_each_replica(model_fn) train_op = d.group(d.unwrap(train_op)) sess.run(variables.global_variables_initializer()) sess.run(train_op) return True - def _test_variable_initialization(self, task_type, task_id, num_gpus): + def _test_variable_initialization(self, + task_type, + task_id, + num_gpus, + use_core_strategy=False): distribution, master_target, config = self._get_test_object( task_type, task_id, num_gpus) with ops.Graph().as_default(), \ @@ -225,7 +289,7 @@ class CollectiveAllReduceStrategyTestBase( 1.0, 10.0, dtype=dtypes.float32)) return array_ops.identity(x) - x = distribution.call_for_each_replica(model_fn) + x = distribution.extended.call_for_each_replica(model_fn) reduced_x = distribution.reduce(reduce_util.ReduceOp.MEAN, x) x = distribution.unwrap(x)[0] @@ -238,8 +302,14 @@ class CollectiveAllReduceStrategyTestBase( reduced_x_value))) return np.allclose(x_value, reduced_x_value, atol=1e-5) - def _test_input_fn_iterator(self, task_type, task_id, num_gpus, input_fn, - expected_values): + def _test_input_fn_iterator(self, + task_type, + task_id, + num_gpus, + input_fn, + expected_values, + test_reinitialize=True, + use_core_strategy=False): distribution, master_target, config = self._get_test_object( task_type, task_id, num_gpus) devices = distribution.extended.worker_devices @@ -262,13 +332,14 @@ class CollectiveAllReduceStrategyTestBase( for r in range(len(devices))]) # After re-initializing the iterator, should be able to iterate again. - sess.run(iterator.initialize()) + if test_reinitialize: + sess.run(iterator.initialize()) - for expected_value in expected_values: - next_element = iterator.get_next() - computed_value = sess.run([values.select_replica(r, next_element) - for r in range(len(devices))]) - self.assertEqual(expected_value, computed_value) + for expected_value in expected_values: + next_element = iterator.get_next() + computed_value = sess.run([values.select_replica(r, next_element) + for r in range(len(devices))]) + self.assertEqual(expected_value, computed_value) class DistributedCollectiveAllReduceStrategyTest( @@ -282,71 +353,114 @@ class DistributedCollectiveAllReduceStrategyTest( cls._cluster_spec = multi_worker_test_base.create_in_process_cluster( num_workers=3, num_ps=0) - def test_num_replicas_in_sync(self): - distribution = collective_all_reduce_strategy.CollectiveAllReduceStrategy( - num_gpus_per_worker=2) - distribution.configure(cluster_spec=self._cluster_spec, task_type='worker', - task_id=0) + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def test_num_replicas_in_sync(self, use_core_strategy): + distribution, _, _ = create_test_objects( + cluster_spec=self._cluster_spec, + task_type='worker', + task_id=0, + num_gpus=2, + use_core_strategy=use_core_strategy) num_workers = len(self._cluster_spec.get('chief', []) + self._cluster_spec.get('worker', [])) self.assertEqual(2 * num_workers, distribution.num_replicas_in_sync) @combinations.generate( - combinations.combine(mode=['graph'], num_gpus=[0, 1, 2], required_gpus=1)) - def testMinimizeLossGraph(self, num_gpus): - self._run_between_graph_clients(self._test_minimize_loss_graph, - self._cluster_spec, num_gpus) + combinations.combine( + mode=['graph'], + num_gpus=[0, 1, 2], + required_gpus=1, + use_core_strategy=[True, False])) + def testMinimizeLossGraph(self, num_gpus, use_core_strategy): + self._run_between_graph_clients( + self._test_minimize_loss_graph, + self._cluster_spec, + num_gpus, + use_core_strategy=use_core_strategy) @combinations.generate( - combinations.combine(mode=['graph'], num_gpus=[0, 1, 2], required_gpus=1)) - def testVariableInitialization(self, num_gpus): + combinations.combine( + mode=['graph'], + num_gpus=[0, 1, 2], + required_gpus=1, + use_core_strategy=[True, False])) + def testVariableInitialization(self, num_gpus, use_core_strategy): if context.num_gpus() < num_gpus: self.skipTest('Not enough GPUs') self._run_between_graph_clients( self._test_variable_initialization, self._cluster_spec, - num_gpus=num_gpus) + num_gpus=num_gpus, + use_core_strategy=use_core_strategy) @combinations.generate( - combinations.combine(mode=['graph'], num_gpus=[0, 1, 2], required_gpus=1)) - def testComplexModel(self, num_gpus): + combinations.combine( + mode=['graph'], + num_gpus=[0, 1, 2], + required_gpus=1, + use_core_strategy=[True, False])) + def testComplexModel(self, num_gpus, use_core_strategy): if context.num_gpus() < num_gpus: self.skipTest('Not enough GPUs') self._run_between_graph_clients( - self._test_complex_model, self._cluster_spec, num_gpus=num_gpus) + self._test_complex_model, + self._cluster_spec, + num_gpus=num_gpus, + use_core_strategy=use_core_strategy) # TODO(yuefengz): Update how we use num_gpus and required_gpus @combinations.generate( - combinations.combine(mode=['graph'], num_gpus=[0, 1, 2], required_gpus=1)) - def testMakeInputFnIterator(self, num_gpus): + combinations.combine( + mode=['graph'], + num_gpus=[0, 1, 2], + required_gpus=1, + use_dataset=[True, False], + use_core_strategy=[True, False])) + def testMakeInputFnIterator(self, num_gpus, use_dataset, use_core_strategy): if context.num_gpus() < num_gpus: self.skipTest('Not enough GPUs') - dataset_fn = lambda: dataset_ops.Dataset.range(100) + if use_dataset: + fn = lambda: dataset_ops.Dataset.range(100) + else: + def fn(): + dataset = dataset_ops.Dataset.range(100) + it = dataset.make_one_shot_iterator() + return it.get_next # We use CPU as the device when num_gpus = 0 devices_per_worker = max(1, num_gpus) expected_values = [[i+j for j in range(devices_per_worker)] for i in range(0, 100, devices_per_worker)] input_fn = self._input_fn_to_test_input_context( - dataset_fn, + fn, expected_num_replicas_in_sync=3*devices_per_worker, expected_num_input_pipelines=3, expected_input_pipeline_id=1) # because task_id = 1 - self._test_input_fn_iterator('worker', 1, num_gpus, - input_fn, expected_values) + self._test_input_fn_iterator( + 'worker', + 1, + num_gpus, + input_fn, + expected_values, + test_reinitialize=use_dataset, + use_core_strategy=use_core_strategy) - def testUpdateConfigProto(self): - distribution = collective_all_reduce_strategy.CollectiveAllReduceStrategy( - num_gpus_per_worker=2) - distribution.configure( - cluster_spec=self._cluster_spec, task_type='worker', task_id=1) + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testUpdateConfigProto(self, use_core_strategy): + strategy, _, _ = self._get_test_object( + task_type='worker', + task_id=1, + num_gpus=2, + use_core_strategy=use_core_strategy) config_proto = config_pb2.ConfigProto(device_filters=['to_be_overridden']) rewrite_options = config_proto.graph_options.rewrite_options rewrite_options.scoped_allocator_opts.enable_op.append('to_be_removed') - new_config = distribution.update_config_proto(config_proto) + new_config = strategy.update_config_proto(config_proto) # Verify group leader self.assertEqual('/job:worker/replica:0/task:0', @@ -397,36 +511,136 @@ class DistributedCollectiveAllReduceStrategyTestWithChief( self._test_complex_model, self._cluster_spec, num_gpus=num_gpus) -class LocalCollectiveAllReduceStrategy(CollectiveAllReduceStrategyTestBase, - strategy_test_lib.DistributionTestBase, - parameterized.TestCase): +class LocalCollectiveAllReduceStrategy( + CollectiveAllReduceStrategyTestBase, + strategy_test_lib.DistributionTestBase, + strategy_test_lib.TwoDeviceDistributionTestBase, + parameterized.TestCase): - def testMinimizeLossGraph(self, num_gpus=2): + @combinations.generate( + combinations.combine( + mode=['graph', 'eager'], + num_gpus=[2, 4], + required_gpus=2, + use_core_strategy=[True, False])) + def testMinimizeLoss(self, num_gpus, use_core_strategy): # Collective ops doesn't support strategy with one device. if context.num_gpus() < num_gpus: self.skipTest('Not enough GPUs') - self._test_minimize_loss_graph(None, None, num_gpus) + if context.executing_eagerly(): + strategy, _, _ = self._get_test_object( + None, None, num_gpus, use_core_strategy=use_core_strategy) + self._test_minimize_loss_eager(strategy) + else: + self._test_minimize_loss_graph( + None, None, num_gpus, use_core_strategy=use_core_strategy) - def testComplexModel(self, num_gpus=2): - # Collective ops doesn't support strategy with one device. + @combinations.generate( + combinations.combine( + mode=['graph'], + num_gpus=[2, 4], + required_gpus=2, + use_core_strategy=[True, False])) + def testComplexModel(self, num_gpus, use_core_strategy): if context.num_gpus() < num_gpus: self.skipTest('Not enough GPUs') - self._test_complex_model(None, None, num_gpus) + self._test_complex_model( + None, None, num_gpus, use_core_strategy=use_core_strategy) - def testMakeInputFnIterator(self, num_gpus=2): - # Collective ops doesn't support strategy with one device. - if context.num_gpus() < num_gpus: - self.skipTest('Not enough GPUs') - dataset_fn = lambda: dataset_ops.Dataset.range(10) - expected_values = [[i, i+1] for i in range(0, 10, 2)] + @combinations.generate( + combinations.combine( + mode=['graph', 'eager'], + required_gpus=2, + use_dataset=[True, False], + use_core_strategy=[True, False])) + def testMakeInputFnIterator(self, use_dataset, use_core_strategy): + num_gpus = 2 + if use_dataset: + fn = lambda: dataset_ops.Dataset.range(5 * num_gpus) + else: + def fn(): + dataset = dataset_ops.Dataset.range(5 * num_gpus) + it = dataset.make_one_shot_iterator() + return it.get_next + expected_values = [range(i, i + num_gpus) for i in range(0, 10, num_gpus)] input_fn = self._input_fn_to_test_input_context( - dataset_fn, + fn, expected_num_replicas_in_sync=num_gpus, expected_num_input_pipelines=1, expected_input_pipeline_id=0) - self._test_input_fn_iterator(None, None, num_gpus, - input_fn, expected_values) + self._test_input_fn_iterator( + None, + None, + num_gpus, + input_fn, + expected_values, + test_reinitialize=use_dataset, + use_core_strategy=use_core_strategy) + + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testAllReduceSum(self, use_core_strategy): + if context.num_gpus() < 2: self.skipTest('Not enough GPUs') + distribution, target, config = self._get_test_object( + None, None, num_gpus=2, use_core_strategy=use_core_strategy) + with self.cached_session(config=config, target=target): + self._test_all_reduce_sum(distribution) + + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testAllReduceSumGradients(self, use_core_strategy): + if context.num_gpus() < 2: self.skipTest('Not enough GPUs') + distribution, target, config = self._get_test_object( + None, None, num_gpus=2, use_core_strategy=use_core_strategy) + with self.cached_session(config=config, target=target): + self._test_all_reduce_sum_gradients(distribution) + + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testAllReduceSumGradientTape(self, use_core_strategy): + if context.num_gpus() < 2: self.skipTest('Not enough GPUs') + distribution, target, config = self._get_test_object( + None, None, num_gpus=2, use_core_strategy=use_core_strategy) + with self.cached_session(config=config, target=target): + self._test_all_reduce_sum_gradient_tape(distribution) + + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testAllReduceMean(self, use_core_strategy): + if context.num_gpus() < 2: self.skipTest('Not enough GPUs') + distribution, target, config = self._get_test_object( + None, None, num_gpus=2, use_core_strategy=use_core_strategy) + with self.cached_session(config=config, target=target): + self._test_all_reduce_mean(distribution) + + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testAllReduceMeanGradients(self, use_core_strategy): + if context.num_gpus() < 2: self.skipTest('Not enough GPUs') + distribution, target, config = self._get_test_object( + None, None, num_gpus=2, use_core_strategy=use_core_strategy) + with self.cached_session(config=config, target=target): + self._test_all_reduce_mean_gradients(distribution) + + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testAllReduceMeanGradientTape(self, use_core_strategy): + if context.num_gpus() < 2: self.skipTest('Not enough GPUs') + distribution, target, config = self._get_test_object( + None, None, num_gpus=2, use_core_strategy=use_core_strategy) + with self.cached_session(config=config, target=target): + self._test_all_reduce_mean_gradient_tape(distribution) + + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testNumpyIterator(self, use_core_strategy): + num_gpus = 2 + if context.num_gpus() < num_gpus: + self.skipTest('Not enough GPUs') + strategy, _, _ = self._get_test_object( + None, None, num_gpus=num_gpus, use_core_strategy=use_core_strategy) + self._test_numpy_iterator(strategy) if __name__ == '__main__': diff --git a/tensorflow/contrib/distribute/python/combinations.py b/tensorflow/contrib/distribute/python/combinations.py index f6c4291659b493b59b102b7ca83454f804898772..7c0f8033fbc046580bc46f90ee9945ffa2a718f9 100644 --- a/tensorflow/contrib/distribute/python/combinations.py +++ b/tensorflow/contrib/distribute/python/combinations.py @@ -49,13 +49,19 @@ import six from tensorflow.contrib import cluster_resolver from tensorflow.contrib.distribute.python import mirrored_strategy as mirrored_lib from tensorflow.contrib.distribute.python import one_device_strategy as one_device_lib +from tensorflow.contrib.distribute.python import parameter_server_strategy from tensorflow.contrib.distribute.python import tpu_strategy as tpu_lib from tensorflow.contrib.optimizer_v2 import adagrad as adagrad_v2 from tensorflow.contrib.optimizer_v2 import adam as adam_v2 from tensorflow.contrib.optimizer_v2 import gradient_descent as gradient_descent_v2 +from tensorflow.contrib.tpu.python.tpu import device_assignment as device_assignment_lib from tensorflow.python.distribute import distribution_strategy_context from tensorflow.python.eager import context from tensorflow.python.framework import ops +from tensorflow.python.keras.optimizer_v2 import adagrad as adagrad_keras_v2 +from tensorflow.python.keras.optimizer_v2 import adam as adam_keras_v2 +from tensorflow.python.keras.optimizer_v2 import gradient_descent as gradient_descent_keras_v2 +from tensorflow.python.keras.optimizer_v2 import rmsprop as rmsprop_keras_v2 from tensorflow.python.training import adagrad from tensorflow.python.training import adam from tensorflow.python.training import gradient_descent @@ -226,7 +232,7 @@ def combine(**kwargs): if not kwargs: return [OrderedDict()] - sort_by_key = lambda k: k[0][0] + sort_by_key = lambda k: k[0] kwargs = OrderedDict(sorted(kwargs.items(), key=sort_by_key)) first = list(kwargs.items())[0] @@ -321,11 +327,19 @@ class NamedDistribution(object): return self._required_tpu -def _get_tpu_strategy_creator(steps_per_run): +def _get_tpu_strategy_creator(steps_per_run, use_single_core=False, **kwargs): def _create_tpu_strategy(): resolver = cluster_resolver.TPUClusterResolver("") - tpu_lib.initialize_tpu_system(resolver) - strategy = tpu_lib.TPUStrategy(resolver, steps_per_run=steps_per_run) + topology = tpu_lib.initialize_tpu_system(resolver) + device_assignment = None + if use_single_core: + device_assignment = device_assignment_lib.DeviceAssignment( + topology, core_assignment=device_assignment_lib. + SINGLE_CORE_ASSIGNMENT) + + strategy = tpu_lib.TPUStrategy(resolver, steps_per_run=steps_per_run, + device_assignment=device_assignment, + **kwargs) return strategy return _create_tpu_strategy @@ -338,12 +352,23 @@ default_strategy = NamedDistribution( one_device_strategy = NamedDistribution( "OneDeviceCPU", lambda: one_device_lib.OneDeviceStrategy("/cpu:0"), required_gpus=None) +one_device_strategy_gpu = NamedDistribution( + "OneDeviceGPU", lambda: one_device_lib.OneDeviceStrategy("/gpu:0"), + required_gpus=1) tpu_strategy = NamedDistribution( "TPU", _get_tpu_strategy_creator(steps_per_run=2), required_tpu=True) tpu_strategy_one_step = NamedDistribution( "TPUOneStep", _get_tpu_strategy_creator(steps_per_run=1), required_tpu=True) +tpu_strategy_one_core = NamedDistribution( + "TPUOneCore", _get_tpu_strategy_creator( + steps_per_run=2, use_single_core=True), + required_tpu=True) +tpu_strategy_one_step_one_core = NamedDistribution( + "TPUOneStepOneCore", _get_tpu_strategy_creator( + steps_per_run=1, use_single_core=True), + required_tpu=True) mirrored_strategy_with_one_cpu = NamedDistribution( "Mirrored1CPU", @@ -375,6 +400,11 @@ core_mirrored_strategy_with_two_gpus = NamedDistribution( "CoreMirrored2GPUs", lambda: mirrored_lib.CoreMirroredStrategy(["/gpu:0", "/gpu:1"]), required_gpus=2) +parameter_server_strategy_with_two_gpus = NamedDistribution( + "ParameterServer2GPUs", + lambda: parameter_server_strategy.ParameterServerStrategy( + num_gpus_per_worker=2), + required_gpus=2) gradient_descent_optimizer_v1_fn = NamedObject( @@ -394,10 +424,20 @@ gradient_descent_optimizer_v2_fn = NamedObject( adagrad_optimizer_v2_fn = NamedObject( "AdagradV2", lambda: adagrad_v2.AdagradOptimizer(0.001)) adam_optimizer_v2_fn = NamedObject( - "AdamV2", lambda: adam_v2.AdamOptimizer(0.001, epsilon=1)) + "AdamV2", lambda: adam_v2.AdamOptimizer(0.001, epsilon=1.0)) optimizers_v2 = [gradient_descent_optimizer_v2_fn, adagrad_optimizer_v2_fn] +gradient_descent_optimizer_keras_v2_fn = NamedObject( + "GradientDescentKerasV2", + lambda: gradient_descent_keras_v2.SGD(0.2)) +adagrad_optimizer_keras_v2_fn = NamedObject( + "AdagradKerasV2", lambda: adagrad_keras_v2.Adagrad(0.001)) +adam_optimizer_keras_v2_fn = NamedObject( + "AdamKerasV2", lambda: adam_keras_v2.Adam(0.001, epsilon=1.0)) +rmsprop_optimizer_keras_v2_fn = NamedObject( + "RmsPropKerasV2", lambda: rmsprop_keras_v2.RMSprop(0.001)) + graph_and_eager_modes = ["graph", "eager"] diff --git a/tensorflow/contrib/distribute/python/combinations_test.py b/tensorflow/contrib/distribute/python/combinations_test.py index 86aa48cea889c6c2ce169b18bcabb6d08890fbed..9f3deadbec98c4f66061ca29b4d29a74b8de40b1 100644 --- a/tensorflow/contrib/distribute/python/combinations_test.py +++ b/tensorflow/contrib/distribute/python/combinations_test.py @@ -42,6 +42,14 @@ class TestingCombinationsTest(test.TestCase): "b": 3 }], combinations.combine(a=[1, 2], b=[2, 3])) + def test_arguments_sorted(self): + self.assertEqual([ + OrderedDict([("aa", 1), ("ab", 2)]), + OrderedDict([("aa", 1), ("ab", 3)]), + OrderedDict([("aa", 2), ("ab", 2)]), + OrderedDict([("aa", 2), ("ab", 3)]) + ], combinations.combine(ab=[2, 3], aa=[1, 2])) + def test_combine_single_parameter(self): self.assertEqual([{ "a": 1, diff --git a/tensorflow/contrib/distribute/python/cross_device_ops_test.py b/tensorflow/contrib/distribute/python/cross_device_ops_test.py index 54cce2988383fcf5e063726948fbbf62c7094ce5..2d56c25f1034acb29e460985a5eee0133397b97f 100644 --- a/tensorflow/contrib/distribute/python/cross_device_ops_test.py +++ b/tensorflow/contrib/distribute/python/cross_device_ops_test.py @@ -204,15 +204,15 @@ class SingleWorkerCrossDeviceOpsTest(CrossDeviceOpsTestBase): reduction_to_one_combinations = combinations.combine( cross_device_ops=[ combinations.NamedObject( - "DefaultReductionToOneDeviceCrossDeviceOps", - cross_device_ops_lib.ReductionToOneDeviceCrossDeviceOps()), + "DefaultReductionToOneDevice", + cross_device_ops_lib.ReductionToOneDevice()), combinations.NamedObject( "ReductionToCPUDeviceCrossDeviceOps", - cross_device_ops_lib.ReductionToOneDeviceCrossDeviceOps( + cross_device_ops_lib.ReductionToOneDevice( reduce_to_device=_cpu_device)), combinations.NamedObject( "AccumulateNCrossDeviceOp", - cross_device_ops_lib.ReductionToOneDeviceCrossDeviceOps( + cross_device_ops_lib.ReductionToOneDevice( accumulation_fn=math_ops.accumulate_n)), ], distribution=[ @@ -228,20 +228,23 @@ class SingleWorkerCrossDeviceOpsTest(CrossDeviceOpsTestBase): combinations.NamedObject( "AllReduce", cross_device_ops_lib.AllReduceCrossDeviceOps("nccl", 1, 0, 0)), - combinations.NamedObject( - "HierarchicalCopy", - cross_device_ops_lib.AllReduceCrossDeviceOps( - "hierarchical_copy", 8, 0, 0)), combinations.NamedObject( "AllReduceNoGradientRepacking", cross_device_ops_lib.AllReduceCrossDeviceOps("nccl", 0, 0, 0)), + combinations.NamedObject("NcclAllReduce", + cross_device_ops_lib.NcclAllReduce()), + combinations.NamedObject( + "HierarchicalCopy", + cross_device_ops_lib.HierarchicalCopyAllReduce(8)), combinations.NamedObject( "HierarchicalCopyAggregateSmallTensors", cross_device_ops_lib.AllReduceCrossDeviceOps( "hierarchical_copy", 0, 100, 10)) ], - distribution=[combinations.mirrored_strategy_with_two_gpus, - combinations.core_mirrored_strategy_with_two_gpus], + distribution=[ + combinations.mirrored_strategy_with_two_gpus, + combinations.core_mirrored_strategy_with_two_gpus + ], mode=["graph", "eager"]) @combinations.generate(reduction_to_one_combinations + allreduce_combinations) @@ -306,8 +309,8 @@ class SingleWorkerCrossDeviceOpsTest(CrossDeviceOpsTestBase): combinations.combine( cross_device_ops_instance=[ combinations.NamedObject( - "ReductionToOneDeviceCrossDeviceOps", - cross_device_ops_lib.ReductionToOneDeviceCrossDeviceOps()), + "ReductionToOneDevice", + cross_device_ops_lib.ReductionToOneDevice()), combinations.NamedObject( "AllReduceCrossDeviceOps", cross_device_ops_lib.AllReduceCrossDeviceOps()) diff --git a/tensorflow/contrib/distribute/python/examples/BUILD b/tensorflow/contrib/distribute/python/examples/BUILD index 84b106545e1326fddd3ed299462534af982dc102..5f89df5824a8d03198987a6fa3d21e2330deedf0 100644 --- a/tensorflow/contrib/distribute/python/examples/BUILD +++ b/tensorflow/contrib/distribute/python/examples/BUILD @@ -31,6 +31,12 @@ py_binary( py_binary( name = "keras_mnist", + srcs = ["keras_mnist.py"], + deps = [":keras_mnist_lib"], +) + +py_library( + name = "keras_mnist_lib", srcs = [ "keras_mnist.py", ], @@ -39,3 +45,14 @@ py_binary( "//third_party/py/numpy", ], ) + +py_binary( + name = "mnist_eager_multigpu", + srcs = [ + "mnist_eager_multigpu.py", + ], + deps = [ + "//tensorflow:tensorflow_py", + "//third_party/py/numpy", + ], +) diff --git a/tensorflow/contrib/distribute/python/examples/keras_mnist.py b/tensorflow/contrib/distribute/python/examples/keras_mnist.py index 60fda996642464135fe1fb8c314bcf7f04d19362..1ce91ecaf22a80a53124c8f00fac05c6b4711ed9 100644 --- a/tensorflow/contrib/distribute/python/examples/keras_mnist.py +++ b/tensorflow/contrib/distribute/python/examples/keras_mnist.py @@ -109,22 +109,21 @@ def main(_): tf.enable_eager_execution() train_ds, eval_ds, input_shape = get_input_datasets() - model = get_model(input_shape) # Instantiate the MirroredStrategy object. If we don't specify `num_gpus` or # the `devices` argument then all the GPUs available on the machine are used. # TODO(priyag): Use `tf.distribute.MirroredStrategy` once available. strategy = mirrored_strategy.MirroredStrategy(['/gpu:0', '/cpu:0']) - optimizer = rmsprop.RMSProp(learning_rate=0.001) - - # Compile the model by passing the distribution strategy object to the - # `distribute` argument. `fit`, `evaluate` and `predict` will be distributed - # based on the strategy instantiated. - model.compile(loss=tf.keras.losses.categorical_crossentropy, - optimizer=optimizer, - metrics=['accuracy'], - distribute=strategy) + # Create and compile the model under Distribution strategy scope. + # `fit`, `evaluate` and `predict` will be distributed based on the strategy + # model was compiled with. + with strategy.scope(): + model = get_model(input_shape) + optimizer = rmsprop.RMSProp(learning_rate=0.001) + model.compile(loss=tf.keras.losses.categorical_crossentropy, + optimizer=optimizer, + metrics=['accuracy']) # Train the model with the train dataset. model.fit(x=train_ds, epochs=20, steps_per_epoch=468) diff --git a/tensorflow/contrib/distribute/python/examples/mnist_eager_multigpu.py b/tensorflow/contrib/distribute/python/examples/mnist_eager_multigpu.py new file mode 100644 index 0000000000000000000000000000000000000000..c045a5586b9dad371d8c505f9cac4b792dd157fd --- /dev/null +++ b/tensorflow/contrib/distribute/python/examples/mnist_eager_multigpu.py @@ -0,0 +1,169 @@ +# 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.compat.v2 as tf + +flags.DEFINE_integer("num_gpus", None, "How many GPUs should we run on?" + "Defaults to all available GPUs, otherwise CPU.") +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.DEFINE_boolean("use_function", False, + "Should we wrap the step in a tf.function.") + +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(unused_argv): + """Run a CNN model on MNIST data to demonstrate DistributedStrategies.""" + + tf.enable_v2_behavior() + + num_gpus = FLAGS.num_gpus + if num_gpus is None: + devices = None + elif num_gpus == 0: + devices = ["/device:CPU:0"] + else: + devices = ["/device:GPU:{}".format(i) for i in range(num_gpus)] + strategy = tf.distribute.MirroredStrategy(devices) + + 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): + images, labels = inputs + with tf.GradientTape() as tape: + logits = model(images, training=True) + loss = compute_loss(logits, labels) + grads = tape.gradient(loss, model.variables) + optimizer.apply_gradients(zip(grads, model.variables)) + training_loss.update_state(loss) + training_accuracy.update_state(labels, logits) + + def test_step(inputs): + images, labels = inputs + logits = model(images, training=False) + loss = compute_loss(logits, labels) + test_loss.update_state(loss) + test_accuracy.update_state(labels, logits) + + train_iterator = strategy.make_dataset_iterator(train_ds) + test_iterator = strategy.make_dataset_iterator(test_ds) + + for epoch in range(0, FLAGS.num_epochs): + # TODO(b/123315763): Create the tf.function outside this loop once we are + # able to initialize iterator in eager mode. + dist_train = lambda it: strategy.experimental_run(train_step, it) + dist_test = lambda it: strategy.experimental_run(test_step, it) + if FLAGS.use_function: + dist_train = tf.function(dist_train) + dist_test = tf.function(dist_test) + + # Train + print("Starting epoch {}".format(epoch)) + train_iterator.initialize() + while True: + try: + dist_train(train_iterator) + except tf.errors.OutOfRangeError: + break + print("Training loss: {:0.4f}, accuracy: {:0.2f}%".format( + training_loss.result(), training_accuracy.result() * 100)) + training_loss.reset_states() + training_accuracy.reset_states() + + # Test + test_iterator.initialize() + while True: + try: + dist_test(test_iterator) + except tf.errors.OutOfRangeError: + break + print("Test loss: {:0.4f}, accuracy: {:0.2f}%".format( + test_loss.result(), test_accuracy.result() * 100)) + test_loss.reset_states() + test_accuracy.reset_states() + + +if __name__ == "__main__": + app.run(main) diff --git a/tensorflow/contrib/distribute/python/input_lib_test.py b/tensorflow/contrib/distribute/python/input_lib_test.py index f589cd6ad54ea8f33002cb067ef8d83d3d33036a..10a58316ec5b3d9d968a88c5c39ff70c277daa65 100644 --- a/tensorflow/contrib/distribute/python/input_lib_test.py +++ b/tensorflow/contrib/distribute/python/input_lib_test.py @@ -22,7 +22,6 @@ from absl.testing import parameterized from tensorflow.contrib.distribute.python import combinations from tensorflow.contrib.distribute.python import multi_worker_test_base -from tensorflow.core.protobuf import config_pb2 from tensorflow.python.data.experimental.ops import batching from tensorflow.python.data.ops import dataset_ops from tensorflow.python.distribute import distribute_lib @@ -31,243 +30,10 @@ from tensorflow.python.distribute import values from tensorflow.python.eager import context from tensorflow.python.eager import test from tensorflow.python.framework import errors -from tensorflow.python.framework import test_util from tensorflow.python.ops import control_flow_ops -from tensorflow.python.ops import random_ops from tensorflow.python.util import nest -class PerReplicaDatasetTest(test.TestCase): - - config = config_pb2.ConfigProto() - config.allow_soft_placement = True - - def _test_iterator(self, devices, dataset, expected_values): - device_map = values.ReplicaDeviceMap(devices) - input_workers = input_lib.InputWorkers(device_map) - per_replica_dataset = input_lib.PerReplicaDataset(dataset, input_workers, 0) - if context.executing_eagerly(): - iterator = per_replica_dataset.make_one_shot_iterator() - else: - iterator = per_replica_dataset.make_initializable_iterator() - self.evaluate([iterator.initializer]) - - for expected_value in expected_values: - next_element = iterator.get_next_as_list() - computed_value = self.evaluate(next_element) - self.assertEqual(expected_value, computed_value) - - with self.assertRaises(errors.OutOfRangeError): - next_element = iterator.get_next_as_list() - self.evaluate(next_element) - - @test_util.run_in_graph_and_eager_modes - def testOneDevice(self): - devices = ["/device:CPU:0"] - dataset = dataset_ops.Dataset.range(10) - - expected_values = [[i] for i in range(10)] - - self._test_iterator(devices, dataset, expected_values) - - @test_util.run_in_graph_and_eager_modes(config=config) - def testMultipleDevices(self): - if context.num_gpus() < 1 and context.executing_eagerly(): - self.skipTest("A GPU is not available for this test in eager mode.") - - devices = ["/device:CPU:0", "/device:GPU:0"] - dataset = dataset_ops.Dataset.range(10) - - expected_values = [[i, i+1] for i in range(0, 10, 2)] - - self._test_iterator(devices, dataset, expected_values) - - @test_util.run_in_graph_and_eager_modes(config=config) - def testTupleDataset(self): - if context.num_gpus() < 1 and context.executing_eagerly(): - self.skipTest("A GPU is not available for this test in eager mode.") - - devices = ["/device:CPU:0", "/device:GPU:0"] - dataset1 = dataset_ops.Dataset.range(10) - dataset2 = dataset_ops.Dataset.range(10).map(lambda x: x**2) - dataset = dataset_ops.Dataset.zip((dataset1, dataset2)) - - expected_values = [[(i, i**2), (i+1, (i+1)**2)] for i in range(0, 10, 2)] - - self._test_iterator(devices, dataset, expected_values) - - @test_util.run_in_graph_and_eager_modes(config=config) - def testUnevenDatasetBatches(self): - if context.num_gpus() < 1 and context.executing_eagerly(): - self.skipTest("A GPU is not available for this test in eager mode.") - - devices = ["/device:CPU:0", "/device:GPU:0"] - dataset = dataset_ops.Dataset.range(11) - - expected_values = [[i, i+1] for i in range(0, 10, 2)] - self._test_iterator(devices, dataset, expected_values) - - def testInitializableIterator(self): - with context.graph_mode(): - devices = ["/device:CPU:0"] - # Using random input since that is only allowed with initializable - # iterator. - dataset = dataset_ops.Dataset.from_tensor_slices( - random_ops.random_uniform((10,))) - - device_map = values.ReplicaDeviceMap(devices) - input_workers = input_lib.InputWorkers(device_map) - per_replica_dataset = input_lib.PerReplicaDataset( - dataset, input_workers, 0) - iterator = per_replica_dataset.make_initializable_iterator() - - self.evaluate(iterator.initializer) - next_element = iterator.get_next_as_list() - for _ in range(10): - self.evaluate(next_element) - - # Should fail after the input is finished. - with self.assertRaises(errors.OutOfRangeError): - self.evaluate(next_element) - - # After re-initializing the iterator, should be able to iterate again. - self.evaluate(iterator.initializer) - for _ in range(10): - self.evaluate(next_element) - - -class MultiWorkerDatasetTest(multi_worker_test_base.MultiWorkerTestBase): - - def _test_iterator(self, sess, iterator, devices, expected_values): - next_element = iterator.get_next() - for r, device in enumerate(devices): - v = values.select_replica(r, next_element) - # The `v` here can be a tuple. - for element in nest.flatten(v): - self.assertTrue(element.device in device) - - for expected_value in expected_values: - t = [values.select_replica(r, next_element) for r in range(len(devices))] - actual = sess.run(t) - self.assertEqual(expected_value, actual) - - with self.assertRaises(errors.OutOfRangeError): - sess.run([values.select_replica(r, next_element) - for r in range(len(devices))]) - - def _test_dataset(self, dataset_fn, worker_devices, devices, - expected_values): - device_map = values.ReplicaDeviceMap(devices) - input_workers = input_lib.InputWorkers(device_map, worker_devices) - multi_worker_dataset = input_lib.MultiWorkerDataset( - dataset_fn, input_workers) - multi_worker_iterator = multi_worker_dataset.make_initializable_iterator() - with self.cached_session() as sess: - sess.run(multi_worker_iterator.initializer) - self._test_iterator(sess, multi_worker_iterator, devices, expected_values) - - def _cpu_devices(self): - worker_devices = ( - ("/job:worker/replica:0/task:0", - ["/job:worker/replica:0/task:0/device:CPU:0"]), - ("/job:worker/replica:0/task:1", - ["/job:worker/replica:0/task:1/device:CPU:0"]) - ) - devices = [ - "/job:worker/replica:0/task:0/device:CPU:0", - "/job:worker/replica:0/task:1/device:CPU:0" - ] - return worker_devices, devices - - def _cpu_and_one_gpu_devices(self): - worker_devices = ( - ("/job:worker/replica:0/task:0", ( - "/job:worker/replica:0/task:0/device:GPU:0", - "/job:worker/replica:0/task:0/device:CPU:0" - )), - ("/job:worker/replica:0/task:1", ( - "/job:worker/replica:0/task:1/device:GPU:0", - "/job:worker/replica:0/task:1/device:CPU:0" - )) - ) - devices = [ - "/job:worker/replica:0/task:0/device:GPU:0", - "/job:worker/replica:0/task:0/device:CPU:0", - "/job:worker/replica:0/task:1/device:GPU:0", - "/job:worker/replica:0/task:1/device:CPU:0" - ] - return worker_devices, devices - - def testDataDistributionOneDevicePerWorker(self): - worker_devices, devices = self._cpu_devices() - with context.graph_mode(): - dataset_fn = lambda: dataset_ops.Dataset.range(8) - self._test_dataset( - dataset_fn, worker_devices, devices, - [[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]]) - - def testDataDistributionTwoDevicePerWorker(self): - if context.num_gpus() < 1: - self.skipTest("A GPU is not available for this test.") - worker_devices, devices = self._cpu_and_one_gpu_devices() - with context.graph_mode(): - dataset_fn = lambda: dataset_ops.Dataset.range(8) - self._test_dataset( - dataset_fn, worker_devices, devices, - [[0, 1, 0, 1], [2, 3, 2, 3], [4, 5, 4, 5], [6, 7, 6, 7]]) - - def testTupleDataset(self): - worker_devices, devices = self._cpu_devices() - - with context.graph_mode(): - - def dataset_fn(): - dataset1 = dataset_ops.Dataset.range(8) - dataset2 = dataset_ops.Dataset.range(8).map(lambda x: x**2) - return dataset_ops.Dataset.zip((dataset1, dataset2)) - - expected_values = [[(i, i**2), (i, i**2)] for i in range(8)] - self._test_dataset(dataset_fn, worker_devices, devices, - expected_values) - - def testInitializableIterator(self): - worker_devices, devices = self._cpu_devices() - with context.graph_mode(), self.cached_session() as sess: - dataset_fn = lambda: dataset_ops.Dataset.range(8) - device_map = values.ReplicaDeviceMap(devices) - input_workers = input_lib.InputWorkers(device_map, worker_devices) - multi_worker_dataset = input_lib.MultiWorkerDataset( - dataset_fn, input_workers) - multi_worker_iterator = multi_worker_dataset.make_initializable_iterator() - - sess.run(multi_worker_iterator.initializer) - self._test_iterator( - sess, multi_worker_iterator, devices, - [[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]]) - - # After re-initializing the iterator, should be able to iterate again. - sess.run(multi_worker_iterator.initializer) - self._test_iterator( - sess, multi_worker_iterator, devices, - [[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]]) - - def testValueErrorForIterator(self): - # Incompatiable arguments. - d1 = "/device:GPU:0" - d2 = "/device:GPU:1" - device_map = values.ReplicaDeviceMap([d1, d2]) - input_workers = input_lib.InputWorkers( - device_map, (("w1", (d1,)), ("w2", (d2,)))) - with self.assertRaises(ValueError): - input_lib.MultiWorkerDataIterator([("w1", None)], input_workers) - - def testDuplicateDevices(self): - _, devices = self._cpu_devices() - devices.append("/job:worker/replica:0/task:0/device:CPU:0") - with self.assertRaises(ValueError): - _ = values.ReplicaDeviceMap(devices) - - class InputIteratorTestBase(test.TestCase): def _test_iterator(self, input_type, dataset_fn, worker_device_pairs, diff --git a/tensorflow/contrib/distribute/python/keras_backward_compat_test.py b/tensorflow/contrib/distribute/python/keras_backward_compat_test.py index 92de8e643e7588365c23dc8513e197c0869c9ecf..c49b5522f9135efd9ae3005e92099caf54a76a3a 100644 --- a/tensorflow/contrib/distribute/python/keras_backward_compat_test.py +++ b/tensorflow/contrib/distribute/python/keras_backward_compat_test.py @@ -31,6 +31,7 @@ from tensorflow.python.framework import random_seed from tensorflow.python.keras import testing_utils 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.keras.utils.mode_keys import ModeKeys from tensorflow.python.ops.parsing_ops import gen_parsing_ops from tensorflow.python.training import gradient_descent from tensorflow.python.training import rmsprop @@ -316,15 +317,19 @@ def all_strategy_combinations(): return strategy_minus_tpu_combinations() + tpu_strategy_combinations() -# TODO(priyag): Add v2 optimizers here. def strategy_and_optimizer_combinations(): return combinations.times( all_strategy_combinations(), - combinations.combine( - optimizer=[combinations.adagrad_optimizer_v1_fn, - combinations.adam_optimizer_v1_fn, - combinations.gradient_descent_optimizer_v1_fn, - combinations.rmsprop_optimizer_v1_fn])) + combinations.combine(optimizer=[ + combinations.adagrad_optimizer_v1_fn, + combinations.adagrad_optimizer_keras_v2_fn, + combinations.adam_optimizer_v1_fn, + combinations.adam_optimizer_keras_v2_fn, + combinations.gradient_descent_optimizer_v1_fn, + combinations.gradient_descent_optimizer_keras_v2_fn, + combinations.rmsprop_optimizer_v1_fn, + combinations.rmsprop_optimizer_keras_v2_fn + ])) def strategy_and_input_combinations(): @@ -741,7 +746,9 @@ class TestDistributionStrategyWithDatasets(test.TestCase, model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0, callbacks=[keras.callbacks.LearningRateScheduler(schedule)]) - grouped_models = distribution.unwrap(model._distributed_model) + grouped_models = distribution.unwrap( + distributed_training_utils.get_distributed_model( + model, ModeKeys.TRAIN)) with distribution.scope(): for m in grouped_models: self.assertAllClose(0.001, keras.backend.get_value( @@ -787,16 +794,21 @@ class TestDistributionStrategyErrorCases(test.TestCase, parameterized.TestCase): verbose=0, sample_weight=sample_weight) - # Test with not specifying the `steps` argument. - with self.assertRaisesRegexp( - ValueError, 'the `steps_per_epoch` argument'): + # Test with not specifying the `steps` argument for dataset with + # infinite cardinality. + dataset = dataset.repeat() + with self.assertRaisesRegexp(ValueError, 'When passing an infinitely ' + 'repeating dataset, you must specify the ' + '`steps_per_epoch` argument'): model.fit(dataset, epochs=1, verbose=0) - with self.assertRaisesRegexp(ValueError, - 'the `steps` argument'): + with self.assertRaisesRegexp(ValueError, 'When passing an infinitely ' + 'repeating dataset, you must specify the ' + '`steps` argument'): model.evaluate(dataset, verbose=0) - with self.assertRaisesRegexp(ValueError, - 'the `steps` argument'): + with self.assertRaisesRegexp(ValueError, 'When passing an infinitely ' + 'repeating dataset, you must specify the ' + '`steps` argument'): model.predict(dataset, verbose=0) @combinations.generate(combinations.combine( diff --git a/tensorflow/contrib/distribute/python/keras_correctness_test.py b/tensorflow/contrib/distribute/python/keras_correctness_test_base.py similarity index 51% rename from tensorflow/contrib/distribute/python/keras_correctness_test.py rename to tensorflow/contrib/distribute/python/keras_correctness_test_base.py index 3abdee2c0ed8ebe60876a6d43bb90d8f4b7a5052..0fb7a18c40484ce01a5acfd6b191de464cfd9840 100644 --- a/tensorflow/contrib/distribute/python/keras_correctness_test.py +++ b/tensorflow/contrib/distribute/python/keras_correctness_test_base.py @@ -19,6 +19,7 @@ from __future__ import print_function from absl.testing import parameterized import numpy as np +import six from tensorflow.contrib.distribute.python import combinations from tensorflow.contrib.distribute.python import mirrored_strategy @@ -30,8 +31,6 @@ from tensorflow.python.eager import context from tensorflow.python.eager import test from tensorflow.python.framework import random_seed 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.training import gradient_descent _RANDOM_SEED = 1337 _EVAL_STEPS = 20 @@ -53,33 +52,52 @@ all_strategies = [ ] -def all_strategy_combinations_with_eager_and_graph_modes(): - return combinations.combine(distribution=all_strategies, - mode=['graph', 'eager']) +def eager_mode_test_configuration(): + return combinations.combine(mode='eager', + use_numpy=False, + use_validation_data=False) -def all_strategy_combinations_with_graph_mode(): - return combinations.combine(distribution=all_strategies, mode=['graph']) +def graph_mode_test_configuration(): + return combinations.combine(mode='graph', + use_numpy=[True, False], + use_validation_data=[True, False]) -def strategy_and_input_combinations(): - def cnn_model_with_batch_norm(**kwargs): - return _create_cnn_model(with_batch_norm=True, **kwargs) - +def all_strategy_and_input_config_combinations(): return ( combinations.times( combinations.combine(distribution=all_strategies), - combinations.combine(mode=['graph', 'eager'], - use_numpy=[True, False], - use_validation_data=[True, False]), - combinations.combine(model_with_data=[ - ModelWithData('dnn', _create_dnn_model, _dnn_training_data), - ModelWithData('cnn', _create_cnn_model, _cnn_training_data), - ModelWithData('cnn_batch_norm', - cnn_model_with_batch_norm, - _cnn_training_data, - with_batch_norm=True), - ]))) + eager_mode_test_configuration() + graph_mode_test_configuration())) + + +def strategies_for_embedding_models(): + """Returns distribution strategies to test for embedding models. + + Since embedding models take longer to train, we disregard OneDeviceStrategy + and DefaultStrategy in order to prevent testing timeouts. + """ + + return [s for s in all_strategies if s.required_tpu or s.required_gpus] + + +def test_combinations_for_embedding_model(): + return ( + combinations.times( + combinations.combine(distribution= + strategies_for_embedding_models()), + (graph_mode_test_configuration() + + eager_mode_test_configuration()))) + + +def test_combinations_with_tpu_strategies(): + tpu_strategies = [combinations.tpu_strategy, + combinations.tpu_strategy_one_step] + + return ( + combinations.times( + combinations.combine(distribution=tpu_strategies), + graph_mode_test_configuration())) class MaybeDistributionScope(object): @@ -100,97 +118,6 @@ class MaybeDistributionScope(object): self._scope = None -class ModelWithData(object): - """An object giving a good name in combinations. - - The model_fn must take two arguments: initial_weights and distribution. - """ - - def __init__(self, name, model_fn, data_fn, with_batch_norm=False): - self.name = name - self.model_fn = model_fn - self.data_fn = data_fn - self.with_batch_norm = with_batch_norm - - def __repr__(self): - return self.name - - -def _dnn_training_data(): - # TODO(xiejw): Change this back to 10000, once we support final partial - # batch. - num_samples = 9984 - x_train = np.random.rand(num_samples, 1) - y_train = 3 * x_train - x_train = x_train.astype('float32') - y_train = y_train.astype('float32') - x_predict = [[1.], [2.], [3.], [4.]] - return x_train, y_train, x_predict - - -def _create_dnn_model(initial_weights=None, distribution=None): - with MaybeDistributionScope(distribution): - # We add few non-linear layers to make it non-trivial. - model = keras.Sequential() - model.add(keras.layers.Dense(10, activation='relu', input_shape=(1,))) - model.add(keras.layers.Dense(10, activation='relu')) - model.add(keras.layers.Dense(10, activation='relu')) - model.add(keras.layers.Dense(1)) - - if initial_weights: - model.set_weights(initial_weights) - - model.compile( - loss=keras.losses.mean_squared_error, - optimizer=gradient_descent_keras.SGD(0.5), - metrics=['mse']) - return model - - -def _cnn_training_data(count=_GLOBAL_BATCH_SIZE * _EVAL_STEPS, - shape=(28, 28, 3), num_classes=10): - centers = np.random.randn(num_classes, *shape) - - features = [] - labels = [] - for _ in range(count): - label = np.random.randint(0, num_classes, size=1)[0] - offset = np.random.normal(loc=0, scale=0.1, size=np.prod(shape)) - offset = offset.reshape(shape) - labels.append(label) - features.append(centers[label] + offset) - - x_train = np.asarray(features, dtype=np.float32) - y_train = np.asarray(labels, dtype=np.float32).reshape((count, 1)) - x_predict = x_train - return x_train, y_train, x_predict - - -def _create_cnn_model(initial_weights=None, distribution=None, - with_batch_norm=False): - with MaybeDistributionScope(distribution): - image = keras.layers.Input(shape=(28, 28, 3), name='image') - c1 = keras.layers.Conv2D( - name='conv1', filters=16, kernel_size=(3, 3), strides=(4, 4))( - image) - if with_batch_norm: - c1 = keras.layers.BatchNormalization(name='bn1')(c1) - c1 = keras.layers.MaxPooling2D(pool_size=(2, 2))(c1) - logits = keras.layers.Dense( - 10, activation='softmax', name='pred')( - keras.layers.Flatten()(c1)) - model = keras.Model(inputs=[image], outputs=[logits]) - - if initial_weights: - model.set_weights(initial_weights) - - model.compile( - optimizer=gradient_descent.GradientDescentOptimizer(learning_rate=0.1), - loss='sparse_categorical_crossentropy', - metrics=['sparse_categorical_accuracy']) - return model - - def batch_wrapper(dataset, batch_size, distribution, repeat=None): if repeat: dataset = dataset.repeat(repeat) @@ -213,6 +140,19 @@ def get_batch_size(global_batch_size, distribution): return batch_size +def get_data_size(data): + """Gets the size of data in list, tuple, dict, or a numpy array.""" + assert isinstance(data, (np.ndarray, list, dict, tuple)) + + if isinstance(data, np.ndarray): + return len(data) + + if isinstance(data, (list, tuple)): + return len(data[0]) + + return len(six.next(six.itervalues(data))) + + def get_correctness_test_inputs(use_numpy, use_validation_data, with_distribution, x_train, y_train, x_predict): """Generates the inputs for correctness check when enable Keras with DS.""" @@ -239,11 +179,12 @@ def get_correctness_test_inputs(use_numpy, use_validation_data, 'y': y_train, } predict_inputs = { - 'x': np.array(x_predict, dtype=np.float32), + 'x': x_predict } else: - if len(x_train) < _GLOBAL_BATCH_SIZE * _EVAL_STEPS: - # Currently, we cannot detech the size of a dataset. So, the eval steps is + training_data_size = get_data_size(x_train) + if training_data_size < _GLOBAL_BATCH_SIZE * _EVAL_STEPS: + # Currently, we cannot detect the size of a dataset. So, the eval steps is # hard coded. raise ValueError('x_train must have at least ' '_GLOBAL_BATCH_SIZE * _EVAL_STEPS samples') @@ -259,7 +200,7 @@ def get_correctness_test_inputs(use_numpy, use_validation_data, 'y': None, 'epochs': training_epochs, 'shuffle': False, - 'steps_per_epoch': len(x_train) // global_batch_size, + 'steps_per_epoch': training_data_size // global_batch_size, } if use_validation_data: eval_inputs = None # Remove the eval_inputs @@ -275,7 +216,8 @@ def get_correctness_test_inputs(use_numpy, use_validation_data, 'steps': _EVAL_STEPS, } - predict_batch_size = get_batch_size(len(x_predict), with_distribution) + predict_batch_size = get_batch_size(get_data_size(x_predict), + with_distribution) predict_dataset = dataset_ops.Dataset.from_tensor_slices(x_predict) predict_dataset = batch_wrapper(predict_dataset, predict_batch_size, with_distribution) @@ -287,8 +229,9 @@ def get_correctness_test_inputs(use_numpy, use_validation_data, return training_inputs, eval_inputs, predict_inputs -def fit_eval_and_predict( - initial_weights, input_fn, model_fn, distribution=None): +def fit_eval_and_predict(initial_weights, input_fn, model_fn, + distribution=None, is_stateful_model=False): + """Generates results for fit/predict/evaluate for given model.""" model = model_fn(initial_weights=initial_weights, distribution=distribution) training_inputs, eval_inputs, predict_inputs = input_fn(distribution) @@ -301,7 +244,15 @@ def fit_eval_and_predict( result['weights_1'] = model.get_weights() if predict_inputs is not None: - result['predict_result_1'] = model.predict(**predict_inputs) + # Check correctness of the result of predict() invoked + # multiple times -- as for stateful models, result of + # predict may differ for each batch. + predict_length = 1 + if is_stateful_model: + predict_length = 3 + for i in range(predict_length): + result_key = 'predict_result_{}'.format(i) + result[result_key] = model.predict(**predict_inputs) # Train and eval again to mimic user's flow. @@ -317,20 +268,23 @@ def fit_eval_and_predict( def compare_results(results_with_ds, results_without_ds, distribution, testcase): + """Compares results of model compiled with/without distribution strategy.""" + default_tolerance = 1e-5 - tol_table = {} - - if isinstance(distribution, ( - mirrored_strategy.MirroredStrategy, - mirrored_strategy.CoreMirroredStrategy, - distribute_lib._DefaultDistributionStrategy)): # pylint: disable=protected-access - # TODO(b/119257215): Weights are not exactly the same, so use larger - # tolerance for now. Predict should be related to weights. - tol_table = { - 'weights_1': 1e-4, - 'weights_2': 1e-4, - 'predict_result_1': 1e-4, - } + relaxed_tolerance = 1e-4 + + def _get_compare_result_tolerance(key): + """Returns tolerance to compare results.""" + # TODO(b/119257215): For MirroredStrategy, weights are not exactly the same, + # so use larger tolerance for now. Predict should be related to weights. + if (isinstance(distribution, ( + mirrored_strategy.MirroredStrategy, + mirrored_strategy.CoreMirroredStrategy, + distribute_lib._DefaultDistributionStrategy)) and # pylint: disable=protected-access + key.startswith(('weights_1', 'weights_2', 'predict_result'))): + return relaxed_tolerance + + return default_tolerance for key in results_with_ds: if (key.startswith('training_history') and @@ -340,8 +294,7 @@ def compare_results(results_with_ds, results_without_ds, distribution, # underlying bug is fixed. continue - tolerance = tol_table.get(key, default_tolerance) - + tolerance = _get_compare_result_tolerance(key) testcase.assertAllClose( results_with_ds[key], results_without_ds[key], @@ -350,7 +303,13 @@ def compare_results(results_with_ds, results_without_ds, distribution, msg='Fail to assert {}.'.format(key)) +def should_skip_tpu_with_eager(distribution): + return (context.executing_eagerly() and + isinstance(distribution, tpu_strategy.TPUStrategy)) + + class LearningRateBatchScheduler(keras.callbacks.Callback): + """Scheduler that dynamically sets the learning rate of model.""" def __init__(self, update_freq=None): self._update_freq = update_freq @@ -364,107 +323,61 @@ class LearningRateBatchScheduler(keras.callbacks.Callback): keras.backend.set_value(self.model.optimizer.lr, lr) -class TestDistributionStrategyCorrectness(test.TestCase, - parameterized.TestCase): +class TestDistributionStrategyCorrectnessBase(test.TestCase, + parameterized.TestCase): + """Model agnostic testing infra to test correctness of Keras models.""" - def _should_skip_tpu_with_eager(self, distribution): - return (context.executing_eagerly() and - isinstance(distribution, tpu_strategy.TPUStrategy)) + def set_up_test_config(self, use_numpy=False, + use_validation_data=False, + with_batch_norm=False): + self.use_numpy = use_numpy + self.use_validation_data = use_validation_data + self.with_batch_norm = with_batch_norm - @combinations.generate(all_strategy_combinations_with_eager_and_graph_modes()) - def test_metric_correctness(self, distribution): - if self._should_skip_tpu_with_eager(distribution): - self.skipTest('TPUStrategy does not support eager mode now.') + keras.backend.set_image_data_format('channels_last') + np.random.seed(_RANDOM_SEED) + random_seed.set_random_seed(_RANDOM_SEED) - with self.cached_session(): - keras.backend.set_image_data_format('channels_last') - num_samples = 10000 - - x_train = np.random.randint(0, 2, num_samples) - x_train = np.reshape(x_train, (num_samples, 1)) - y_train = x_train - x_train = x_train.astype('float32') - y_train = y_train.astype('float32') - - # Create identity model. - with distribution.scope(): - model = keras.Sequential() - model.add( - keras.layers.Dense(1, input_shape=(1,), kernel_initializer='ones')) - model.compile( - loss=keras.losses.mean_squared_error, - optimizer=gradient_descent.GradientDescentOptimizer(0.5), - metrics=[keras.metrics.BinaryAccuracy()]) - - batch_size = 64 - batch_size = get_batch_size(batch_size, distribution) - train_dataset = dataset_ops.Dataset.from_tensor_slices((x_train, y_train)) - train_dataset = batch_wrapper(train_dataset, batch_size, distribution) - - history = model.fit(x=train_dataset, epochs=2, steps_per_epoch=10) - self.assertEqual(history.history['binary_accuracy'], [1.0, 1.0]) - - @combinations.generate(all_strategy_combinations_with_eager_and_graph_modes()) - def test_eval_metrics_correctness(self, distribution): - if self._should_skip_tpu_with_eager(distribution): - self.skipTest('TPUStrategy does not support eager mode now.') + def get_data(self): + num_samples = 10000 + x_train = np.random.randint(0, 2, num_samples) + x_train = np.reshape(x_train, (num_samples, 1)) + y_train = x_train + return (x_train.astype('float32'), y_train.astype('float32'), None) - with self.cached_session(): - with distribution.scope(): - model = keras.Sequential() - model.add( - keras.layers.Dense( - 3, activation='relu', input_dim=4, kernel_initializer='ones')) - model.add( - keras.layers.Dense( - 1, activation='sigmoid', kernel_initializer='ones')) - model.compile( - loss='mae', - metrics=['accuracy', keras.metrics.BinaryAccuracy()], - optimizer=gradient_descent.GradientDescentOptimizer(0.001)) - - # verify correctness of stateful and stateless metrics. - x = np.ones((100, 4)).astype('float32') - y = np.ones((100, 1)).astype('float32') - dataset = dataset_ops.Dataset.from_tensor_slices((x, y)).repeat() - dataset = batch_wrapper(dataset, 4, distribution) - outs = model.evaluate(dataset, steps=10) - self.assertEqual(outs[1], 1.) - self.assertEqual(outs[2], 1.) - - y = np.zeros((100, 1)).astype('float32') - dataset = dataset_ops.Dataset.from_tensor_slices((x, y)).repeat() - dataset = batch_wrapper(dataset, 4, distribution) - outs = model.evaluate(dataset, steps=10) - self.assertEqual(outs[1], 0.) - self.assertEqual(outs[2], 0.) - - @combinations.generate(strategy_and_input_combinations()) - def test_correctness(self, distribution, use_numpy, use_validation_data, - model_with_data): - if self._should_skip_tpu_with_eager(distribution): + def get_model(self, distribution=None): + raise NotImplementedError + + def skip_unsupported_test_configuration(self, distribution): + if should_skip_tpu_with_eager(distribution): self.skipTest('TPUStrategy does not support eager mode now.') - if context.executing_eagerly() and use_numpy: + if context.executing_eagerly() and self.use_numpy: self.skipTest('Numpy as inputs is not supported with strategy in eager.') - if context.executing_eagerly() and use_validation_data: - self.skipTest('TODO') - + if context.executing_eagerly() and self.use_validation_data: + self.skipTest('TODO(hongjunchoi): Add test logic for using validation ' + 'data for eager execution.') + return + + def run_correctness_test(self, + distribution, + use_numpy, + use_validation_data, + with_batch_norm=False, + is_stateful_model=False): with self.cached_session(): - keras.backend.set_image_data_format('channels_last') - np.random.seed(_RANDOM_SEED) - random_seed.set_random_seed(_RANDOM_SEED) + self.set_up_test_config(use_numpy, use_validation_data, with_batch_norm) + self.skip_unsupported_test_configuration(distribution) - model_fn, data_fn = model_with_data.model_fn, model_with_data.data_fn # Train, eval, and predict datasets are created with the same input numpy # arrays. - x_train, y_train, x_predict = data_fn() + x_train, y_train, x_predict = self.get_data() # The model is built once and the initial weights are saved. # This is used to initialize the model for both the distribution and # non-distribution run. - model = model_fn() + model = self.get_model() initial_weights = model.get_weights() def input_fn(dist): @@ -472,15 +385,15 @@ class TestDistributionStrategyCorrectness(test.TestCase, use_numpy, use_validation_data, dist, x_train, y_train, x_predict) results_with_ds = fit_eval_and_predict( - initial_weights, input_fn=input_fn, model_fn=model_fn, - distribution=distribution) + initial_weights, input_fn=input_fn, model_fn=self.get_model, + distribution=distribution, is_stateful_model=is_stateful_model) results_without_ds = fit_eval_and_predict( - initial_weights, input_fn=input_fn, model_fn=model_fn, - distribution=None) + initial_weights, input_fn=input_fn, model_fn=self.get_model, + distribution=None, is_stateful_model=is_stateful_model) # First, special case, for multi-replica distributed training, batch norm # is not aggregated globally. So it is expected to have different weights. - if (model_with_data.with_batch_norm and + if (self.with_batch_norm and distribution.num_replicas_in_sync > 1): with self.assertRaises(AssertionError): compare_results(results_with_ds, results_without_ds, distribution, @@ -489,21 +402,16 @@ class TestDistributionStrategyCorrectness(test.TestCase, compare_results(results_with_ds, results_without_ds, distribution, testcase=self) - @combinations.generate(all_strategy_combinations_with_graph_mode()) - def test_dynamic_lr(self, distribution): - + def run_dynamic_lr_test(self, distribution): with self.cached_session(): + self.set_up_test_config() + self.skip_unsupported_test_configuration(distribution) - keras.backend.set_image_data_format('channels_last') - np.random.seed(_RANDOM_SEED) - random_seed.set_random_seed(_RANDOM_SEED) - - x_train, y_train, _ = _dnn_training_data() - - model = _create_dnn_model() + x_train, y_train, _ = self.get_data() + model = self.get_model() initial_weights = model.get_weights() - update_freq = None + if (isinstance(distribution, tpu_strategy.TPUStrategy) and distribution.extended.steps_per_run > 1): # For TPUStrategy with steps_per_run > 1, the callback is not invoked @@ -512,6 +420,7 @@ class TestDistributionStrategyCorrectness(test.TestCase, update_freq = distribution.extended.steps_per_run def input_fn(dist): + """Generates training test given test configuration.""" training_epochs = 2 global_batch_size = 64 batch_size = get_batch_size(global_batch_size, dist) @@ -530,14 +439,49 @@ class TestDistributionStrategyCorrectness(test.TestCase, return training_inputs, eval_inputs, predict_inputs results_with_ds = fit_eval_and_predict( - initial_weights, input_fn=input_fn, model_fn=_create_dnn_model, + initial_weights, input_fn=input_fn, model_fn=self.get_model, distribution=distribution) results_without_ds = fit_eval_and_predict( - initial_weights, input_fn=input_fn, model_fn=_create_dnn_model, + initial_weights, input_fn=input_fn, model_fn=self.get_model, distribution=None) compare_results(results_with_ds, results_without_ds, distribution, testcase=self) +class TestDistributionStrategyEmbeddingModelCorrectnessBase( + TestDistributionStrategyCorrectnessBase): + """Base class to test correctness of Keras models with embedding layers.""" + + def get_data(self, + count=(_GLOBAL_BATCH_SIZE * _EVAL_STEPS), + min_words=5, + max_words=10, + max_word_id=19, + num_classes=2): + distribution = [] + for _ in range(num_classes): + dist = np.abs(np.random.randn(max_word_id)) + dist /= np.sum(dist) + distribution.append(dist) + + features = [] + labels = [] + for _ in range(count): + label = np.random.randint(0, num_classes, size=1)[0] + num_words = np.random.randint(min_words, max_words, size=1)[0] + word_ids = np.random.choice( + max_word_id, size=num_words, replace=True, p=distribution[label]) + word_ids = word_ids + labels.append(label) + features.append(word_ids) + + features = keras.preprocessing.sequence.pad_sequences( + features, maxlen=max_words) + x_train = np.asarray(features, dtype=np.float32) + y_train = np.asarray(labels, dtype=np.int32).reshape((count, 1)) + x_predict = x_train[:_GLOBAL_BATCH_SIZE] + return x_train, y_train, x_predict + + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/distribute/python/keras_dnn_correctness_test.py b/tensorflow/contrib/distribute/python/keras_dnn_correctness_test.py new file mode 100644 index 0000000000000000000000000000000000000000..61202e30c4f33892d2675080fae07cc4d7102337 --- /dev/null +++ b/tensorflow/contrib/distribute/python/keras_dnn_correctness_test.py @@ -0,0 +1,173 @@ +# 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. +# ============================================================================== +"""Correctness tests for tf.keras DNN model using DistributionStrategy.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.distribute.python import combinations +from tensorflow.contrib.distribute.python import keras_correctness_test_base +from tensorflow.python import keras +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.eager import test +from tensorflow.python.keras.optimizer_v2 import gradient_descent as gradient_descent_keras +from tensorflow.python.training import gradient_descent + + +def all_strategy_combinations_with_eager_and_graph_modes(): + return combinations.combine(distribution=keras_correctness_test_base. + all_strategies, + mode=['graph', 'eager']) + + +def all_strategy_combinations_with_graph_mode(): + return combinations.combine(distribution=keras_correctness_test_base. + all_strategies, mode=['graph']) + + +class TestDistributionStrategyDnnCorrectness( + keras_correctness_test_base.TestDistributionStrategyCorrectnessBase): + + def get_model(self, initial_weights=None, distribution=None): + with keras_correctness_test_base.MaybeDistributionScope(distribution): + # We add few non-linear layers to make it non-trivial. + model = keras.Sequential() + model.add(keras.layers.Dense(10, activation='relu', input_shape=(1,))) + model.add(keras.layers.Dense( + 10, activation='relu', + kernel_regularizer=keras.regularizers.l2(1e-4))) + model.add(keras.layers.Dense(10, activation='relu')) + model.add(keras.layers.Dense(1)) + + if initial_weights: + model.set_weights(initial_weights) + + model.compile( + loss=keras.losses.mean_squared_error, + optimizer=gradient_descent_keras.SGD(0.5), + metrics=['mse']) + return model + + def get_data(self): + # TODO(xiejw): Change this back to 10000, once we support final partial + # batch. + num_samples = 9984 + x_train = np.random.rand(num_samples, 1) + y_train = 3 * x_train + x_train = x_train.astype('float32') + y_train = y_train.astype('float32') + x_predict = np.array([[1.], [2.], [3.], [4.]], dtype=np.float32) + return x_train, y_train, x_predict + + @combinations.generate(keras_correctness_test_base. + all_strategy_and_input_config_combinations()) + def test_dnn_correctness(self, distribution, use_numpy, use_validation_data): + self.run_correctness_test(distribution, use_numpy, use_validation_data) + + @combinations.generate(all_strategy_combinations_with_graph_mode()) + def test_dnn_with_dynamic_learning_rate(self, distribution): + self.run_dynamic_lr_test(distribution) + + +class TestDistributionStrategyDnnMetricCorrectness( + keras_correctness_test_base.TestDistributionStrategyCorrectnessBase): + + def get_model(self, distribution=None): + with distribution.scope(): + model = keras.Sequential() + model.add(keras.layers.Dense(1, + input_shape=(1,), + kernel_initializer='ones')) + model.compile( + loss=keras.losses.mean_squared_error, + optimizer=gradient_descent.GradientDescentOptimizer(0.5), + metrics=[keras.metrics.BinaryAccuracy()]) + return model + + def run_metric_correctness_test(self, distribution): + with self.cached_session(): + self.set_up_test_config() + self.skip_unsupported_test_configuration(distribution) + + x_train, y_train, _ = self.get_data() + model = self.get_model(distribution=distribution) + + batch_size = 64 + batch_size = (keras_correctness_test_base. + get_batch_size(batch_size, distribution)) + train_dataset = dataset_ops.Dataset.from_tensor_slices((x_train, y_train)) + train_dataset = (keras_correctness_test_base. + batch_wrapper(train_dataset, batch_size, distribution)) + + history = model.fit(x=train_dataset, epochs=2, steps_per_epoch=10) + self.assertEqual(history.history['binary_accuracy'], [1.0, 1.0]) + + @combinations.generate(all_strategy_combinations_with_eager_and_graph_modes()) + def test_simple_dnn_metric_correctness(self, distribution): + self.run_metric_correctness_test(distribution) + + +class TestDistributionStrategyDnnMetricEvalCorrectness( + keras_correctness_test_base.TestDistributionStrategyCorrectnessBase): + + def get_model(self, distribution=None): + with distribution.scope(): + model = keras.Sequential() + model.add( + keras.layers.Dense( + 3, activation='relu', input_dim=4, kernel_initializer='ones')) + model.add( + keras.layers.Dense( + 1, activation='sigmoid', kernel_initializer='ones')) + model.compile( + loss='mae', + metrics=['accuracy', keras.metrics.BinaryAccuracy()], + optimizer=gradient_descent.GradientDescentOptimizer(0.001)) + return model + + def run_eval_metrics_correctness_test(self, distribution): + with self.cached_session(): + self.set_up_test_config() + self.skip_unsupported_test_configuration(distribution) + + model = self.get_model(distribution=distribution) + + # verify correctness of stateful and stateless metrics. + x = np.ones((100, 4)).astype('float32') + y = np.ones((100, 1)).astype('float32') + dataset = dataset_ops.Dataset.from_tensor_slices((x, y)).repeat() + dataset = (keras_correctness_test_base. + batch_wrapper(dataset, 4, distribution)) + outs = model.evaluate(dataset, steps=10) + self.assertEqual(outs[1], 1.) + self.assertEqual(outs[2], 1.) + + y = np.zeros((100, 1)).astype('float32') + dataset = dataset_ops.Dataset.from_tensor_slices((x, y)).repeat() + dataset = (keras_correctness_test_base. + batch_wrapper(dataset, 4, distribution)) + outs = model.evaluate(dataset, steps=10) + self.assertEqual(outs[1], 0.) + self.assertEqual(outs[2], 0.) + + @combinations.generate(all_strategy_combinations_with_eager_and_graph_modes()) + def test_identity_model_metric_eval_correctness(self, distribution): + self.run_eval_metrics_correctness_test(distribution) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/distribute/python/keras_embedding_model_correctness_test.py b/tensorflow/contrib/distribute/python/keras_embedding_model_correctness_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e881bb70ecc428e3f972cde5f19c1b61b1dc0f0b --- /dev/null +++ b/tensorflow/contrib/distribute/python/keras_embedding_model_correctness_test.py @@ -0,0 +1,150 @@ +# 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. +# ============================================================================== +"""Correctness test for tf.keras Embedding models using DistributionStrategy.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.distribute.python import combinations +from tensorflow.contrib.distribute.python import keras_correctness_test_base +from tensorflow.python import keras +from tensorflow.python.eager import test +from tensorflow.python.training import gradient_descent + + +class DistributionStrategyEmbeddingModelCorrectnessTest( + keras_correctness_test_base. + TestDistributionStrategyEmbeddingModelCorrectnessBase): + + def get_model(self, max_words=10, initial_weights=None, distribution=None): + with keras_correctness_test_base.MaybeDistributionScope(distribution): + word_ids = keras.layers.Input( + shape=(max_words,), dtype=np.int32, name='words') + word_embed = keras.layers.Embedding(input_dim=20, + output_dim=10)(word_ids) + if self.use_distributed_dense: + word_embed = keras.layers.TimeDistributed(keras.layers.Dense(4))( + word_embed) + avg = keras.layers.GlobalAveragePooling1D()(word_embed) + preds = keras.layers.Dense(2, activation='softmax')(avg) + model = keras.Model(inputs=[word_ids], outputs=[preds]) + + if initial_weights: + model.set_weights(initial_weights) + + model.compile( + optimizer=gradient_descent.GradientDescentOptimizer( + learning_rate=0.1), + loss='sparse_categorical_crossentropy', + metrics=['sparse_categorical_accuracy']) + return model + + @combinations.generate(keras_correctness_test_base. + test_combinations_for_embedding_model()) + def test_embedding_model_correctness(self, distribution, use_numpy, + use_validation_data): + + self.use_distributed_dense = False + self.run_correctness_test(distribution, use_numpy, use_validation_data) + + @combinations.generate(keras_correctness_test_base. + test_combinations_for_embedding_model()) + def test_embedding_time_distributed_model_correctness(self, + distribution, + use_numpy, + use_validation_data): + self.use_distributed_dense = True + self.run_correctness_test(distribution, use_numpy, use_validation_data) + + +class DistributionStrategySiameseEmbeddingModelCorrectnessTest( + keras_correctness_test_base. + TestDistributionStrategyEmbeddingModelCorrectnessBase): + + def get_model(self, max_words=10, initial_weights=None, distribution=None): + with keras_correctness_test_base.MaybeDistributionScope(distribution): + word_ids_a = keras.layers.Input( + shape=(max_words,), dtype=np.int32, name='words_a') + word_ids_b = keras.layers.Input( + shape=(max_words,), dtype=np.int32, name='words_b') + + def submodel(embedding, word_ids): + word_embed = embedding(word_ids) + rep = keras.layers.GlobalAveragePooling1D()(word_embed) + return keras.Model(inputs=[word_ids], outputs=[rep]) + + word_embed = keras.layers.Embedding( + input_dim=20, + output_dim=10, + input_length=max_words, + embeddings_initializer=keras.initializers.RandomUniform(0, 1)) + + a_rep = submodel(word_embed, word_ids_a).outputs[0] + b_rep = submodel(word_embed, word_ids_b).outputs[0] + sim = keras.layers.Dot(axes=1, normalize=True)([a_rep, b_rep]) + + model = keras.Model(inputs=[word_ids_a, word_ids_b], outputs=[sim]) + + if initial_weights: + model.set_weights(initial_weights) + + model.compile( + optimizer=gradient_descent.GradientDescentOptimizer( + learning_rate=0.1), + loss='mse', + metrics=['mse']) + return model + + def get_data(self, + count=(keras_correctness_test_base._GLOBAL_BATCH_SIZE * + keras_correctness_test_base._EVAL_STEPS), + min_words=5, + max_words=10, + max_word_id=19, + num_classes=2): + features_a, labels_a, _ = (super( + DistributionStrategySiameseEmbeddingModelCorrectnessTest, self). + get_data(count, min_words, max_words, + max_word_id, num_classes)) + + features_b, labels_b, _ = (super( + DistributionStrategySiameseEmbeddingModelCorrectnessTest, self). + get_data(count, min_words, max_words, + max_word_id, num_classes)) + + y_train = np.zeros((count, 1), dtype=np.float32) + y_train[labels_a == labels_b] = 1.0 + y_train[labels_a != labels_b] = -1.0 + # TODO(b/123360757): Add tests for using list as inputs for multi-input + # models. + x_train = { + 'words_a': features_a, + 'words_b': features_b, + } + x_predict = x_train + + return x_train, y_train, x_predict + + @combinations.generate(keras_correctness_test_base. + test_combinations_for_embedding_model()) + def test_siamese_embedding_model_correctness(self, distribution, use_numpy, + use_validation_data): + self.run_correctness_test(distribution, use_numpy, use_validation_data) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/distribute/python/keras_image_model_correctness_test.py b/tensorflow/contrib/distribute/python/keras_image_model_correctness_test.py new file mode 100644 index 0000000000000000000000000000000000000000..3c2961456b2eede9570ce29f7a8900834f2ccfb7 --- /dev/null +++ b/tensorflow/contrib/distribute/python/keras_image_model_correctness_test.py @@ -0,0 +1,93 @@ +# 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. +# ============================================================================== +"""Correctness tests for tf.keras CNN models using DistributionStrategy.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.distribute.python import combinations +from tensorflow.contrib.distribute.python import keras_correctness_test_base +from tensorflow.python import keras +from tensorflow.python.eager import test +from tensorflow.python.keras.optimizer_v2 import gradient_descent + + +class DistributionStrategyCnnCorrectnessTest( + keras_correctness_test_base.TestDistributionStrategyCorrectnessBase): + + def get_model(self, initial_weights=None, distribution=None): + with keras_correctness_test_base.MaybeDistributionScope(distribution): + image = keras.layers.Input(shape=(28, 28, 3), name='image') + c1 = keras.layers.Conv2D( + name='conv1', filters=16, kernel_size=(3, 3), strides=(4, 4), + kernel_regularizer=keras.regularizers.l2(1e-4))( + image) + if self.with_batch_norm: + c1 = keras.layers.BatchNormalization(name='bn1')(c1) + c1 = keras.layers.MaxPooling2D(pool_size=(2, 2))(c1) + logits = keras.layers.Dense( + 10, activation='softmax', name='pred')( + keras.layers.Flatten()(c1)) + model = keras.Model(inputs=[image], outputs=[logits]) + + if initial_weights: + model.set_weights(initial_weights) + + model.compile( + optimizer=gradient_descent.SGD( + learning_rate=0.1), + loss='sparse_categorical_crossentropy', + metrics=['sparse_categorical_accuracy']) + + return model + + def get_data(self, + count=keras_correctness_test_base._GLOBAL_BATCH_SIZE + * keras_correctness_test_base._EVAL_STEPS, + shape=(28, 28, 3), + num_classes=10): + centers = np.random.randn(num_classes, *shape) + + features = [] + labels = [] + for _ in range(count): + label = np.random.randint(0, num_classes, size=1)[0] + offset = np.random.normal(loc=0, scale=0.1, size=np.prod(shape)) + offset = offset.reshape(shape) + labels.append(label) + features.append(centers[label] + offset) + + x_train = np.asarray(features, dtype=np.float32) + y_train = np.asarray(labels, dtype=np.float32).reshape((count, 1)) + x_predict = x_train + return x_train, y_train, x_predict + + @combinations.generate(keras_correctness_test_base. + all_strategy_and_input_config_combinations()) + def test_cnn_correctness(self, distribution, use_numpy, use_validation_data): + self.run_correctness_test(distribution, use_numpy, use_validation_data) + + @combinations.generate(keras_correctness_test_base. + all_strategy_and_input_config_combinations()) + def test_cnn_with_batch_norm_correctness(self, distribution, use_numpy, + use_validation_data): + self.run_correctness_test(distribution, use_numpy, use_validation_data, + with_batch_norm=True) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/distribute/python/keras_lstm_model_correctness_test.py b/tensorflow/contrib/distribute/python/keras_lstm_model_correctness_test.py new file mode 100644 index 0000000000000000000000000000000000000000..9ed2dfa206cdf4be24a88b1d54090487c1873399 --- /dev/null +++ b/tensorflow/contrib/distribute/python/keras_lstm_model_correctness_test.py @@ -0,0 +1,65 @@ +# 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. +# ============================================================================== +"""Correctness tests for tf.keras LSTM model using DistributionStrategy.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.distribute.python import combinations +from tensorflow.contrib.distribute.python import keras_correctness_test_base +from tensorflow.python import keras +from tensorflow.python.eager import test +from tensorflow.python.training import gradient_descent + + +class DistributionStrategyLstmModelCorrectnessTest( + keras_correctness_test_base. + TestDistributionStrategyEmbeddingModelCorrectnessBase): + + def get_model(self, max_words=10, initial_weights=None, distribution=None): + with keras_correctness_test_base.MaybeDistributionScope(distribution): + word_ids = keras.layers.Input( + shape=(max_words,), dtype=np.int32, name='words') + word_embed = keras.layers.Embedding(input_dim=20, + output_dim=10)(word_ids) + lstm_embed = keras.layers.LSTM(units=4, + return_sequences=False)(word_embed) + + preds = keras.layers.Dense(2, activation='softmax')(lstm_embed) + model = keras.Model(inputs=[word_ids], outputs=[preds]) + + if initial_weights: + model.set_weights(initial_weights) + + model.compile( + optimizer=gradient_descent.GradientDescentOptimizer( + learning_rate=0.1), + loss='sparse_categorical_crossentropy', + metrics=['sparse_categorical_accuracy']) + return model + + @combinations.generate(keras_correctness_test_base. + test_combinations_for_embedding_model()) + def test_lstm_model_correctness(self, + distribution, + use_numpy, + use_validation_data): + self.run_correctness_test(distribution, use_numpy, use_validation_data) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/distribute/python/keras_optimizer_v2_test.py b/tensorflow/contrib/distribute/python/keras_optimizer_v2_test.py index cce93b3c10a2ac7bd1c594a5027b9d51629bb915..c93d7afa7ceef2c9c272e91997e2871655cea079 100644 --- a/tensorflow/contrib/distribute/python/keras_optimizer_v2_test.py +++ b/tensorflow/contrib/distribute/python/keras_optimizer_v2_test.py @@ -24,6 +24,7 @@ import numpy as np from tensorflow.contrib.distribute.python import combinations from tensorflow.python import keras from tensorflow.python.distribute import distribution_strategy_context as ds_context +from tensorflow.python.eager import context from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -44,107 +45,71 @@ def get_model(): class MirroredStrategyOptimizerV2Test(test.TestCase, parameterized.TestCase): - @combinations.generate(combinations.combine( - distribution=[ - combinations.mirrored_strategy_with_gpu_and_cpu, - combinations.core_mirrored_strategy_with_gpu_and_cpu], - mode=['graph'])) + @combinations.generate( + combinations.combine( + distribution=[ + combinations.core_mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_two_gpus, + combinations.parameter_server_strategy_with_two_gpus, + ], + mode=['graph', 'eager'])) def testKerasOptimizerWithUnequalInput(self, distribution): - def create_fn(): + with distribution.scope(): var = variables.Variable( 2.0, name='var', aggregation=variable_scope.VariableAggregation.SUM) - # grad for cpu is 1, grad for gpu is 2, avg grad is 1.5. - def loss(): - return math_ops.cast(_replica_id() + 1, dtype=dtypes.float32) * var - optimizer = adam.Adam(learning_rate=0.01, beta_1=0.2, beta_2=0.2) - train_op = optimizer.minimize(loss, var_list=[var]) - m = optimizer.get_slot(var, 'm') - v = optimizer.get_slot(var, 'v') - return (var, m, v, train_op, optimizer.iterations) + all_vars = [] - devices = ['/device:GPU:0', '/device:CPU:0'] - with distribution.scope(): - (var, m, v, op, counter) = distribution.call_for_each_replica(create_fn) + def model_fn(): + + def loss_fn(): + replica_id = _replica_id() + return math_ops.cast(replica_id + 1, dtype=dtypes.float32) * 0.5 * var + + train_op = optimizer.minimize(loss_fn, var_list=[var]) + + return train_op, optimizer + + def train_fn(): + train_op, optimizer = distribution.extended.call_for_each_replica( + model_fn) + if not all_vars: + all_vars.append(var) + all_vars.append(optimizer.get_slot(var, 'm')) + all_vars.append(optimizer.get_slot(var, 'v')) + return distribution.group(train_op) + + if not context.executing_eagerly(): + with self.cached_session() as sess: + train_fn = sess.make_callable(train_fn()) self.evaluate(variables.global_variables_initializer()) - var_val = [2.0, 2.0, 2.0] - self.assertAllClose( - var_val, - self.evaluate( - [distribution.extended.read_var(var), - var.get(devices[0]), - var.get(devices[1])])) - self.assertAllClose([0, 0, 0], - self.evaluate([ - distribution.extended.read_var(counter), - counter.get(devices[0]), - counter.get(devices[1]) - ])) - - train_op = distribution.unwrap(op) - self.evaluate(train_op) - # m(1) = beta1 * m(0) + (1-beta1) * grad = 0.2 * 0 + 0.8 * (1 + 2) / 2 - m_val = [1.2, 1.2, 1.2] - # assert slot variables in both replicas are the same. - self.assertAllClose( - m_val, - self.evaluate( - [distribution.extended.read_var(m), - m.get(devices[0]), - m.get(devices[1])])) - # v(1) = beta2 * v(0) + (1-beta2) * grad^2 = 0.2 * 0 + 0.8 * 2.25 - v_val = [1.8, 1.8, 1.8] - self.assertAllClose( - v_val, - self.evaluate( - [distribution.extended.read_var(v), - v.get(devices[0]), - v.get(devices[1])])) + + # first step. + train_fn() # var(1) = var(0) - lr * m(1) * sqrt(1 - beta2) / sqrt(v(1)) / (1 - beta1) # = 2.0 - 0.01 * 1.2 * sqrt(0.8) / sqrt(1.8) / 0.8 - var_val = [1.99, 1.99, 1.99] - self.assertAllClose( - var_val, - self.evaluate( - [distribution.extended.read_var(var), - var.get(devices[0]), - var.get(devices[1])])) - self.assertAllClose([1, 1, 1], - self.evaluate([ - distribution.extended.read_var(counter), - counter.get(devices[0]), - counter.get(devices[1]) - ])) - - self.evaluate(train_op) + self.assertAllClose(1.99, self.evaluate(all_vars[0])) + # m(1) = beta1 * m(0) + (1-beta1) * grad = 0.2 * 0 + 0.8 * (1 + 2) / 2 + self.assertAllClose(1.2, self.evaluate(all_vars[1])) + # v(1) = beta2 * v(0) + (1-beta2) * grad^2 = 0.2 * 0 + 0.8 * 2.25 + self.assertAllClose(1.8, self.evaluate(all_vars[2])) + + # second step. + train_fn() + # var(1) = var(0) - lr * 2 = 1.98 + self.assertAllClose(1.98, self.evaluate(all_vars[0])) # m(2) = beta1 * m(1) + (1-beta1) * grad = 0.2 * 1.2 + 0.8 * 1.5 - m_val = [1.44, 1.44, 1.44] - self.assertAllClose( - m_val, - self.evaluate( - [distribution.extended.read_var(m), - m.get(devices[0]), - m.get(devices[1])])) + self.assertAllClose(1.44, self.evaluate(all_vars[1])) # v(2) = beta2 * v(1) + (1-beta2) * grad^2 = 0.2 * 1.8 + 0.8 * 2.25 - v_val = [2.16, 2.16, 2.16] - self.assertAllClose( - v_val, - self.evaluate( - [distribution.extended.read_var(v), - v.get(devices[0]), - v.get(devices[1])])) - self.assertAllClose([2, 2, 2], - self.evaluate([ - distribution.extended.read_var(counter), - counter.get(devices[0]), - counter.get(devices[1]) - ])) - - @combinations.generate(combinations.combine( - distribution=[ - combinations.mirrored_strategy_with_gpu_and_cpu, - combinations.core_mirrored_strategy_with_gpu_and_cpu], - mode=['graph'])) + self.assertAllClose(2.16, self.evaluate(all_vars[2])) + + @combinations.generate( + combinations.combine( + distribution=[ + combinations.core_mirrored_strategy_with_gpu_and_cpu, + combinations.parameter_server_strategy_with_two_gpus, + ], + mode=['graph', 'eager'])) def testOptimizerWithKerasModelAndNumpyArrays(self, distribution): with self.cached_session(): diff --git a/tensorflow/contrib/distribute/python/keras_stateful_lstm_model_correctness_test.py b/tensorflow/contrib/distribute/python/keras_stateful_lstm_model_correctness_test.py new file mode 100644 index 0000000000000000000000000000000000000000..f5faf6c36b880a72bafc8d082cff2816f3b11a76 --- /dev/null +++ b/tensorflow/contrib/distribute/python/keras_stateful_lstm_model_correctness_test.py @@ -0,0 +1,99 @@ +# 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. +# ============================================================================== +"""Tests for stateful tf.keras LSTM models using DistributionStrategy.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.distribute.python import combinations +from tensorflow.contrib.distribute.python import keras_correctness_test_base +from tensorflow.python import keras +from tensorflow.python.eager import test +from tensorflow.python.training import gradient_descent + + +def strategies_for_stateful_embedding_model(): + """Returns TPUStrategy with single core device assignment.""" + + return [combinations.tpu_strategy_one_core, + combinations.tpu_strategy_one_step_one_core] + + +def test_combinations_for_stateful_embedding_model(): + return ( + combinations.combine( + distribution=strategies_for_stateful_embedding_model(), + mode='graph', + use_numpy=False, + use_validation_data=False + )) + + +class DistributionStrategyStatefulLstmModelCorrectnessTest( + keras_correctness_test_base. + TestDistributionStrategyEmbeddingModelCorrectnessBase): + + def get_model(self, max_words=10, initial_weights=None, distribution=None): + batch_size = keras_correctness_test_base._GLOBAL_BATCH_SIZE + + with keras_correctness_test_base.MaybeDistributionScope(distribution): + word_ids = keras.layers.Input( + shape=(max_words,), + batch_size=batch_size, + dtype=np.int32, name='words') + word_embed = keras.layers.Embedding(input_dim=20, + output_dim=10)(word_ids) + lstm_embed = keras.layers.LSTM(units=4, + return_sequences=False, + stateful=True)(word_embed) + + preds = keras.layers.Dense(2, activation='softmax')(lstm_embed) + model = keras.Model(inputs=[word_ids], outputs=[preds]) + + if initial_weights: + model.set_weights(initial_weights) + + model.compile( + optimizer=gradient_descent.GradientDescentOptimizer( + learning_rate=0.1), + loss='sparse_categorical_crossentropy', + metrics=['sparse_categorical_accuracy']) + return model + + @combinations.generate(test_combinations_for_stateful_embedding_model()) + def test_stateful_lstm_model_correctness(self, + distribution, + use_numpy, + use_validation_data): + self.run_correctness_test(distribution, use_numpy, use_validation_data, + is_stateful_model=True) + + @combinations.generate(keras_correctness_test_base. + test_combinations_with_tpu_strategies()) + def test_incorrectly_use_multiple_cores_for_stateful_lstm_model( + self, distribution, use_numpy, use_validation_data): + with self.assertRaisesRegexp(ValueError, + 'Single core must be used for computation ' + 'on stateful models. Consider adding ' + '`device_assignment` parameter to ' + 'TPUStrategy'): + self.run_correctness_test(distribution, use_numpy, use_validation_data, + is_stateful_model=True) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/distribute/python/keras_test.py b/tensorflow/contrib/distribute/python/keras_test.py index b99d11b923155b1362b434ffb82a96a11b08d352..2eca1d1877f36b0dacdf8abef3b3527c5db061ec 100644 --- a/tensorflow/contrib/distribute/python/keras_test.py +++ b/tensorflow/contrib/distribute/python/keras_test.py @@ -25,13 +25,11 @@ from tensorflow.contrib.distribute.python import combinations from tensorflow.contrib.distribute.python import mirrored_strategy from tensorflow.contrib.distribute.python import tpu_strategy from tensorflow.python import keras +from tensorflow.python.data.experimental.ops import cardinality from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.distribute import values from tensorflow.python.eager import test from tensorflow.python.estimator import keras as keras_lib from tensorflow.python.estimator import run_config as run_config_lib -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes from tensorflow.python.framework import test_util from tensorflow.python.keras import testing_utils from tensorflow.python.keras.engine import distributed_training_utils @@ -70,6 +68,18 @@ def simple_functional_model(): return model +def simple_multi_inputs_multi_outputs_model(): + input_a = keras.layers.Input(shape=(16,), name='input_a') + input_b = keras.layers.Input(shape=(16,), name='input_b') + + merged = keras.layers.concatenate([input_a, input_b], name='merge') + output_c = keras.layers.Dense(3, activation='softmax', name='dense_2')(merged) + output_d = keras.layers.Dense(2, activation='softmax', name='dense_3')(merged) + model = keras.models.Model( + inputs=[input_a, input_b], outputs=[output_c, output_d]) + return model + + def multi_inputs_multi_outputs_model(): input_a = keras.layers.Input(shape=(16,), name='input_a') input_b = keras.layers.Input(shape=(16,), name='input_b') @@ -202,6 +212,22 @@ def get_predict_dataset(distribution): return dataset +def convert_numpy_to_dataset_with_unknown_cardinality(inputs, + targets=None): + if targets is not None: + input_slices = (inputs, targets) + dummy_op = (lambda inp, target: True) + else: + input_slices = inputs + dummy_op = (lambda inp: True) + + original_dataset = (dataset_ops.Dataset.from_tensor_slices( + input_slices)) + ds_with_unknown_cardinality = (original_dataset.filter(dummy_op). + batch(10, drop_remainder=True)) + return ds_with_unknown_cardinality + + def multi_input_output_model(): a = keras.layers.Input(shape=(3,), name='input_a') b = keras.layers.Input(shape=(5,), name='input_b') @@ -216,9 +242,12 @@ def multi_input_output_model(): return model +# TODO(josh11b): Add combinations.one_device_strategy_gpu once it works with +# TestDistributionStrategyWithCallbacks.test_callbacks_in_predict. strategies_minus_tpu = [ combinations.default_strategy, combinations.one_device_strategy, + combinations.one_device_strategy_gpu, combinations.mirrored_strategy_with_gpu_and_cpu, combinations.mirrored_strategy_with_two_gpus, combinations.core_mirrored_strategy_with_gpu_and_cpu, @@ -230,36 +259,45 @@ tpu_strategies = [ def strategy_minus_tpu_combinations(): - return combinations.combine( - distribution=strategies_minus_tpu, - mode=['graph', 'eager']) + return combinations.combine(distribution=strategies_minus_tpu, + mode=['graph', 'eager']) def tpu_strategy_combinations(): - return combinations.combine( - distribution=tpu_strategies, - mode=['graph']) + return combinations.combine(distribution=tpu_strategies, + mode=['graph']) def all_strategy_combinations(): return strategy_minus_tpu_combinations() + tpu_strategy_combinations() -# TODO(priyag): Add v2 optimizers here. +def all_strategy_combinations_minus_default(): + strategy_minus_default_combinations = combinations.combine( + distribution=[ + combinations.one_device_strategy, + combinations.one_device_strategy_gpu, + 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], + mode=['graph', 'eager']) + return strategy_minus_default_combinations + tpu_strategy_combinations() + + def strategy_and_optimizer_combinations(): return combinations.times( all_strategy_combinations(), - combinations.combine( - optimizer=[combinations.adagrad_optimizer_v1_fn, - combinations.adam_optimizer_v1_fn, - combinations.gradient_descent_optimizer_v1_fn, - combinations.rmsprop_optimizer_v1_fn])) - - -def strategy_for_numpy_input_combinations(): - return combinations.combine( - distribution=strategies_minus_tpu + tpu_strategies, - mode=['graph']) + combinations.combine(optimizer=[ + combinations.adagrad_optimizer_v1_fn, + combinations.adagrad_optimizer_keras_v2_fn, + combinations.adam_optimizer_v1_fn, + combinations.adam_optimizer_keras_v2_fn, + combinations.gradient_descent_optimizer_v1_fn, + combinations.gradient_descent_optimizer_keras_v2_fn, + combinations.rmsprop_optimizer_v1_fn, + combinations.rmsprop_optimizer_keras_v2_fn + ])) class TestEstimatorDistributionStrategy(test_util.TensorFlowTestCase, @@ -417,16 +455,7 @@ class TestEstimatorDistributionStrategy(test_util.TensorFlowTestCase, class TestDistributionStrategyWithNumpyArrays(test.TestCase, parameterized.TestCase): - @combinations.generate(strategy_for_numpy_input_combinations()) - def test_creating_var_with_numpy_arrays(self, distribution): - with self.cached_session(): - x = np.asarray(np.random.random((64, 3)), dtype=np.float32) - var_x = distributed_training_utils.get_var_for_numpy(distribution, x) - val = self.evaluate(var_x.value()) - # Verify that the numpy value is copied to the variable. - self.assertAllEqual(x, val) - - @combinations.generate(strategy_for_numpy_input_combinations()) + @combinations.generate(all_strategy_combinations()) def test_calculating_input_params_no_steps_no_batch_size(self, distribution): # Calculate the per_replica_batch_size scaling factor for strategies # that use per_core_batch_size @@ -457,7 +486,7 @@ class TestDistributionStrategyWithNumpyArrays(test.TestCase, distributed_training_utils.get_input_params( distribution, input_63_samples, steps=None, batch_size=None) - @combinations.generate(strategy_for_numpy_input_combinations()) + @combinations.generate(all_strategy_combinations()) def test_calculating_input_params_with_steps_no_batch_size(self, distribution): # Calculate the per_replica_batch_size scaling factor for strategies @@ -503,7 +532,7 @@ class TestDistributionStrategyWithNumpyArrays(test.TestCase, distributed_training_utils.get_input_params( distribution, input_63_samples, steps=1, batch_size=None) - @combinations.generate(strategy_for_numpy_input_combinations()) + @combinations.generate(all_strategy_combinations()) def test_calculating_input_params_no_steps_with_batch_size(self, distribution): # Calculate the per_replica_batch_size scaling factor for strategies @@ -537,7 +566,7 @@ class TestDistributionStrategyWithNumpyArrays(test.TestCase, distributed_training_utils.get_input_params( distribution, input_64_samples, steps=None, batch_size=3) - @combinations.generate(strategy_for_numpy_input_combinations()) + @combinations.generate(all_strategy_combinations()) def test_calculating_input_params_with_steps_with_batch_size(self, distribution): with self.cached_session(): @@ -554,7 +583,7 @@ class TestDistributionStrategyWithNumpyArrays(test.TestCase, distributed_training_utils.get_input_params( distribution, input_64_samples, steps=10, batch_size=13) - @combinations.generate(strategy_for_numpy_input_combinations()) + @combinations.generate(all_strategy_combinations()) def test_calling_model_with_numpy_arrays(self, distribution): with self.cached_session(): with distribution.scope(): @@ -564,28 +593,28 @@ class TestDistributionStrategyWithNumpyArrays(test.TestCase, metrics = ['mae'] model.compile(optimizer, loss, metrics=metrics) - inputs = np.zeros((64, 3), dtype=np.float32) - targets = np.zeros((64, 4), dtype=np.float32) + inputs = np.zeros((64, 3), dtype=np.float32) + targets = np.zeros((64, 4), dtype=np.float32) - # Call fit with validation data - model.fit(inputs, targets, epochs=1, batch_size=2, verbose=0, - validation_data=(inputs, targets)) + # Call fit with validation data + model.fit(inputs, targets, epochs=1, batch_size=2, verbose=0, + validation_data=(inputs, targets)) - # TODO(anjalisridhar): We need tests for when the batch size and steps are - # smaller and results in a 0 batch_size and steps value. - model.evaluate(inputs, targets) - # with steps - model.evaluate(inputs, targets, steps=2) - # with batch_size - model.evaluate(inputs, targets, batch_size=8) + # TODO(anjalisridhar): We need tests for when the batch size and steps + # are smaller and results in a 0 batch_size and steps value. + model.evaluate(inputs, targets) + # with steps + model.evaluate(inputs, targets, steps=2) + # with batch_size + model.evaluate(inputs, targets, batch_size=8) - model.predict(inputs) - # with steps - model.predict(inputs, steps=2) - # with batch_size - model.predict(inputs, batch_size=8) + model.predict(inputs) + # with steps + model.predict(inputs, steps=2) + # with batch_size + model.predict(inputs, batch_size=8) - @combinations.generate(strategy_for_numpy_input_combinations()) + @combinations.generate(all_strategy_combinations()) def test_calling_model_with_nested_numpy_arrays(self, distribution): with self.cached_session(): with distribution.scope(): @@ -637,7 +666,7 @@ class TestDistributionStrategyWithNumpyArrays(test.TestCase, model.fit(inputs, targets, sample_weight=sample_weights, epochs=1, steps_per_epoch=2, verbose=1) - @combinations.generate(strategy_for_numpy_input_combinations()) + @combinations.generate(all_strategy_combinations()) def test_flatten_predict_outputs(self, distribution): with self.cached_session(): with distribution.scope(): @@ -662,6 +691,61 @@ class TestDistributionStrategyWithNumpyArrays(test.TestCase, self.assertAllEqual([6, 7], outs[0].shape) self.assertAllEqual([6, 7], outs[1].shape) + @combinations.generate(tpu_strategy_combinations()) + def test_predict_with_partial_batch(self, distribution): + with self.cached_session(): + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + + with distribution.scope(): + model_with_ds_strategy = get_model() + model_with_ds_strategy.compile(optimizer, loss) + + cpu_model = get_model() + cpu_model.compile(optimizer, loss) + + 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 + # distribution strategy uses entire sample as a single batch. As so, + # we remove parameters `batch_size` and `steps`. + 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) + + @combinations.generate(tpu_strategy_combinations()) + def test_predict_multi_output_model_with_partial_batch( + self, distribution): + with self.cached_session(): + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + + with distribution.scope(): + model_with_ds_strategy = simple_multi_inputs_multi_outputs_model() + model_with_ds_strategy.compile(optimizer, loss) + + cpu_model = simple_multi_inputs_multi_outputs_model() + cpu_model.compile(optimizer, loss) + + input_data, _ = get_multi_inputs_multi_outputs_data() + input_dict = { + 'input_a': input_data['input_a'], + 'input_b': input_data['input_b'], + } + + # As sample size is 200, we batch by 18 so that the last batch is + # a partial batch. Also `fit()` using numpy array as inputs without + # distribution strategy uses entire sample as a single batch. As so, + # we remove parameters `batch_size` and `steps`. + cpu_model.set_weights(model_with_ds_strategy.get_weights()) + self.assertAllClose( + model_with_ds_strategy.predict(input_dict, batch_size=18, steps=12), + cpu_model.predict(input_dict), + atol=1e-4, rtol=1e-4) + class TestDistributionStrategyWithDatasets(test.TestCase, parameterized.TestCase): @@ -765,6 +849,95 @@ class TestDistributionStrategyWithDatasets(test.TestCase, model.fit(dataset_dict, epochs=1, steps_per_epoch=2, verbose=1) + @combinations.generate(all_strategy_combinations()) + def test_fit_eval_and_predict_methods_on_dataset_without_steps( + self, distribution): + with self.cached_session(): + with distribution.scope(): + model = get_model() + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + metrics = ['mae', keras.metrics.CategoricalAccuracy()] + model.compile(optimizer, loss, metrics=metrics) + + inputs = np.zeros((1000, 3), dtype=np.float32) + targets = np.zeros((1000, 4), dtype=np.float32) + # steps/steps_per_epoch are calculated when using numpy arrays as + # input data. + fit_with_numpy = model.fit(inputs, targets, epochs=1, + batch_size=10).history + eval_with_numpy = model.evaluate(inputs, targets, batch_size=10) + predict_with_numpy = model.predict(inputs, batch_size=10) + + dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets)) + dataset = dataset.batch(10, drop_remainder=True) + fit_with_ds = model.fit(dataset, epochs=1).history + eval_with_ds = model.evaluate(dataset) + predict_dataset = dataset_ops.Dataset.from_tensor_slices(inputs) + predict_dataset = predict_dataset.batch(10, drop_remainder=True) + predict_with_ds = model.predict(predict_dataset) + self.assertAllClose( + fit_with_numpy, fit_with_ds, atol=1e-4, rtol=1e-4) + self.assertAllClose( + eval_with_numpy, eval_with_ds, atol=1e-4, rtol=1e-4) + self.assertAllClose( + predict_with_numpy, predict_with_ds, atol=1e-4, rtol=1e-4) + + @combinations.generate(all_strategy_combinations()) + def test_on_dataset_with_unknown_cardinality_without_steps( + self, distribution): + with self.cached_session(): + with distribution.scope(): + model = get_model() + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + metrics = ['mae', keras.metrics.CategoricalAccuracy()] + model.compile(optimizer, loss, metrics=metrics) + + inputs = np.zeros((1000, 3), dtype=np.float32) + targets = np.zeros((1000, 4), dtype=np.float32) + # steps/steps_per_epoch are calculated when using numpy arrays as + # input data. + fit_with_numpy = model.fit(inputs, targets, epochs=1, + batch_size=10).history + fit_with_numpy_multiple_epochs = model.fit( + inputs, targets, epochs=2, batch_size=10).history + eval_with_numpy = model.evaluate(inputs, targets, batch_size=10) + predict_with_numpy = model.predict(inputs, batch_size=10) + + dataset = convert_numpy_to_dataset_with_unknown_cardinality( + inputs, targets) + predict_dataset = convert_numpy_to_dataset_with_unknown_cardinality( + inputs) + + self.assertEqual(keras.backend.get_value(cardinality.cardinality( + dataset)), cardinality.UNKNOWN) + self.assertEqual(keras.backend.get_value(cardinality.cardinality( + predict_dataset)), cardinality.UNKNOWN) + + eval_with_ds = model.evaluate(dataset) + predict_with_ds = model.predict(predict_dataset) + self.assertAllClose( + eval_with_numpy, eval_with_ds, atol=1e-4, rtol=1e-4) + self.assertAllClose( + predict_with_numpy, predict_with_ds, atol=1e-4, rtol=1e-4) + + if (distributed_training_utils.is_tpu_strategy(distribution) and + distribution.extended.steps_per_run != 1): + with self.assertRaisesRegexp(ValueError, '`steps_per_epoch` ' + 'should be specified'): + fit_with_ds = model.fit(dataset, epochs=1) + else: + fit_with_ds = model.fit(dataset, + epochs=1).history + fit_with_ds_multiple_epochs = model.fit(dataset, + epochs=2).history + self.assertAllClose( + fit_with_numpy, fit_with_ds, atol=1e-4, rtol=1e-4) + self.assertAllClose( + fit_with_numpy_multiple_epochs, + fit_with_ds_multiple_epochs, atol=1e-4, rtol=1e-4) + @combinations.generate(all_strategy_combinations()) def test_fit_eval_and_predict_methods_on_dataset(self, distribution): with self.cached_session(): @@ -937,9 +1110,6 @@ class TestDistributionStrategyWithDatasets(test.TestCase, @combinations.generate(all_strategy_combinations()) def testOptimizerWithCallbacks(self, distribution): with self.cached_session(): - # TODO(b/120946189): Investigate why default strategy + eager fails. - if '_Default' in distribution.__class__.__name__: - self.skipTest('Disable the test for default strategy.') with distribution.scope(): model = get_model() optimizer = gradient_descent_keras.SGD(0.01) @@ -955,198 +1125,63 @@ class TestDistributionStrategyWithDatasets(test.TestCase, callbacks=[keras.callbacks.LearningRateScheduler(schedule)]) self.assertAllClose(0.001, keras.backend.get_value(model.optimizer.lr)) - -class TestDistributionStrategyErrorCases(test.TestCase, parameterized.TestCase): - - @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_validating_dataset_input_tensors_with_shape_mismatch(self, - distribution): - with self.cached_session(): - a = constant_op.constant([1, 2], shape=(1, 2)) - b = constant_op.constant([[1, 2], [1, 2]], shape=(2, 2)) - device_map = values.ReplicaDeviceMap(('/device:CPU:0', '/device:GPU:0')) - x = values.DistributedValues(device_map, (a, b)) - y = values.DistributedValues(device_map, (a, a)) - # Removed device and input tensor shape details from the error message - # since the order of the device and the corresponding input tensor shape - # is not deterministic over different runs. - with self.assertRaisesRegexp(ValueError, - 'Input tensor shapes do not match for ' - 'distributed tensor inputs ' - 'DistributedValues:.+'): - with distribution.scope(): - distributed_training_utils.validate_distributed_dataset_inputs( - distribution, x, y) - - @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_validating_dataset_input_tensors_with_dtype_mismatch(self, - distribution): + @combinations.generate(tpu_strategy_combinations()) + def test_predict_with_dataset_with_partial_batch(self, distribution): with self.cached_session(): - a = constant_op.constant([1, 2], shape=(1, 2), dtype=dtypes.int32) - b = constant_op.constant([1, 2], shape=(1, 2), dtype=dtypes.float64) - device_map = values.ReplicaDeviceMap(('/device:CPU:0', '/device:GPU:0')) - x = values.DistributedValues(device_map, (a, b)) - y = values.DistributedValues(device_map, (a, a)) - # Removed device and input tensor dtype details from the error message - # since the order of the device and the corresponding input tensor dtype - # is not deterministic over different runs. - with self.assertRaisesRegexp(ValueError, - 'Input tensor dtypes do not match for ' - 'distributed tensor inputs ' - 'DistributedValues:.+'): - with distribution.scope(): - distributed_training_utils.validate_distributed_dataset_inputs( - distribution, x, y) + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' - @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_unsupported_features(self, distribution): - with self.cached_session(): with distribution.scope(): - model = get_model() - optimizer = gradient_descent.GradientDescentOptimizer(0.001) - loss = 'mse' - metrics = ['mae'] - model.compile(optimizer, loss, metrics=metrics) + model_with_ds_strategy = get_model() + model_with_ds_strategy.compile(optimizer, loss) - dataset = get_dataset(distribution) + cpu_model = get_model() + cpu_model.compile(optimizer, loss) - # Test with validation split - with self.assertRaisesRegexp( - ValueError, '`validation_split` argument is not ' - 'supported when input `x` is a dataset or a ' - 'dataset iterator.+'): - model.fit(dataset, - epochs=1, steps_per_epoch=2, verbose=0, - validation_split=0.5, validation_steps=2) - - # Test with sample weight. - sample_weight = np.random.random((10,)) - with self.assertRaisesRegexp( - ValueError, '`sample_weight` argument is not supported when input ' - '`x` is a dataset or a dataset iterator.'): - model.fit( - dataset, - epochs=1, - steps_per_epoch=2, - verbose=0, - sample_weight=sample_weight) - - # Test with not specifying the `steps` argument. - with self.assertRaisesRegexp( - ValueError, 'the `steps_per_epoch` argument'): - model.fit(dataset, epochs=1, verbose=0) - with self.assertRaisesRegexp(ValueError, 'the `steps` argument'): - model.evaluate(dataset, verbose=0) - - with self.assertRaisesRegexp(ValueError, 'the `steps` argument'): - model.predict(dataset, verbose=0) - - @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_calling_with_unsupported_predefined_callbacks(self, distribution): - with self.cached_session(): - with distribution.scope(): - model = get_model() - optimizer = gradient_descent.GradientDescentOptimizer(0.001) - loss = 'mse' - metrics = ['mae'] - model.compile(optimizer, loss, metrics=metrics) + inputs = np.zeros((10, 3), dtype=np.float32) + dataset = dataset_ops.Dataset.from_tensor_slices((inputs)) - dataset = get_dataset(distribution) + # As sample size is 10, we batch by 4 so that the last batch is + # a partial batch. + dataset_with_partial_batch = dataset.batch(4) + cpu_model.set_weights(model_with_ds_strategy.get_weights()) - def schedule(_): - return 0.001 - with self.assertRaisesRegexp(ValueError, - 'You must specify a Keras Optimizer V2 when ' - 'using'): - model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0, - callbacks=[keras.callbacks.LearningRateScheduler(schedule)]) + self.assertAllClose( + model_with_ds_strategy.predict(dataset_with_partial_batch, steps=3), + cpu_model.predict(dataset_with_partial_batch, steps=3), + atol=1e-5, rtol=1e-5) - with self.assertRaisesRegexp(ValueError, - 'You must specify a Keras Optimizer V2 when ' - 'using'): - model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0, - callbacks=[keras.callbacks.ReduceLROnPlateau()]) + @combinations.generate(tpu_strategy_combinations()) + def test_predict_multi_output_model_with_dataset_with_partial_batch( + self, distribution): + with self.cached_session(): + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + with distribution.scope(): + model_with_ds_strategy = simple_multi_inputs_multi_outputs_model() + model_with_ds_strategy.compile(optimizer, loss) -class TestDistributionStrategyWithLossMasking(test.TestCase, - parameterized.TestCase): + cpu_model = simple_multi_inputs_multi_outputs_model() + cpu_model.compile(optimizer, loss) - # TODO(priyag): Enable all strategies for this test. Currently it does not - # work for TPU due to some invalid datatype. - @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_masking(self, distribution): - with self.cached_session(): - np.random.seed(1337) - x = np.array([[[1], [1]], [[0], [0]]]) - with distribution.scope(): - model = keras.models.Sequential() - model.add(keras.layers.Masking(mask_value=0, input_shape=(2, 1))) - model.add( - keras.layers.TimeDistributed( - keras.layers.Dense(1, kernel_initializer='one'))) - model.compile(loss='mse', - optimizer=gradient_descent.GradientDescentOptimizer(0.01)) - y = np.array([[[1], [1]], [[1], [1]]]) - dataset = dataset_ops.Dataset.from_tensor_slices((x, y)) - dataset = dataset.repeat(100) - dataset = dataset.batch(10) - hist = model.fit(x=dataset, epochs=1, steps_per_epoch=2) - self.assertEqual(hist.history['loss'][0], 0) + input_data, _ = get_multi_inputs_multi_outputs_data() + input_dict = { + 'input_a': input_data['input_a'], + 'input_b': input_data['input_b'], + } + dataset = dataset_ops.Dataset.from_tensor_slices(input_dict) -class TestDistributionStrategyWithNormalizationLayer( - test.TestCase, parameterized.TestCase): + # As sample size is 200, we batch by 18 using 12 steps per epoch so + # that the last batch is a partial batch. + dataset_with_partial_batch = dataset.batch(18) + cpu_model.set_weights(model_with_ds_strategy.get_weights()) - @combinations.generate(combinations.times( - all_strategy_combinations(), - combinations.combine(fused=[True, False]))) - def test_batchnorm_correctness(self, distribution, fused): - with self.cached_session(): - with distribution.scope(): - model = keras.models.Sequential() - norm = keras.layers.BatchNormalization( - input_shape=(10,), momentum=0.8, fused=fused) - model.add(norm) - model.compile(loss='mse', - optimizer=gradient_descent.GradientDescentOptimizer(0.01)) - - # centered on 5.0, variance 10.0 - x = np.random.normal(loc=5.0, scale=10.0, size=(1000, 10)) - x = x.astype('float32') - dataset = dataset_ops.Dataset.from_tensor_slices((x, x)) - dataset = dataset.repeat(100) - dataset = batch_wrapper(dataset, 32, distribution) - - predict_dataset = dataset_ops.Dataset.from_tensor_slices(x) - predict_dataset = predict_dataset.repeat(100) - predict_dataset = batch_wrapper(predict_dataset, 32, distribution) - - model.fit(dataset, epochs=4, verbose=0, steps_per_epoch=10) - out = model.predict(predict_dataset, steps=2) - out -= keras.backend.eval(norm.beta) - out /= keras.backend.eval(norm.gamma) - np.testing.assert_allclose(out.mean(), 0.0, atol=1e-1) - np.testing.assert_allclose(out.std(), 1.0, atol=1e-1) + self.assertAllClose( + model_with_ds_strategy.predict(dataset_with_partial_batch, steps=12), + cpu_model.predict(dataset_with_partial_batch, steps=12), + atol=1e-4, rtol=1e-4) if __name__ == '__main__': diff --git a/tensorflow/contrib/distribute/python/keras_utils_test.py b/tensorflow/contrib/distribute/python/keras_utils_test.py new file mode 100644 index 0000000000000000000000000000000000000000..3e5b422f512c0de5ed90ba73b56448ebbf4fc8a7 --- /dev/null +++ b/tensorflow/contrib/distribute/python/keras_utils_test.py @@ -0,0 +1,471 @@ +# Copyright 2016 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 tf.keras models with callbacks, checkpointing with dist strategy.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import collections +import tempfile +from absl.testing import parameterized +import numpy as np + +from tensorflow.contrib.distribute.python import combinations +from tensorflow.contrib.distribute.python import keras_test as keras_test_lib +from tensorflow.contrib.distribute.python import tpu_strategy +from tensorflow.python import keras +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.distribute import values +from tensorflow.python.eager import test +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.training import gradient_descent + + +class Counter(keras.callbacks.Callback): + """Counts the number of times each callback method was run. + + Attributes: + method_counts: dict. Contains the counts of time each callback method was + run. + """ + + def __init__(self): + self.method_counts = collections.defaultdict(int) + methods_to_count = [ + 'on_batch_begin', 'on_batch_end', 'on_epoch_begin', 'on_epoch_end', + 'on_predict_batch_begin', 'on_predict_batch_end', 'on_predict_begin', + 'on_predict_end', 'on_test_batch_begin', 'on_test_batch_end', + 'on_test_begin', 'on_test_end', 'on_train_batch_begin', + 'on_train_batch_end', 'on_train_begin', 'on_train_end' + ] + for method_name in methods_to_count: + setattr(self, method_name, + self.wrap_with_counts(method_name, getattr(self, method_name))) + + def wrap_with_counts(self, method_name, method): + + def _call_and_count(*args, **kwargs): + self.method_counts[method_name] += 1 + return method(*args, **kwargs) + + return _call_and_count + + +class TestDistributionStrategyWithCallbacks(test.TestCase, + parameterized.TestCase): + + @combinations.generate(keras_test_lib.all_strategy_combinations()) + def test_callbacks_in_fit(self, distribution): + with distribution.scope(): + model = keras_test_lib.get_model() + model.compile(optimizer='sgd', loss='mse', metrics=['mae']) + + dataset = keras_test_lib.get_dataset(distribution) + counter = Counter() + + epochs = 2 + steps_per_epoch = 5 + validation_steps = 3 + + model.fit( + dataset, + epochs=epochs, + steps_per_epoch=steps_per_epoch, + verbose=0, + validation_data=dataset, + validation_steps=validation_steps, + callbacks=[counter]) + + if isinstance(distribution, tpu_strategy.TPUStrategy): + # TPU Strategy can have multi step training, from extended.steps_per_run + # if steps_per_run = 1, then num_batch_call_per_epoch = steps_per_epoch + steps_per_run = distribution.extended.steps_per_run + num_batch_call_per_epoch = steps_per_epoch // steps_per_run + if steps_per_epoch % steps_per_run: + num_batch_call_per_epoch += 1 + else: + num_batch_call_per_epoch = steps_per_epoch + + self.assertDictEqual( + counter.method_counts, { + 'on_batch_begin': epochs * num_batch_call_per_epoch, + 'on_batch_end': epochs * num_batch_call_per_epoch, + 'on_epoch_begin': epochs, + 'on_epoch_end': epochs, + 'on_test_batch_begin': epochs * validation_steps, + 'on_test_batch_end': epochs * validation_steps, + 'on_test_begin': epochs, + 'on_test_end': epochs, + 'on_train_batch_begin': epochs * num_batch_call_per_epoch, + 'on_train_batch_end': epochs * num_batch_call_per_epoch, + 'on_train_begin': 1, + 'on_train_end': 1 + }) + + @combinations.generate(keras_test_lib.all_strategy_combinations()) + def test_callbacks_in_eval(self, distribution): + with distribution.scope(): + model = keras_test_lib.get_model() + model.compile(optimizer='sgd', loss='mse', metrics=['mae']) + + dataset = keras_test_lib.get_dataset(distribution) + counter = Counter() + + model.evaluate(dataset, steps=5, callbacks=[counter]) + + self.assertDictEqual( + counter.method_counts, { + 'on_test_batch_begin': 5, + 'on_test_batch_end': 5, + 'on_test_begin': 1, + 'on_test_end': 1 + }) + + @combinations.generate(keras_test_lib.all_strategy_combinations()) + def test_callbacks_in_predict(self, distribution): + with distribution.scope(): + model = keras_test_lib.get_model() + model.compile(optimizer='sgd', loss='mse', metrics=['mae']) + + dataset = keras_test_lib.get_dataset(distribution) + counter = Counter() + + model.predict( + keras_test_lib.get_predict_dataset(dataset), + steps=5, + callbacks=[counter]) + + self.assertDictEqual( + counter.method_counts, { + 'on_predict_batch_begin': 5, + 'on_predict_batch_end': 5, + 'on_predict_begin': 1, + 'on_predict_end': 1 + }) + + +class TestDistributionStrategyErrorCases(test.TestCase, parameterized.TestCase): + + @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_validating_dataset_input_tensors_with_shape_mismatch( + self, distribution): + with self.cached_session(): + a = constant_op.constant([1, 2], shape=(1, 2)) + b = constant_op.constant([[1, 2], [1, 2]], shape=(2, 2)) + device_map = values.ReplicaDeviceMap(('/device:CPU:0', '/device:GPU:0')) + x = values.DistributedValues(device_map, (a, b)) + y = values.DistributedValues(device_map, (a, a)) + # Removed device and input tensor shape details from the error message + # since the order of the device and the corresponding input tensor shape + # is not deterministic over different runs. + with self.assertRaisesRegexp( + ValueError, 'Input tensor shapes do not match for ' + 'distributed tensor inputs ' + 'DistributedValues:.+'): + with distribution.scope(): + distributed_training_utils.validate_distributed_dataset_inputs( + distribution, x, y) + + @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_validating_dataset_input_tensors_with_dtype_mismatch( + self, distribution): + with self.cached_session(): + a = constant_op.constant([1, 2], shape=(1, 2), dtype=dtypes.int32) + b = constant_op.constant([1, 2], shape=(1, 2), dtype=dtypes.float64) + device_map = values.ReplicaDeviceMap(('/device:CPU:0', '/device:GPU:0')) + x = values.DistributedValues(device_map, (a, b)) + y = values.DistributedValues(device_map, (a, a)) + # Removed device and input tensor dtype details from the error message + # since the order of the device and the corresponding input tensor dtype + # is not deterministic over different runs. + with self.assertRaisesRegexp( + ValueError, 'Input tensor dtypes do not match for ' + 'distributed tensor inputs ' + 'DistributedValues:.+'): + with distribution.scope(): + distributed_training_utils.validate_distributed_dataset_inputs( + distribution, x, y) + + @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_unsupported_features(self, distribution): + with self.cached_session(): + with distribution.scope(): + model = keras_test_lib.get_model() + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + metrics = ['mae'] + model.compile(optimizer, loss, metrics=metrics) + + dataset = keras_test_lib.get_dataset(distribution) + + # Test with validation split + with self.assertRaisesRegexp( + ValueError, '`validation_split` argument is not ' + 'supported when input `x` is a dataset or a ' + 'dataset iterator.+'): + model.fit( + dataset, + epochs=1, + steps_per_epoch=2, + verbose=0, + validation_split=0.5, + validation_steps=2) + + # Test with sample weight. + sample_weight = np.random.random((10,)) + with self.assertRaisesRegexp( + ValueError, '`sample_weight` argument is not supported when input ' + '`x` is a dataset or a dataset iterator.'): + model.fit( + dataset, + epochs=1, + steps_per_epoch=2, + verbose=0, + sample_weight=sample_weight) + + # Test with not specifying the `steps` argument for dataset with infinite + # cardinality. + dataset = dataset.repeat() + with self.assertRaisesRegexp( + ValueError, 'When passing an infinitely ' + 'repeating dataset, you must specify the ' + '`steps_per_epoch` argument'): + model.fit(dataset, epochs=1, verbose=0) + with self.assertRaisesRegexp( + ValueError, 'When passing an infinitely ' + 'repeating dataset, you must specify the ' + '`steps` argument'): + model.evaluate(dataset, verbose=0) + + with self.assertRaisesRegexp( + ValueError, 'When passing an infinitely ' + 'repeating dataset, you must specify the ' + '`steps` argument'): + model.predict(dataset, verbose=0) + + @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_calling_with_unsupported_predefined_callbacks(self, distribution): + with self.cached_session(): + with distribution.scope(): + model = keras_test_lib.get_model() + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + metrics = ['mae'] + model.compile(optimizer, loss, metrics=metrics) + + dataset = keras_test_lib.get_dataset(distribution) + + def schedule(_): + return 0.001 + + with self.assertRaisesRegexp( + ValueError, 'You must specify a Keras Optimizer V2 when ' + 'using'): + model.fit( + dataset, + epochs=1, + steps_per_epoch=2, + verbose=0, + callbacks=[keras.callbacks.LearningRateScheduler(schedule)]) + + with self.assertRaisesRegexp( + ValueError, 'You must specify a Keras Optimizer V2 when ' + 'using'): + model.fit( + dataset, + epochs=1, + steps_per_epoch=2, + verbose=0, + callbacks=[keras.callbacks.ReduceLROnPlateau()]) + + +class TestDistributionStrategyWithLossMasking(test.TestCase, + parameterized.TestCase): + + # TODO(priyag): Enable all strategies for this test. Currently it does not + # work for TPU due to some invalid datatype. + @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_masking(self, distribution): + with self.cached_session(): + np.random.seed(1337) + x = np.array([[[1], [1]], [[0], [0]]]) + with distribution.scope(): + model = keras.models.Sequential() + model.add(keras.layers.Masking(mask_value=0, input_shape=(2, 1))) + model.add( + keras.layers.TimeDistributed( + keras.layers.Dense(1, kernel_initializer='one'))) + model.compile( + loss='mse', + optimizer=gradient_descent.GradientDescentOptimizer(0.01)) + y = np.array([[[1], [1]], [[1], [1]]]) + dataset = dataset_ops.Dataset.from_tensor_slices((x, y)) + dataset = dataset.repeat(100) + dataset = dataset.batch(10) + hist = model.fit(x=dataset, epochs=1, steps_per_epoch=2) + self.assertEqual(hist.history['loss'][0], 0) + + +class TestDistributionStrategyWithNormalizationLayer(test.TestCase, + parameterized.TestCase): + + @combinations.generate( + combinations.times(keras_test_lib.all_strategy_combinations(), + combinations.combine(fused=[True, False]))) + def test_batchnorm_correctness(self, distribution, fused): + with self.cached_session(): + with distribution.scope(): + model = keras.models.Sequential() + norm = keras.layers.BatchNormalization( + input_shape=(10,), momentum=0.8, fused=fused) + model.add(norm) + model.compile( + loss='mse', + optimizer=gradient_descent.GradientDescentOptimizer(0.01)) + + # centered on 5.0, variance 10.0 + x = np.random.normal(loc=5.0, scale=10.0, size=(1000, 10)) + x = x.astype('float32') + dataset = dataset_ops.Dataset.from_tensor_slices((x, x)) + dataset = dataset.repeat(100) + dataset = keras_test_lib.batch_wrapper(dataset, 32, distribution) + + predict_dataset = dataset_ops.Dataset.from_tensor_slices(x) + predict_dataset = predict_dataset.repeat(100) + predict_dataset = keras_test_lib.batch_wrapper(predict_dataset, 32, + distribution) + + model.fit(dataset, epochs=4, verbose=0, steps_per_epoch=10) + out = model.predict(predict_dataset, steps=2) + out -= keras.backend.eval(norm.beta) + out /= keras.backend.eval(norm.gamma) + np.testing.assert_allclose(out.mean(), 0.0, atol=1e-1) + np.testing.assert_allclose(out.std(), 1.0, atol=1e-1) + + +class TestDistributionStrategySaveLoadWeights(test.TestCase, + parameterized.TestCase): + + @combinations.generate( + keras_test_lib.all_strategy_combinations_minus_default()) + def test_save_load_h5(self, distribution): + 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('.h5') + 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) + + @combinations.generate( + keras_test_lib.all_strategy_combinations_minus_default()) + def test_save_load_checkpointable(self, distribution): + # TODO(sourabhbajaj): Test fails with optimizer v2 without h5 + with self.cached_session(): + dataset = keras_test_lib.get_dataset(distribution) + with distribution.scope(): + model = keras_test_lib.get_model() + model.compile(gradient_descent.GradientDescentOptimizer(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.GradientDescentOptimizer(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): + + @combinations.generate( + keras_test_lib.all_strategy_combinations_minus_default()) + def test_layer_outside_scope(self, distribution): + with self.cached_session(): + with self.assertRaisesRegexp( + ValueError, 'was not created in the distribution strategy'): + x = keras.layers.Input(shape=(3,), name='input') + y = keras.layers.Dense(4, name='dense')(x) + with distribution.scope(): + model = keras.Model(x, y) + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + metrics = ['mae', keras.metrics.CategoricalAccuracy()] + model.compile(optimizer, loss, metrics=metrics) + + @combinations.generate( + keras_test_lib.all_strategy_combinations_minus_default()) + def test_model_outside_scope(self, distribution): + with self.cached_session(): + with self.assertRaisesRegexp( + ValueError, 'was not created in the distribution strategy'): + x = keras.layers.Input(shape=(3,), name='input') + y = keras.layers.Dense(4, name='dense')(x) + model = keras.Model(x, y) + with distribution.scope(): + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + metrics = ['mae', keras.metrics.CategoricalAccuracy()] + model.compile(optimizer, loss, metrics=metrics) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/distribute/python/metrics_v1_test.py b/tensorflow/contrib/distribute/python/metrics_v1_test.py index 32a0d199434e0627122fd4e47cf8894079ef3a1e..a663e809dd45ea099e1d8a08e681d07b05bee3c9 100644 --- a/tensorflow/contrib/distribute/python/metrics_v1_test.py +++ b/tensorflow/contrib/distribute/python/metrics_v1_test.py @@ -95,16 +95,15 @@ class MetricsV1Test(test.TestCase, parameterized.TestCase): def _test_metric(self, distribution, dataset_fn, metric_fn, expected_fn): with ops.Graph().as_default(), distribution.scope(): - iterator = distribution.distribute_dataset( - dataset_fn).make_initializable_iterator() + iterator = distribution.make_input_fn_iterator(lambda _: dataset_fn()) if isinstance(distribution, tpu_strategy.TPUStrategy): def step_fn(ctx, inputs): - value, update = distribution.call_for_each_replica( + value, update = distribution.extended.call_for_each_replica( metric_fn, args=(inputs,)) ctx.set_non_tensor_output(name="value", output=value) return distribution.group(update) - ctx = distribution.run_steps_on_dataset( + ctx = distribution.extended.experimental_run_steps_on_iterator( step_fn, iterator, iterations=distribution.extended.steps_per_run) update = ctx.run_op value = ctx.non_tensor_outputs["value"] @@ -114,15 +113,14 @@ class MetricsV1Test(test.TestCase, parameterized.TestCase): distribution.num_replicas_in_sync * distribution.extended.steps_per_run) else: - value, update = distribution.call_for_each_replica( + value, update = distribution.extended.call_for_each_replica( metric_fn, args=(iterator.get_next(),)) update = distribution.group(update) # TODO(josh11b): Once we switch to using a global batch size for input, # replace "distribution.num_replicas_in_sync" with "1". batches_per_update = distribution.num_replicas_in_sync - self.evaluate(iterator.initializer) - self.evaluate(distribution.initialize()) + self.evaluate(iterator.initialize()) self.evaluate(variables.local_variables_initializer()) batches_consumed = 0 @@ -136,8 +134,6 @@ class MetricsV1Test(test.TestCase, parameterized.TestCase): if batches_consumed >= 4: # Consume 4 input batches in total. break - self.evaluate(distribution.finalize()) - @combinations.generate(all_combinations() + tpu_combinations()) def testMean(self, distribution): def _dataset_fn(): diff --git a/tensorflow/contrib/distribute/python/minimize_loss_test.py b/tensorflow/contrib/distribute/python/minimize_loss_test.py index 824c4b09371fcc8d590f2d2b2be8f39b4a585b27..f06c9b75644b2890b7657f75e74e4e20a6f15705 100644 --- a/tensorflow/contrib/distribute/python/minimize_loss_test.py +++ b/tensorflow/contrib/distribute/python/minimize_loss_test.py @@ -41,12 +41,9 @@ from tensorflow.python.ops.losses import losses_impl class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): - def _get_iterator(self, ds): - if context.executing_eagerly(): - iterator = ds.make_one_shot_iterator() - else: - iterator = ds.make_initializable_iterator() - self.evaluate(iterator.initializer) + def _get_iterator(self, strategy, input_fn): + iterator = strategy.make_input_fn_iterator(lambda _: input_fn()) + self.evaluate(iterator.initialize()) return iterator @combinations.generate( @@ -67,15 +64,15 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): def step_fn(ctx, inputs): del ctx # Unused return distribution.group( - distribution.call_for_each_replica(model_fn, args=(inputs,))) + distribution.extended.call_for_each_replica( + model_fn, args=(inputs,))) - iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn)) + iterator = self._get_iterator(distribution, dataset_fn) def run_step(): - return distribution.run_steps_on_dataset( + return distribution.extended.experimental_run_steps_on_iterator( step_fn, iterator, iterations=2).run_op - self.evaluate(distribution.initialize()) if not context.executing_eagerly(): with self.cached_session() as sess: run_step = sess.make_callable(run_step()) @@ -84,12 +81,9 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): weights, biases = [], [] for _ in range(5): run_step() - weights.append(self.evaluate(layer.kernel)) biases.append(self.evaluate(layer.bias)) - self.evaluate(distribution.finalize()) - error = abs(numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1) is_not_increasing = all(y <= x for x, y in zip(error, error[1:])) self.assertTrue(is_not_increasing) @@ -105,11 +99,11 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): model_fn, dataset_fn, layer = minimize_loss_example( optimizer_fn, use_bias=True, use_callable_loss=use_callable_loss) - iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn)) + iterator = self._get_iterator(distribution, dataset_fn) def run_step(): return distribution.group( - distribution.call_for_each_replica( + distribution.extended.call_for_each_replica( model_fn, args=(iterator.get_next(),))) if not context.executing_eagerly(): @@ -152,7 +146,7 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): # `distribution.scope`. with variable_scope.variable_creator_scope( appending_creator), distribution.scope(): - model_fn, dataset_fn, layer = minimize_loss_example( + model_fn, dataset_fn, _ = minimize_loss_example( optimizer_fn, use_bias=True, use_callable_loss=True, @@ -161,24 +155,21 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): def step_fn(ctx, inputs): del ctx # Unused return distribution.group( - distribution.call_for_each_replica(model_fn, args=(inputs,))) + distribution.extended.call_for_each_replica( + model_fn, args=(inputs,))) - iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn)) + iterator = self._get_iterator(distribution, dataset_fn) def run_step(): - return distribution.run_steps_on_dataset( + return distribution.extended.experimental_run_steps_on_iterator( step_fn, iterator, iterations=1).run_op - self.evaluate(distribution.initialize()) if not context.executing_eagerly(): with self.cached_session() as sess: run_step = sess.make_callable(run_step()) self.evaluate(variables_lib.global_variables_initializer()) - run_step() - self.evaluate(distribution.finalize()) - def get_expected_variables(optimizer_fn, num_parameter_devices): variables_map = { "GradientDescent": ["dense/kernel", "dense/bias"], @@ -197,7 +188,7 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): self.assertEqual( get_expected_variables(optimizer_fn, - len(distribution.parameter_devices)), + len(distribution.extended.parameter_devices)), set(created_variables)) @combinations.generate( @@ -230,18 +221,18 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): def step_fn(ctx, inputs): del ctx # Unused fetches = distribution.unwrap( - distribution.call_for_each_replica(model_fn, args=(inputs,))) + distribution.extended.call_for_each_replica( + model_fn, args=(inputs,))) if update_ops_in_cross_replica_mode: fetches += tuple(ops.get_collection(ops.GraphKeys.UPDATE_OPS)) return control_flow_ops.group(fetches) - iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn)) + iterator = self._get_iterator(distribution, dataset_fn) def run_step(): - return distribution.run_steps_on_dataset( + return distribution.extended.experimental_run_steps_on_iterator( step_fn, iterator, iterations=1).run_op - self.evaluate(distribution.initialize()) if not context.executing_eagerly(): with self.cached_session() as sess: run_step = sess.make_callable(run_step()) @@ -267,8 +258,6 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): expected_moving_mean - averaged_batch_mean(i)) * (1.0 - momentum)) self.assertNear(expected_moving_means[i], moving_means[i], 0.0001) - self.evaluate(distribution.finalize()) - @combinations.generate( combinations.times( combinations.combine( @@ -327,15 +316,15 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): def step_fn(ctx, inputs): del ctx # Unused return distribution.group( - distribution.call_for_each_replica(model_fn, args=(inputs,))) + distribution.extended.call_for_each_replica( + model_fn, args=(inputs,))) - iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn)) + iterator = self._get_iterator(distribution, dataset_fn) def run_step(): - return distribution.run_steps_on_dataset( + return distribution.extended.experimental_run_steps_on_iterator( step_fn, iterator, iterations=1).run_op - self.evaluate(distribution.initialize()) if not context.executing_eagerly(): with self.cached_session() as sess: run_step = sess.make_callable(run_step()) @@ -370,8 +359,6 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): # One of the mean loss reductions. self.assertNear(weight, 2 + 10.6, 0.0001) - self.evaluate(distribution.finalize()) - @combinations.generate( combinations.times( combinations.distributions_and_v1_optimizers(), @@ -412,7 +399,7 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): return (train_op, loss) def step_fn(output_context, inputs): - (train_op, loss) = distribution.call_for_each_replica( + (train_op, loss) = distribution.extended.call_for_each_replica( model_fn, args=(output_context, inputs)) output_context.set_last_step_output( name="cross_replica_loss_reduced", @@ -423,7 +410,7 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): output=loss) return distribution.group(train_op) - iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn)) + iterator = self._get_iterator(distribution, dataset_fn) def run_step(): initial_loss = lambda: constant_op.constant(1e7) @@ -439,7 +426,7 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): "cross_replica_loss_not_reduced": distribution.unwrap(distribution.broadcast(initial_loss())) } - ctx = distribution.run_steps_on_dataset( + ctx = distribution.extended.experimental_run_steps_on_iterator( step_fn, iterator, iterations=2, initial_loop_values=initial_loop_values) @@ -458,7 +445,6 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): reduced=False, distribution=distribution) return (ctx.run_op, ctx.last_step_outputs["replica_loss_reduced"]) - self.evaluate(distribution.initialize()) if not context.executing_eagerly(): with self.cached_session() as sess: run_step = sess.make_callable(run_step()) @@ -471,8 +457,6 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): weights.append(self.evaluate(layer.kernel)) biases.append(self.evaluate(layer.bias)) - self.evaluate(distribution.finalize()) - loss_is_not_increasing = all(y <= x for x, y in zip(losses, losses[1:])) self.assertTrue(loss_is_not_increasing) diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy.py b/tensorflow/contrib/distribute/python/mirrored_strategy.py index db8fd9830785277b5d4c2ddbf1369752bae16c50..5391e083fc9b3ed99cc64bbed11bdeb8dea07f93 100644 --- a/tensorflow/contrib/distribute/python/mirrored_strategy.py +++ b/tensorflow/contrib/distribute/python/mirrored_strategy.py @@ -18,8 +18,6 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import functools - from tensorflow.python.distribute import distribute_lib from tensorflow.python.distribute import input_lib from tensorflow.python.distribute import mirrored_strategy @@ -48,8 +46,8 @@ class MirroredStrategy(distribute_lib.DistributionStrategy): distributed environment. There are several important concepts for distributed TensorFlow, e.g. - `client`, `job`, 'task', `cluster`, `in-graph replication` and - 'synchronous training' and they have already been defined in the + `client`, `job`, `task`, `cluster`, `in-graph replication` and + `synchronous training` and they have already been defined in the [TensorFlow's documentation](https://www.tensorflow.org/deploy/distributed). The distribution strategy inherits these concepts as well and in addition to that we also clarify several more concepts: @@ -104,6 +102,61 @@ class MirroredStrategy(distribute_lib.DistributionStrategy): auto_shard_dataset) super(MirroredStrategy, self).__init__(extended) + # Override to change the documentation to reflect the different handling of + # global vs. local batch size between core and contrib. + def make_dataset_iterator(self, dataset): # pylint: disable=useless-super-delegation + """Makes an iterator for input provided via `dataset`. + + NOTE: The batch size of the `dataset` argument is treated differently for + this contrib version of `MirroredStrategy`. + + Data from the given dataset will be distributed evenly across all the + compute replicas. We will assume that the input dataset is batched by the + per-replica batch size. + + The user could also use `make_input_fn_iterator` if they want to + customize which input is fed to which replica/worker etc. + + Args: + dataset: `tf.data.Dataset` that will be distributed evenly across all + replicas. + + Returns: + An `tf.distribute.InputIterator` which returns inputs for each step of the + computation. User should call `initialize` on the returned iterator. + """ + return super(MirroredStrategy, self).make_dataset_iterator(dataset) + + # Override to change the documentation to reflect the different handling of + # global vs. local batch size between core and contrib. + def experimental_make_numpy_iterator( # pylint: disable=useless-super-delegation + self, numpy_input, batch_size, num_epochs=1, shuffle=1024, session=None): + """Makes an iterator for input provided via a nest of numpy arrays. + + NOTE: The `batch_size` argument here has different behavior for this + contrib version of `MirroredStrategy`. + + Args: + numpy_input: A nest of NumPy input arrays that will be distributed evenly + across all replicas. + batch_size: The number of entries from the array we should consume in one + step of the computation, across all replicas. This is the per-replica + batch size. The global batch size will be this times + `num_replicas_in_sync`. + num_epochs: The number of times to iterate through the examples. A value + of `None` means repeat forever. + shuffle: Size of buffer to use for shuffling the input examples. + Use `None` to disable shuffling. + session: (TensorFlow v1.x graph execution only) A session used for + initialization. + + Returns: + An `tf.distribute.InputIterator` which returns inputs for each step of the + computation. User should call `initialize` on the returned iterator. + """ + return super(MirroredStrategy, self).experimental_make_numpy_iterator( + numpy_input, batch_size, num_epochs, shuffle, session) + class MirroredExtended(CoreMirroredExtended): """Implementation of (contrib) MirroredStrategy.""" @@ -137,17 +190,8 @@ class MirroredExtended(CoreMirroredExtended): """ return input_lib.DatasetIterator(dataset, self._input_workers) - def _distribute_dataset(self, dataset_fn): - if self._local_mode: - return input_lib.PerReplicaDataset( - self._call_dataset_fn(dataset_fn), self._input_workers, 0) - else: - return input_lib.MultiWorkerDataset( - functools.partial(self._call_dataset_fn, dataset_fn), - self._input_workers, - auto_shard=self._auto_shard_dataset) - # TODO(priyag): Delete this once all strategies use global batch size. @property def _global_batch_size(self): + """The contrib version of Mirrored strategy uses per-replica batch size.""" return False diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py index 9cbc34412da2d498fa1f0624f5cba466f663c194..bc0572bb4618967aa13599320218b63a5eec8d10 100644 --- a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py +++ b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py @@ -38,8 +38,8 @@ from tensorflow.python.eager import context from tensorflow.python.eager import function from tensorflow.python.eager import test from tensorflow.python.framework import constant_op -from tensorflow.python.framework import func_graph from tensorflow.python.framework import dtypes +from tensorflow.python.framework import func_graph from tensorflow.python.framework import ops from tensorflow.python.keras.engine import training as keras_training from tensorflow.python.keras.layers import core as keras_core @@ -66,8 +66,10 @@ GPU_TEST = "test_gpu" in sys.argv[0] combinations.core_mirrored_strategy_with_gpu_and_cpu, combinations.core_mirrored_strategy_with_two_gpus], mode=["graph", "eager"])) -class MirroredTwoDeviceDistributionTest(strategy_test_lib.DistributionTestBase, - parameterized.TestCase): +class MirroredTwoDeviceDistributionTest( + strategy_test_lib.DistributionTestBase, + strategy_test_lib.TwoDeviceDistributionTestBase, + parameterized.TestCase): def testMinimizeLoss(self, distribution): if context.executing_eagerly(): @@ -101,7 +103,7 @@ class MirroredTwoDeviceDistributionTest(strategy_test_lib.DistributionTestBase, expected = sum(range(distribution.num_replicas_in_sync)) self.assertEqual(expected, self.evaluate(reduced)) - def testMakeInputFnIterator(self, distribution): + def testMakeInputFnIteratorWithDataset(self, distribution): dataset_fn = lambda: dataset_ops.Dataset.range(10) expected_values = [[i, i+1] for i in range(0, 10, 2)] @@ -114,9 +116,47 @@ class MirroredTwoDeviceDistributionTest(strategy_test_lib.DistributionTestBase, self._test_input_fn_iterator(iterator, distribution.extended.worker_devices, expected_values) + def testMakeInputFnIteratorWithCallable(self, distribution): + def fn(): + dataset = dataset_ops.Dataset.range(2).interleave( + (lambda _: dataset_ops.Dataset.range(10)), cycle_length=2) + it = dataset.make_one_shot_iterator() + return it.get_next + expected_values = [[i, i] for i in range(0, 10)] + + input_fn = self._input_fn_to_test_input_context( + fn, + expected_num_replicas_in_sync=2, + expected_num_input_pipelines=1, + expected_input_pipeline_id=0) + iterator = distribution.make_input_fn_iterator(input_fn) + self._test_input_fn_iterator(iterator, distribution.extended.worker_devices, + expected_values, test_reinitialize=False) + + def testNumpyIterator(self, distribution): + self._test_numpy_iterator(distribution) + def testGlobalStepUpdate(self, distribution): self._test_global_step_update(distribution) + def testAllReduceSum(self, distribution): + self._test_all_reduce_sum(distribution) + + def testAllReduceSumGradients(self, distribution): + self._test_all_reduce_sum_gradients(distribution) + + def testAllReduceSumGradientTape(self, distribution): + self._test_all_reduce_sum_gradient_tape(distribution) + + def testAllReduceMean(self, distribution): + self._test_all_reduce_mean(distribution) + + def testAllReduceMeanGradients(self, distribution): + self._test_all_reduce_mean_gradients(distribution) + + def testAllReduceMeanGradientTape(self, distribution): + self._test_all_reduce_mean_gradient_tape(distribution) + def one_device_combinations(): return combinations.combine( @@ -128,25 +168,42 @@ def one_device_combinations(): mode=["graph", "eager"]) +@combinations.generate(one_device_combinations()) class MirroredOneDeviceDistributionTest( strategy_test_lib.DistributionTestBase, + strategy_test_lib.OneDeviceDistributionTestBase, parameterized.TestCase): - @combinations.generate(one_device_combinations()) def testMinimizeLoss(self, distribution): if context.executing_eagerly(): self._test_minimize_loss_eager(distribution) else: self._test_minimize_loss_graph(distribution) - @combinations.generate(one_device_combinations()) def testReplicaId(self, distribution): self._test_replica_id(distribution) - @combinations.generate(one_device_combinations()) def testCallAndMergeExceptions(self, distribution): self._test_call_and_merge_exceptions(distribution) + def testAllReduceSum(self, distribution): + self._test_all_reduce_sum(distribution) + + def testAllReduceSumGradients(self, distribution): + self._test_all_reduce_sum_gradients(distribution) + + def testAllReduceSumGradientTape(self, distribution): + self._test_all_reduce_sum_gradient_tape(distribution) + + def testAllReduceMean(self, distribution): + self._test_all_reduce_mean(distribution) + + def testAllReduceMeanGradients(self, distribution): + self._test_all_reduce_mean_gradients(distribution) + + def testAllReduceMeanGradientTape(self, distribution): + self._test_all_reduce_mean_gradient_tape(distribution) + class MirroredStrategyVariableCreatorStackTest( test.TestCase, parameterized.TestCase): @@ -227,7 +284,7 @@ class MirroredStrategyVariableCreationTest(test.TestCase): self.assertIs(strategy, var.distribute_strategy) for d in var.devices: self.assertEqual(d, var.get(d).device) - self.assertIs(strategy, var.get(d).distribute_strategy) + self.assertIs(strategy, var.get(d)._distribute_strategy) # pylint: disable=protected-access def testVariableInFuncGraph(self, distribution): def model_fn(): @@ -326,14 +383,9 @@ class MirroredStrategyVariableCreationTest(test.TestCase): (layer2.kernel, layer2.bias), (layer3.kernel, layer3.bias)] - ds = distribution.distribute_dataset( - lambda: dataset_ops.Dataset.from_tensors([[1.]]).repeat(10)) - if context.executing_eagerly(): - iterator = ds.make_one_shot_iterator() - else: - iterator = ds.make_initializable_iterator() - self.evaluate([iterator.initializer]) - + iterator = distribution.make_input_fn_iterator( + lambda _: dataset_ops.Dataset.from_tensors([[1.]]).repeat(10)) + self.evaluate(iterator.initialize()) features = iterator.get_next() with distribution.scope(): @@ -554,6 +606,7 @@ class MirroredStrategyVariableCreationTest(test.TestCase): aggregation="invalid") def testNonMatchingVariableCreation(self, distribution): + self.skipTest("b/123075960") def model_fn(name): v = variable_scope.variable(1.0, name=name) ds_context.get_replica_context().merge_call(lambda _: _) @@ -1226,7 +1279,7 @@ class MirroredStrategyDefunTest(test.TestCase): self.evaluate(device_result)) for defun in defuns: - # PolymorphicFunctions are specialized to the current device stack, so + # `Function`s are specialized to the current device stack, so # call_for_each has one trace per device. To check that the expected set # of variables was accessed on each trace, we first retrieve each # device-specific graph function. @@ -1381,7 +1434,7 @@ class MultiWorkerMirroredStrategyTest( self.assertEqual(a.device, "/job:worker/task:0") self.assertEqual(b.device, "/job:worker/task:0/device:CPU:0") - def testMakeInputFnIterator(self, distribution): + def testMakeInputFnIteratorWithDataset(self, distribution): self._configure_distribution_strategy(distribution) dataset_fn = lambda: dataset_ops.Dataset.range(100) num_gpus = context.num_gpus() @@ -1402,6 +1455,32 @@ class MultiWorkerMirroredStrategyTest( self._test_input_fn_iterator( iterator, distribution.extended.worker_devices, expected_values, sess) + def testMakeInputFnIteratorWithCallable(self, distribution): + self._configure_distribution_strategy(distribution) + def fn(): + dataset = dataset_ops.Dataset.range(100) + it = dataset.make_one_shot_iterator() + return it.get_next + num_gpus = context.num_gpus() + num_workers = 2 + + expected_values = [] + for i in range(0, 100, num_gpus): + expected_values.append([i+j for j in range(num_gpus)] * num_workers) + + with context.graph_mode(), self.cached_session() as sess: + # `expected_input_pipeline_id` is None because the input_fn will be called + # multiple times, each with a different input_pipeline_id. + input_fn = self._input_fn_to_test_input_context( + fn, + expected_num_replicas_in_sync=num_workers*num_gpus, + expected_num_input_pipelines=num_workers, + expected_input_pipeline_id=None) + iterator = distribution.make_input_fn_iterator(input_fn) + self._test_input_fn_iterator( + iterator, distribution.extended.worker_devices, expected_values, sess, + test_reinitialize=False) + def testUpdateConfigProto(self, distribution): distribution.configure(cluster_spec={"worker": ["fake1", "fake2"]}) diff --git a/tensorflow/contrib/distribute/python/monitor.py b/tensorflow/contrib/distribute/python/monitor.py index 17b7ab74f63f42e1ee14a82d3bffdd1df9b25857..53e35ea6b75088a3de9866973f872da4a4ce25d6 100644 --- a/tensorflow/contrib/distribute/python/monitor.py +++ b/tensorflow/contrib/distribute/python/monitor.py @@ -51,7 +51,7 @@ class Monitor(object): else: if session is None: raise ValueError("Should provide a `session` in Graph mode.") - session.run(step_callable._iterator.initializer) # pylint: disable=protected-access + session.run(step_callable.initialize()) self._run_step = session.make_callable(step_callable()) session.run(variables.global_variables_initializer()) diff --git a/tensorflow/contrib/distribute/python/monitor_test.py b/tensorflow/contrib/distribute/python/monitor_test.py index 16be839e1d155003b9490fbe3da6ab85b7d2d78a..c0651610cafc06a6d5f4206f4e64d27020fae30b 100644 --- a/tensorflow/contrib/distribute/python/monitor_test.py +++ b/tensorflow/contrib/distribute/python/monitor_test.py @@ -23,9 +23,9 @@ import numpy from tensorflow.contrib.distribute.python import combinations from tensorflow.contrib.distribute.python import monitor as monitor_lib -from tensorflow.contrib.distribute.python import one_device_strategy from tensorflow.contrib.distribute.python.single_loss_example import single_loss_example from tensorflow.python.client import session +from tensorflow.python.distribute import one_device_strategy from tensorflow.python.eager import context from tensorflow.python.eager import test from tensorflow.python.framework import ops diff --git a/tensorflow/contrib/distribute/python/moving_averages_test.py b/tensorflow/contrib/distribute/python/moving_averages_test.py index 8f13e9153ea7a951dd722c4549882c97e79b57fe..c4622cdd2af2f6a9c936fe554bcc2eb76f805fdc 100644 --- a/tensorflow/contrib/distribute/python/moving_averages_test.py +++ b/tensorflow/contrib/distribute/python/moving_averages_test.py @@ -53,7 +53,7 @@ class AssignMovingAveragesTest(test.TestCase, parameterized.TestCase): return var, assign with distribution.scope(), self.cached_session() as sess: - var, assign = distribution.call_for_each_replica(replica_fn) + var, assign = distribution.extended.call_for_each_replica(replica_fn) variables.global_variables_initializer().run() self.assertAllClose([10.0, 11.0], var.eval()) sess.run(distribution.unwrap(assign)) @@ -79,7 +79,7 @@ class AssignMovingAveragesTest(test.TestCase, parameterized.TestCase): return var, assign.op with distribution.scope(), self.cached_session() as sess: - var, assign_op = distribution.call_for_each_replica(replica_fn) + var, assign_op = distribution.extended.call_for_each_replica(replica_fn) variables.global_variables_initializer().run() self.assertAllClose([0.0, 0.0], var.eval()) sess.run(distribution.unwrap(assign_op)) @@ -152,7 +152,7 @@ class AssignMovingAveragesTest(test.TestCase, parameterized.TestCase): return var, assign with distribution.scope(), self.cached_session() as sess: - var, assign = distribution.call_for_each_replica(replica_fn) + var, assign = distribution.extended.call_for_each_replica(replica_fn) variables.global_variables_initializer().run() self.assertAllClose([10.0, 11.0], var.eval()) sess.run(distribution.unwrap(assign)) diff --git a/tensorflow/contrib/distribute/python/one_device_strategy.py b/tensorflow/contrib/distribute/python/one_device_strategy.py index fb470f8546fb19425cb77c201582794b2c192271..13a501394ee1fec2dfc1427f6d16d3a4624d7747 100644 --- a/tensorflow/contrib/distribute/python/one_device_strategy.py +++ b/tensorflow/contrib/distribute/python/one_device_strategy.py @@ -18,193 +18,6 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -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 values -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import control_flow_ops -from tensorflow.python.util import nest +from tensorflow.python.distribute import one_device_strategy - -# TODO(josh11b): Replace asserts in this file with if ...: raise ... - - -class OneDeviceStrategy(distribute_lib.DistributionStrategy): - """A distribution strategy for running on a single device.""" - # TODO(josh11b): Do we wrap values in types to generate errors if you are - # doing something that won't work with other DistributionStrategy - # implementations? - - def __init__(self, device): - super(OneDeviceStrategy, self).__init__(OneDeviceExtended(self, device)) - - -class OneDeviceExtended(distribute_lib.DistributionStrategyExtended): - """Implementation of OneDeviceStrategy.""" - - def __init__(self, container_strategy, device): - super(OneDeviceExtended, self).__init__(container_strategy) - self._device = device - self._default_device = device - worker = device_util.canonicalize("/device:CPU:0") - worker_device_pairs = [(worker, [self._device])] - device_map = values.SingleDeviceMap(device) - self._input_workers = input_lib.InputWorkers( - device_map, worker_device_pairs) - - def _create_variable(self, next_creator, *args, **kwargs): - colocate_with = kwargs.pop("colocate_with", None) - if colocate_with is None: - with ops.device(self._device): - return next_creator(*args, **kwargs) - with ops.colocate_with(colocate_with): - return next_creator(*args, **kwargs) - - def _validate_colocate_with_variable(self, colocate_with_variable): - values.validate_colocate(colocate_with_variable, self) - - def _make_dataset_iterator(self, dataset): - """Make iterator from dataset without splitting the batch.""" - return input_lib.DatasetIterator(dataset, self._input_workers) - - def _distribute_dataset(self, dataset_fn): - return input_lib.PerReplicaDataset( - self._call_dataset_fn(dataset_fn), self._input_workers, 0) - - def _make_input_fn_iterator( - self, - input_fn, - replication_mode=distribute_lib.InputReplicationMode.PER_WORKER): - return input_lib.InputFunctionIterator( - input_fn, self._input_workers, [distribute_lib.InputContext()]) - - def _broadcast_to(self, tensor, destinations): - del destinations - return tensor - - # TODO(priyag): Deal with OutOfRange errors once b/111349762 is fixed. - def _experimental_run_steps_on_iterator(self, fn, iterator, iterations, - initial_loop_values=None): - if initial_loop_values is None: - initial_loop_values = {} - initial_loop_values = nest.flatten(initial_loop_values) - - ctx = input_lib.MultiStepContext() - def body(i, *args): - """A wrapper around `fn` to create the while loop body.""" - del args - fn_result = fn(ctx, iterator.get_next()) - flat_last_step_outputs = nest.flatten(ctx.last_step_outputs) - with ops.control_dependencies([fn_result]): - return [i + 1] + flat_last_step_outputs - - # We capture the control_flow_context at this point, before we run `fn` - # inside a while_loop. 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 - - # TODO(priyag): Use max_iterations instead of an explicit counter. - cond = lambda i, *args: i < iterations - i = constant_op.constant(0) - loop_result = control_flow_ops.while_loop( - cond, body, [i] + initial_loop_values, name="", - parallel_iterations=1, back_prop=False, swap_memory=False, - return_same_structure=True) - del self._outer_control_flow_context - - ctx.run_op = control_flow_ops.group(loop_result) - - # Convert the last_step_outputs from a list to the original dict structure - # of last_step_outputs. - last_step_tensor_outputs = loop_result[1:] - last_step_tensor_outputs_dict = nest.pack_sequence_as( - ctx.last_step_outputs, last_step_tensor_outputs) - - ctx._set_last_step_outputs(last_step_tensor_outputs_dict) # pylint: disable=protected-access - return ctx - - def _call_for_each_replica(self, fn, args, kwargs): - strategy = self._container_strategy() - with ops.device(self._device), _OneDeviceReplicaContext(strategy): - return fn(*args, **kwargs) - - def _reduce_to(self, reduce_op, value, destinations): - del reduce_op, destinations - return value - - def _update(self, var, fn, args, kwargs, group): - # The implementations of _update() and _update_non_slot() are identical - # except _update() passes `var` as the first argument to `fn()`. - return self._update_non_slot(var, fn, (var,) + tuple(args), kwargs, group) - - def _update_non_slot(self, colocate_with, fn, args, kwargs, group): - del colocate_with - with ops.device(self._device), distribute_lib.UpdateContext(self._device): - result = fn(*args, **kwargs) - if group: - return result - else: - return nest.map_structure(self._unwrap, result) - - def read_var(self, replica_local_var): - """Read the aggregate value of a replica-local variable.""" - return array_ops.identity(replica_local_var) - - def _unwrap(self, value): - return (value,) - - def value_container(self, value): - return value - - @property - def _num_replicas_in_sync(self): - return 1 - - @property - def worker_devices(self): - return (self._device,) - - @property - def parameter_devices(self): - return (self._device,) - - def non_slot_devices(self, var_list): - del var_list - return (self._device,) - - @property - def experimental_should_init(self): - return True - - @property - def should_checkpoint(self): - return True - - @property - def should_save_summary(self): - return True - - # TODO(priyag): Delete this once all strategies use global batch size. - @property - def _global_batch_size(self): - return True - - -class _OneDeviceReplicaContext(distribute_lib.ReplicaContext): - """ReplicaContext for OneDeviceStrategy.""" - - def __init__(self, strategy): - zero = constant_op.constant(0, dtypes.int32) - distribute_lib.ReplicaContext.__init__( - self, strategy, replica_id_in_sync_group=zero) - - @property - def devices(self): - return self._strategy.extended.worker_devices +OneDeviceStrategy = one_device_strategy.OneDeviceStrategy diff --git a/tensorflow/contrib/distribute/python/one_device_strategy_test.py b/tensorflow/contrib/distribute/python/one_device_strategy_test.py index d46cd6f529e363f76bfa2b22339add63530cfde8..0e56f663d6a1ed7945befd933f2f4a83c5f64342 100644 --- a/tensorflow/contrib/distribute/python/one_device_strategy_test.py +++ b/tensorflow/contrib/distribute/python/one_device_strategy_test.py @@ -18,34 +18,35 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.distribute.python import one_device_strategy +from tensorflow.contrib.distribute.python import combinations from tensorflow.contrib.distribute.python import strategy_test_lib from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.eager import context from tensorflow.python.eager import test -from tensorflow.python.framework import test_util -class OneDeviceStrategyTest(strategy_test_lib.DistributionTestBase): +@combinations.generate(combinations.combine( + distribution=[ + combinations.one_device_strategy, + combinations.one_device_strategy_gpu], + mode=["eager", "graph"])) +class OneDeviceStrategyTest( + strategy_test_lib.DistributionTestBase, + strategy_test_lib.OneDeviceDistributionTestBase): - def _get_distribution_strategy(self): - return one_device_strategy.OneDeviceStrategy("/device:CPU:0") + def testMinimizeLoss(self, distribution): + if context.executing_eagerly(): + self._test_minimize_loss_eager(distribution) + else: + self._test_minimize_loss_graph(distribution) - def testMinimizeLossEager(self): - self._test_minimize_loss_eager(self._get_distribution_strategy()) + def testReplicaId(self, distribution): + self._test_replica_id(distribution) - def testMinimizeLossGraph(self): - self._test_minimize_loss_graph(self._get_distribution_strategy()) + def testCallAndMergeExceptions(self, distribution): + self._test_call_and_merge_exceptions(distribution) - def testReplicaId(self): - self._test_replica_id(self._get_distribution_strategy()) - - @test_util.run_in_graph_and_eager_modes - def testCallAndMergeExceptions(self): - self._test_call_and_merge_exceptions(self._get_distribution_strategy()) - - @test_util.run_in_graph_and_eager_modes - def testMakeInputFnIterator(self): - d = one_device_strategy.OneDeviceStrategy("/device:CPU:0") + def testMakeInputFnIteratorWithDataset(self, distribution): dataset_fn = lambda: dataset_ops.Dataset.range(10) expected_values = [[i] for i in range(10)] input_fn = self._input_fn_to_test_input_context( @@ -53,9 +54,46 @@ class OneDeviceStrategyTest(strategy_test_lib.DistributionTestBase): expected_num_replicas_in_sync=1, expected_num_input_pipelines=1, expected_input_pipeline_id=0) - iterator = d.make_input_fn_iterator(input_fn) + iterator = distribution.make_input_fn_iterator(input_fn) + self._test_input_fn_iterator( + iterator, distribution.extended.worker_devices, expected_values) + + def testMakeInputFnIteratorWithCallable(self, distribution): + def fn(): + dataset = dataset_ops.Dataset.range(10) + it = dataset.make_one_shot_iterator() + return it.get_next + expected_values = [[i] for i in range(10)] + input_fn = self._input_fn_to_test_input_context( + fn, + expected_num_replicas_in_sync=1, + expected_num_input_pipelines=1, + expected_input_pipeline_id=0) + iterator = distribution.make_input_fn_iterator(input_fn) self._test_input_fn_iterator( - iterator, d.extended.worker_devices, expected_values) + iterator, distribution.extended.worker_devices, expected_values, + test_reinitialize=False) + + def testNumpyIterator(self, distribution): + self._test_numpy_iterator(distribution) + + def testAllReduceSum(self, distribution): + self._test_all_reduce_sum(distribution) + + def testAllReduceSumGradients(self, distribution): + self._test_all_reduce_sum_gradients(distribution) + + def testAllReduceSumGradientTape(self, distribution): + self._test_all_reduce_sum_gradient_tape(distribution) + + def testAllReduceMean(self, distribution): + self._test_all_reduce_mean(distribution) + + def testAllReduceMeanGradients(self, distribution): + self._test_all_reduce_mean_gradients(distribution) + + def testAllReduceMeanGradientTape(self, distribution): + self._test_all_reduce_mean_gradient_tape(distribution) if __name__ == "__main__": diff --git a/tensorflow/contrib/distribute/python/optimizer_v2_test.py b/tensorflow/contrib/distribute/python/optimizer_v2_test.py index fa4705af7cb592119f56686d1f693a156f7b4b13..e388061b17a9b92dedbbf9839049b13c8575a22c 100644 --- a/tensorflow/contrib/distribute/python/optimizer_v2_test.py +++ b/tensorflow/contrib/distribute/python/optimizer_v2_test.py @@ -41,21 +41,17 @@ class MinimizeLossOptimizerV2Test(test.TestCase, parameterized.TestCase): with distribution.scope(): model_fn, dataset_fn, layer = minimize_loss_example( optimizer_fn, use_bias=True, use_callable_loss=use_callable_loss) - - ds = distribution.distribute_dataset(dataset_fn) - if context.executing_eagerly(): - iterator = ds.make_one_shot_iterator() - else: - iterator = ds.make_initializable_iterator() + iterator = distribution.make_input_fn_iterator(lambda _: dataset_fn()) def run_step(): - return control_flow_ops.group(distribution.unwrap( - distribution.call_for_each_replica( - model_fn, args=(iterator.get_next(),)))) + return control_flow_ops.group( + distribution.unwrap( + distribution.extended.call_for_each_replica( + model_fn, args=(iterator.get_next(),)))) if not context.executing_eagerly(): with self.cached_session() as sess: - sess.run(iterator.initializer) + sess.run(iterator.initialize()) run_step = sess.make_callable(run_step()) self.evaluate(variables.global_variables_initializer()) diff --git a/tensorflow/contrib/distribute/python/parameter_server_strategy.py b/tensorflow/contrib/distribute/python/parameter_server_strategy.py index 461e1bca21ea4ad7e1f534efd8f4ce53bb021d31..e42bc50fdc4e5e93c998708b0790fdea7768faf2 100644 --- a/tensorflow/contrib/distribute/python/parameter_server_strategy.py +++ b/tensorflow/contrib/distribute/python/parameter_server_strategy.py @@ -18,35 +18,24 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import copy - -from tensorflow.contrib.distribute.python import mirrored_strategy -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 multi_worker_util -from tensorflow.python.distribute import values -from tensorflow.python.eager import context -from tensorflow.python.framework import device as tf_device -from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import resource_variable_ops -from tensorflow.python.ops import variable_scope as vs -from tensorflow.python.platform import tf_logging as logging -from tensorflow.python.training import device_setter -from tensorflow.python.util import nest +from tensorflow.python.distribute import parameter_server_strategy +from tensorflow.python.distribute.cluster_resolver import SimpleClusterResolver +from tensorflow.python.distribute.cluster_resolver import TFConfigClusterResolver + +# pylint: disable=protected-access,invalid-name,line-too-long +CoreParameterServerStrategy = parameter_server_strategy.ParameterServerStrategy +CoreParameterServerExtended = parameter_server_strategy.ParameterServerStrategyExtended -_LOCAL_CPU = "/device:CPU:0" -_LOCAL_GPU_0 = "/device:GPU:0" +# pylint: enable=protected-access,invalid-name,line-too-long -# TODO(yuefengz): maybe cache variables on local CPU. -# TODO(yuefengz): we may want to set session options to disallow communication -# between workers. class ParameterServerStrategy(distribute_lib.DistributionStrategy): """A parameter server DistributionStrategy. + *** contrib version *** + This strategy class works for both local training and between-graph replicated training for multiple workers. If `cluster_spec` is specified, either passed in to __init__() method or parsed from the @@ -81,9 +70,9 @@ class ParameterServerStrategy(distribute_lib.DistributionStrategy): variables. 3) It is also not recommended to open a colocation scope (i.e. calling - `tf.colocate_with`) under the strategy's scope. For colocating variables, - use `distribution.colocate_vars_with` instead. Colocation of ops will possibly - create conflicts of device assignment. + `tf.colocate_with`) under the strategy's scope. For colocating variables, use + `strategy.extended.colocate_vars_with` instead. Colocation of ops will + possibly create conflicts of device assignment. """ def __init__(self, num_gpus_per_worker=0): @@ -100,437 +89,84 @@ class ParameterServerStrategy(distribute_lib.DistributionStrategy): super(ParameterServerStrategy, self).__init__( ParameterServerExtended(self, num_gpus_per_worker)) + # Override to change the documentation to reflect the different handling of + # global vs. local batch size between core and contrib. + def make_dataset_iterator(self, dataset): # pylint: disable=useless-super-delegation + """Makes an iterator for input provided via `dataset`. -class ParameterServerExtended(distribute_lib.DistributionStrategyExtended): - """Implementation of ParameterServerStrategy.""" + NOTE: The batch size of the `dataset` argument is treated differently for + this contrib version of `ParameterServerStrategy`. - def __init__(self, container_strategy, num_gpus_per_worker): - super(ParameterServerExtended, self).__init__(container_strategy) - self._num_gpus_per_worker = num_gpus_per_worker - self._initialize_local(num_gpus_per_worker) + Data from the given dataset will be distributed evenly across all the + compute replicas. We will assume that the input dataset is batched by the + per-replica batch size. - # We typically don't need to do all-reduce in this strategy. - self._cross_device_ops = ( - cross_device_ops_lib.ReductionToOneDeviceCrossDeviceOps( - reduce_to_device=_LOCAL_CPU)) - - def _initialize_multi_worker(self, num_gpus_per_worker, cluster_spec, - task_type, task_id): - """Initialize devices for multiple workers. - - It creates variable devices and compute devices. Variables and operations - will be assigned to them respectively. We have one compute device per - replica. The variable device is a device function or device string. The - default variable device assigns variables to parameter servers in a - round-robin fashion. + The user could also use `make_input_fn_iterator` if they want to + customize which input is fed to which replica/worker etc. Args: - num_gpus_per_worker: number of local GPUs or GPUs per worker. - cluster_spec: a dict, ClusterDef or ClusterSpec object specifying the - cluster configurations. - task_type: the current task type. - task_id: the current task id. + dataset: `tf.data.Dataset` that will be distributed evenly across all + replicas. - Raises: - ValueError: if the cluster_spec doesn't have ps jobs. + Returns: + An `tf.distribute.InputIterator` which returns inputs for each step of the + computation. User should call `initialize` on the returned iterator. """ - assert cluster_spec - if not task_type or task_id is None: - raise ValueError("When `cluster_spec` is given, you must also specify " - "`task_type` and `task_id`") - cluster_spec = multi_worker_util.normalize_cluster_spec(cluster_spec) - - worker_device = "/job:%s/task:%d" % (self._task_type, self._task_id) - - # Define compute devices which is a list of device strings and one for each - # replica. When there are GPUs, replicate operations on these GPUs. - # Otherwise, place operations on CPU. - if num_gpus_per_worker > 0: - compute_devices = tuple( - "%s/device:GPU:%d" % (worker_device, i) - for i in range(num_gpus_per_worker) - ) - else: - compute_devices = (worker_device,) - - self._device_map = values.ReplicaDeviceMap(compute_devices) - self._input_workers = input_lib.InputWorkers( - self._device_map, [(worker_device, compute_devices)]) - - # In distributed mode, place variables on ps jobs in a round-robin fashion. - # Note that devices returned from `replica_device_setter` are not - # canonical and therefore we don't canonicalize all variable devices to - # make them consistent. - # TODO(yuefengz): support passing a strategy object to control variable - # assignment. - # TODO(yuefengz): merge the logic of replica_device_setter into this - # class. - num_ps_replicas = len(cluster_spec.as_dict().get("ps", [])) - if num_ps_replicas == 0: - raise ValueError("The cluster spec needs to have `ps` jobs.") - self._variable_device = device_setter.replica_device_setter( - ps_tasks=num_ps_replicas, - worker_device=worker_device, - merge_devices=True, - cluster=cluster_spec) - - # The `_parameter_devices` is needed for the `parameter_devices` property - # and is a list of all variable devices. Here parameter devices are all - # tasks of the "ps" job. - self._parameter_devices = tuple(map("/job:ps/task:{}".format, - range(num_ps_replicas))) - - # Add a default device so that ops without specified devices will not end up - # on other workers. - self._default_device = worker_device - - self._is_chief = multi_worker_util.is_chief(cluster_spec, task_type, - task_id) - self._cluster_spec = cluster_spec - self._task_type = task_type - self._task_id = task_id - - logging.info( - "Multi-worker ParameterServerStrategy with " - "cluster_spec = %r, task_type = %r, task_id = %r, " - "num_ps_replicas = %r, is_chief = %r, device_map = %r, " - "variable_device = %r", cluster_spec.as_dict(), task_type, task_id, - num_ps_replicas, self._is_chief, self._device_map, - self._variable_device) - - def _initialize_local(self, num_gpus_per_worker): - """Initialize internal devices for local training.""" - worker_device = device_util.canonicalize("/device:CPU:0") - # Define compute devices which is a list of device strings and one for each - # replica. When there are GPUs, replicate operations on these GPUs. - # Otherwise, place operations on CPU. - if num_gpus_per_worker > 0: - compute_devices = tuple( - map("/device:GPU:{}".format, range(num_gpus_per_worker))) - else: - compute_devices = (_LOCAL_CPU,) - - self._device_map = values.ReplicaDeviceMap(compute_devices) - self._input_workers = input_lib.InputWorkers( - self._device_map, [(worker_device, compute_devices)]) - - # If there is only one GPU, put everything on that GPU. Otherwise, place - # variables on CPU. - if num_gpus_per_worker == 1: - assert len(compute_devices) == 1 - self._variable_device = _LOCAL_GPU_0 - self._parameter_devices = (_LOCAL_GPU_0,) - else: - self._variable_device = _LOCAL_CPU - self._parameter_devices = (_LOCAL_CPU,) - - self._is_chief = True - self._cluster_spec = None - self._task_type = None - self._task_id = None - - logging.info( - "ParameterServerStrategy with compute_devices = %r, " - "variable_device = %r", compute_devices, self._variable_device) - - def _validate_colocate_with_variable(self, colocate_with_variable): - values.validate_colocate(colocate_with_variable, self) - - def _distribute_dataset(self, dataset_fn): - """Distributes the dataset to each local GPU.""" - return input_lib.PerReplicaDataset( - self._call_dataset_fn(dataset_fn), self._input_workers, 0, - prefetch_on_device=True) - - def _make_dataset_iterator(self, dataset): - 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): - """Distributes the dataset to each local GPU.""" - if self._cluster_spec: - input_pipeline_id = multi_worker_util.id_in_cluster( - self._cluster_spec, self._task_type, self._task_id) - num_input_pipelines = multi_worker_util.worker_count( - self._cluster_spec, self._task_type) - else: - input_pipeline_id = 0 - num_input_pipelines = 1 - input_context = distribute_lib.InputContext( - num_input_pipelines=num_input_pipelines, - input_pipeline_id=input_pipeline_id, - num_replicas_in_sync=self._num_replicas_in_sync) - return input_lib.InputFunctionIterator( - input_fn, self._input_workers, [input_context]) - - def _broadcast_to(self, tensor, destinations): - # This is both a fast path for Python constants, and a way to delay - # converting Python values to a tensor until we know what type it - # should be converted to. Otherwise we have trouble with: - # global_step.assign_add(1) - # since the `1` gets broadcast as an int32 but global_step is int64. - if isinstance(tensor, (float, int)): - return tensor - if not cross_device_ops_lib.check_destinations(destinations): - # TODO(josh11b): Use current logical device instead of 0 here. - destinations = values.LogicalDeviceSpec( - device_map=self._device_map, logical_device=0) - return self._cross_device_ops.broadcast(tensor, destinations) - - def _allow_variable_partition(self): - return not context.executing_eagerly() - - # TODO(yuefengz): not all ops in device_setter.STANDARD_PS_OPS will go through - # this creator, such as "MutableHashTable". - def _create_variable(self, next_creator, *args, **kwargs): - if self._num_replicas_in_sync > 1: - 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: " + aggregation + - " for variable: " + kwargs["name"]) - - def var_creator(*args, **kwargs): - """Create an AggregatingVariable and fix up collections.""" - # Record what collections this variable should be added to. - collections = kwargs.pop("collections", None) - if collections is None: - collections = [ops.GraphKeys.GLOBAL_VARIABLES] - kwargs["collections"] = [] - - # Create and wrap the variable. - v = next_creator(*args, **kwargs) - wrapped = values.AggregatingVariable( - self._container_strategy(), v, aggregation) - - # Add the wrapped variable to the requested collections. - # The handling of eager mode and the global step matches - # ResourceVariable._init_from_args(). - if not context.executing_eagerly(): - g = ops.get_default_graph() - # If "trainable" is True, next_creator() will add the contained - # variable to the TRAINABLE_VARIABLES collection, so we manually - # remove it and replace with the wrapper. 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): - collections.append(ops.GraphKeys.TRAINABLE_VARIABLES) - l = g.get_collection_ref(ops.GraphKeys.TRAINABLE_VARIABLES) - l.remove(v) - g.add_to_collections(collections, wrapped) - elif ops.GraphKeys.GLOBAL_STEP in collections: - ops.add_to_collections(ops.GraphKeys.GLOBAL_STEP, wrapped) - - return wrapped - else: - var_creator = next_creator - - if "colocate_with" in kwargs: - with ops.device(None): - with ops.colocate_with(kwargs["colocate_with"]): - return var_creator(*args, **kwargs) - - with ops.colocate_with(None, ignore_existing=True): - with ops.device(self._variable_device): - return var_creator(*args, **kwargs) - - def _call_for_each_replica(self, fn, args, kwargs): - # pylint: disable=protected-access - return mirrored_strategy._call_for_each_replica( - self._container_strategy(), self._device_map, fn, args, kwargs) + return super(ParameterServerStrategy, self).make_dataset_iterator(dataset) - def _verify_destinations_not_different_worker(self, destinations): - if not self._cluster_spec: - return - if destinations is None: - return - for d in cross_device_ops_lib.get_devices_from(destinations): - d_spec = tf_device.DeviceSpec.from_string(d) - if d_spec.job == self._task_type and d_spec.task != self._task_id: - raise ValueError( - "Cannot reduce to another worker: %r, current worker is %r" % - (d, self._input_workers.worker_devices[0])) + # Override to change the documentation to reflect the different handling of + # global vs. local batch size between core and contrib. + def experimental_make_numpy_iterator( # pylint: disable=useless-super-delegation + self, numpy_input, batch_size, num_epochs=1, shuffle=1024, session=None): + """Makes an iterator for input provided via a nest of numpy arrays. - def _reduce_to(self, reduce_op, value, destinations): - self._verify_destinations_not_different_worker(destinations) - if not isinstance(value, values.DistributedValues): - # pylint: disable=protected-access - return cross_device_ops_lib.reduce_non_distributed_value( - reduce_op, self._device_map, value, destinations) - return self._cross_device_ops.reduce( - reduce_op, value, destinations=destinations) - - def _batch_reduce_to(self, reduce_op, value_destination_pairs): - for _, destinations in value_destination_pairs: - self._verify_destinations_not_different_worker(destinations) - return self._cross_device_ops.batch_reduce(reduce_op, - value_destination_pairs) - - def _select_single_value(self, structured): - """Select any single values in `structured`.""" - - def _select_fn(x): # pylint: disable=g-missing-docstring - if isinstance(x, values.Mirrored): - if len(x.devices) == 1: - return x.primary - else: - raise ValueError( - "You cannot update variable with a Mirrored object with multiple " - "components %r when using ParameterServerStrategy. You must " - "specify a single value or a Mirrored with a single value." % x) - elif isinstance(x, values.PerReplica): - raise ValueError( - "You cannot update variable with a PerReplica object %r when using " - "ParameterServerStrategy. You must specify a single value or a " - "Mirrored with a single value" % x) - else: - return x - - return nest.map_structure(_select_fn, structured) - - def _update(self, var, fn, args, kwargs, group): - if isinstance(var, values.AggregatingVariable): - var = var.get() - if not isinstance(var, resource_variable_ops.ResourceVariable): - raise ValueError( - "You can not update `var` %r. It must be a Variable." % var) - with ops.colocate_with(var), distribute_lib.UpdateContext(var.device): - result = fn(var, *self._select_single_value(args), - **self._select_single_value(kwargs)) - if group: - return result - else: - return nest.map_structure(self._unwrap, result) - - # TODO(yuefengz): does it need to call _select_single_value? - def _update_non_slot(self, colocate_with, fn, args, kwargs, group): - with ops.device( - colocate_with.device), distribute_lib.UpdateContext(colocate_with): - result = fn(*args, **kwargs) - if group: - return result - else: - return nest.map_structure(self._unwrap, result) - - def _unwrap(self, val): - if isinstance(val, values.DistributedValues): - return val.values - return (val,) - - def value_container(self, val): - if (hasattr(val, "_aggregating_container") and - not isinstance(val, values.AggregatingVariable)): - wrapper = val._aggregating_container() # pylint: disable=protected-access - if wrapper is not None: - return wrapper - return val - - def read_var(self, var): - # No need to distinguish between normal variables and replica-local - # variables. - return array_ops.identity(var) - - def _configure(self, - session_config=None, - cluster_spec=None, - task_type=None, - task_id=None): - """Configures the strategy class. - - The strategy object will be re-initialized if `cluster_spec` is given but - was not passed in the constructor. + NOTE: The `batch_size` argument here has different behavior for this + contrib version of `ParameterServerStrategy`. Args: - session_config: not used currently. - cluster_spec: a dict, ClusterDef or ClusterSpec object specifying the - cluster configurations. - task_type: the current task type. - task_id: the current task id. - - Raises: - ValueError: if `cluster_spec` is given but `task_type` or `task_id` is - not. + numpy_input: A nest of NumPy input arrays that will be distributed evenly + across all replicas. + batch_size: The number of entries from the array we should consume in one + step of the computation, across all replicas. This is the per-replica + batch size. The global batch size will be this times + `num_replicas_in_sync`. + num_epochs: The number of times to iterate through the examples. A value + of `None` means repeat forever. + shuffle: Size of buffer to use for shuffling the input examples. + Use `None` to disable shuffling. + session: (TensorFlow v1.x graph execution only) A session used for + initialization. + + Returns: + An `tf.distribute.InputIterator` which returns inputs for each step of the + computation. User should call `initialize` on the returned iterator. """ - if not self._cluster_spec and cluster_spec: - # If a `cluster_spec` is already passed in, do nothing here. - # TODO(yuefengz): check `cluster_spec` is the same if this object has - # already been initialized with a `cluster_spec`. - if task_type is None or task_id is None: - raise ValueError("When `cluster_spec` is given, must also specify " - "`task_type` and `task_id`.") - self._cluster_spec = multi_worker_util.normalize_cluster_spec( - cluster_spec) - self._task_type = task_type - self._task_id = task_id - self._initialize_multi_worker(self._num_gpus_per_worker, - self._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) - if not self._cluster_spec: - updated_config.isolate_session_state = True - return updated_config - - updated_config.isolate_session_state = False + return super(ParameterServerStrategy, + self).experimental_make_numpy_iterator( + numpy_input, batch_size, num_epochs, shuffle, session) - assert self._task_type - assert self._task_id is not None - # The device filters prevent communication between workers. - if self._task_type not in ["chief", "worker"]: - return updated_config - del updated_config.device_filters[:] - updated_config.device_filters.extend( - ["/job:%s/task:%d" % (self._task_type, self._task_id), "/job:ps"]) - return updated_config - - @property - def _num_replicas_in_sync(self): - return self._device_map.num_replicas_in_graph - - @property - def worker_devices(self): - return self._device_map.all_devices - - @property - def worker_devices_by_replica(self): - return self._device_map.devices_by_replica - - @property - def parameter_devices(self): - return self._parameter_devices - - def non_slot_devices(self, var_list): - return min(var_list, key=lambda x: x.name) - - @property - def experimental_between_graph(self): - # TODO(yuefengz): Should this return False in the local case? - return True - - @property - def experimental_should_init(self): - return self._is_chief +class ParameterServerExtended(CoreParameterServerExtended): + """Implementation of ParameterServerStrategy.""" - @property - def should_checkpoint(self): - return self._is_chief + def __init__(self, container_strategy, num_gpus_per_worker): + # 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. + tfconfig = TFConfigClusterResolver() + cluster_resolver = SimpleClusterResolver( + cluster_spec=tfconfig.cluster_spec(), + task_type=tfconfig.task_type, + task_id=tfconfig.task_id, + num_accelerators=num_gpus_per_worker) + super(ParameterServerExtended, self).__init__( + container_strategy, cluster_resolver=cluster_resolver) - @property - def should_save_summary(self): - return self._is_chief + def _make_dataset_iterator(self, dataset): + return input_lib.DatasetIterator(dataset, self._input_workers) # TODO(priyag): Delete this once all strategies use global batch size. @property def _global_batch_size(self): + """The contrib version of PS strategy uses per-replica batch size.""" return False diff --git a/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py b/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py index ce7065f220541ce2437f4768c312365a35197ab4..fede253d13804087476fef8b7211a6bfe5789906 100644 --- a/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py +++ b/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py @@ -29,10 +29,13 @@ 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 device_util +from tensorflow.python.distribute import distribute_lib from tensorflow.python.distribute import distribution_strategy_context as ds_context from tensorflow.python.distribute import multi_worker_util +from tensorflow.python.distribute import parameter_server_strategy as core_parameter_server_strategy from tensorflow.python.distribute import reduce_util from tensorflow.python.distribute import values +from tensorflow.python.distribute.cluster_resolver import SimpleClusterResolver from tensorflow.python.eager import backprop from tensorflow.python.eager import context from tensorflow.python.estimator import run_config @@ -50,6 +53,7 @@ from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables from tensorflow.python.platform import test from tensorflow.python.training import training_util +from tensorflow.python.training.server_lib import ClusterSpec CHIEF = run_config.TaskType.CHIEF WORKER = run_config.TaskType.WORKER @@ -63,6 +67,57 @@ def _get_replica_id_integer(): return replica_id +class MockCoreParameterServerStrategy(distribute_lib.DistributionStrategy): + """Mock the strategy to allow cluster resolver as an argument.""" + + def __init__(self, cluster_resolver): + super(MockCoreParameterServerStrategy, self).__init__( + core_parameter_server_strategy.ParameterServerStrategyExtended( + self, cluster_resolver=cluster_resolver)) + + +def create_test_objects(cluster_spec=None, + task_type=None, + task_id=None, + num_gpus=None, + sess_config=None, + use_core_strategy=False): + sess_config = sess_config or config_pb2.ConfigProto() + if num_gpus is None: + num_gpus = context.num_gpus() + if use_core_strategy: + if cluster_spec and task_type and task_id is not None: + cluster_resolver = SimpleClusterResolver( + cluster_spec=multi_worker_util.normalize_cluster_spec(cluster_spec), + task_type=task_type, + task_id=task_id, + num_accelerators=num_gpus) + target = 'grpc://' + cluster_spec[WORKER][task_id] + else: + cluster_resolver = SimpleClusterResolver( + ClusterSpec({}), num_accelerators=num_gpus) + target = '' + + distribution = MockCoreParameterServerStrategy(cluster_resolver) + sess_config = copy.deepcopy(sess_config) + sess_config = distribution.update_config_proto(sess_config) + else: + distribution = parameter_server_strategy.ParameterServerStrategy( + num_gpus_per_worker=num_gpus) + if task_type: + sess_config = copy.deepcopy(sess_config) + distribution.configure( + session_config=sess_config, + cluster_spec=cluster_spec, + task_type=task_type, + task_id=task_id) + target = 'grpc://' + cluster_spec[WORKER][task_id] + else: + target = '' + + return distribution, target, sess_config + + class ParameterServerStrategyTestBase( multi_worker_test_base.MultiWorkerTestBase): @@ -76,24 +131,27 @@ class ParameterServerStrategyTestBase( self._sess_config = config_pb2.ConfigProto(allow_soft_placement=True) super(ParameterServerStrategyTestBase, self).setUp() - def _get_test_objects(self, task_type, task_id, num_gpus): - distribution = parameter_server_strategy.ParameterServerStrategy( - num_gpus_per_worker=num_gpus) - if not task_type: - return distribution, '', self._sess_config - - sess_config = copy.deepcopy(self._sess_config) - distribution.configure( - session_config=sess_config, + def _get_test_objects(self, + task_type, + task_id, + num_gpus, + use_core_strategy=False): + return create_test_objects( cluster_spec=self._cluster_spec, task_type=task_type, - task_id=task_id) - return (distribution, 'grpc://' + self._cluster_spec[WORKER][task_id], - sess_config) - - def _test_device_assignment_distributed(self, task_type, task_id, num_gpus): + task_id=task_id, + num_gpus=num_gpus, + sess_config=self._sess_config, + use_core_strategy=use_core_strategy) + + def _test_device_assignment_distributed(self, + task_type, + task_id, + num_gpus, + use_core_strategy=False): worker_device = '/job:%s/replica:0/task:%d' % (task_type, task_id) - d, _, sess_config = self._get_test_objects(task_type, task_id, num_gpus) + d, _, sess_config = self._get_test_objects( + task_type, task_id, num_gpus, use_core_strategy=use_core_strategy) with ops.Graph().as_default(), \ self.cached_session(target=self._default_target, config=sess_config) as sess, \ @@ -132,7 +190,7 @@ class ParameterServerStrategyTestBase( '/job:worker/replica:0/task:0/%s' % last_part_device) # The colocate_vars_with can override the distribution's device. - with d.colocate_vars_with(x): + with d.extended.colocate_vars_with(x): y = variable_scope.get_variable( 'y', initializer=20.0, aggregation=variable_scope.VariableAggregation.SUM) @@ -178,7 +236,7 @@ class ParameterServerStrategyTestBase( self.assertIn('/job:ps/', h.device) return y_add, z_add, f - y, z, f = d.call_for_each_replica(model_fn) + y, z, f = d.extended.call_for_each_replica(model_fn) self.assertNotEqual(y, None) self.assertNotEqual(z, None) self.assertNotEqual(f, None) @@ -191,9 +249,10 @@ class ParameterServerStrategyTestBase( self.assertEqual(f_val, 46.0) def _test_device_assignment_distributed_enable_partitioner( - self, task_type, task_id, num_gpus): - d, _, sess_config = self._get_test_objects(task_type, task_id, num_gpus) - num_shards = len(d.parameter_devices) + self, task_type, task_id, num_gpus, use_core_strategy=False): + d, _, sess_config = self._get_test_objects( + task_type, task_id, num_gpus, use_core_strategy=use_core_strategy) + num_shards = len(d.extended.parameter_devices) partitioner = partitioned_variables.fixed_size_partitioner(num_shards) with ops.Graph().as_default(), \ self.cached_session(target=self._default_target, @@ -227,7 +286,7 @@ class ParameterServerStrategyTestBase( return x_add - x = d.call_for_each_replica(model_fn) + x = d.extended.call_for_each_replica(model_fn) if context.num_gpus() >= 1: variables.global_variables_initializer().run() @@ -285,7 +344,7 @@ class ParameterServerStrategyTestBase( self.assertEqual(e.device, device_util.canonicalize('/device:GPU:2')) # The colocate_vars_with can override the distribution's device. - with d.colocate_vars_with(x): + with d.extended.colocate_vars_with(x): y = variable_scope.get_variable( 'y', initializer=20.0, aggregation=variable_scope.VariableAggregation.SUM) @@ -328,7 +387,7 @@ class ParameterServerStrategyTestBase( device_util.canonicalize(h.device)) return y_add, z_add, f - y, z, f = d.call_for_each_replica(model_fn) + y, z, f = d.extended.call_for_each_replica(model_fn) self.assertNotEqual(y, None) self.assertNotEqual(z, None) self.assertNotEqual(f, None) @@ -340,9 +399,13 @@ class ParameterServerStrategyTestBase( self.assertEqual(z_val, 43.0) self.assertEqual(f_val, 46.0) - def _test_simple_increment(self, task_type, task_id, num_gpus): + def _test_simple_increment(self, + task_type, + task_id, + num_gpus, + use_core_strategy=False): d, master_target, sess_config = self._get_test_objects( - task_type, task_id, num_gpus) + task_type, task_id, num_gpus, use_core_strategy=use_core_strategy) if d.extended._cluster_spec: num_workers = len(d.extended._cluster_spec.as_dict().get(WORKER)) if 'chief' in d.extended._cluster_spec.as_dict(): @@ -375,7 +438,7 @@ class ParameterServerStrategyTestBase( train_op = control_flow_ops.group(x_add, y_add, z_add) return x, y, z, train_op - x, y, z, train_op = d.call_for_each_replica(model_fn) + x, y, z, train_op = d.extended.call_for_each_replica(model_fn) train_op = d.group(train_op) if context.num_gpus() < d.extended._num_gpus_per_worker: @@ -410,9 +473,13 @@ class ParameterServerStrategyTestBase( y_val == 20.0 + 1.0 * num_workers * d.num_replicas_in_sync and z_val == 30.0 + 1.0 * num_workers) - def _test_minimize_loss_graph(self, task_type, task_id, num_gpus): + def _test_minimize_loss_graph(self, + task_type, + task_id, + num_gpus, + use_core_strategy=False): d, master_target, sess_config = self._get_test_objects( - task_type, task_id, num_gpus) + task_type, task_id, num_gpus, use_core_strategy=use_core_strategy) if task_type: # Multi-worker assert hasattr(d.extended, '_cluster_spec') and d.extended._cluster_spec @@ -452,7 +519,7 @@ class ParameterServerStrategyTestBase( def step(): """Perform one optimization step.""" # Run forward & backward to get gradients, variables list. - g_v = d.call_for_each_replica(grad_fn, args=(one,)) + g_v = d.extended.call_for_each_replica(grad_fn, args=(one,)) # Update the variables using the gradients and the update() function. before_list = [] after_list = [] @@ -464,7 +531,7 @@ class ParameterServerStrategyTestBase( g = d.extended.reduce_to( reduce_util.ReduceOp.SUM, g, destinations=v) with ops.control_dependencies( - d.update(v, update, g, grouped=False)): + d.extended.update(v, update, args=(g,), group=False)): after_list.append(d.extended.read_var(v)) return before_list, after_list @@ -498,10 +565,16 @@ class ParameterServerStrategyTestBase( self.assertLess(error_after, error_before) return error_after < error_before - def _test_input_fn_iterator(self, task_type, task_id, num_gpus, input_fn, - expected_values): + def _test_input_fn_iterator(self, + task_type, + task_id, + num_gpus, + input_fn, + expected_values, + test_reinitialize=True, + use_core_strategy=False): distribution, master_target, config = self._get_test_objects( - task_type, task_id, num_gpus) + task_type, task_id, num_gpus, use_core_strategy=use_core_strategy) devices = distribution.extended.worker_devices with ops.Graph().as_default(), \ @@ -522,18 +595,21 @@ class ParameterServerStrategyTestBase( for r in range(len(devices))]) # After re-initializing the iterator, should be able to iterate again. - sess.run(iterator.initialize()) + if test_reinitialize: + sess.run(iterator.initialize()) - for expected_value in expected_values: - next_element = iterator.get_next() - computed_value = sess.run([values.select_replica(r, next_element) - for r in range(len(devices))]) - self.assertEqual(expected_value, computed_value) + for expected_value in expected_values: + next_element = iterator.get_next() + computed_value = sess.run([values.select_replica(r, next_element) + for r in range(len(devices))]) + self.assertEqual(expected_value, computed_value) -class ParameterServerStrategyTest(ParameterServerStrategyTestBase, - strategy_test_lib.DistributionTestBase, - parameterized.TestCase): +class ParameterServerStrategyTest( + ParameterServerStrategyTestBase, + strategy_test_lib.DistributionTestBase, + strategy_test_lib.TwoDeviceDistributionTestBase, + parameterized.TestCase): @classmethod def setUpClass(cls): @@ -541,111 +617,173 @@ class ParameterServerStrategyTest(ParameterServerStrategyTestBase, num_workers=3, num_ps=2) cls._default_target = 'grpc://' + cls._cluster_spec[WORKER][0] - def test_num_replicas_in_sync(self): - distribution = parameter_server_strategy.ParameterServerStrategy( - num_gpus_per_worker=2) + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def test_num_replicas_in_sync(self, use_core_strategy): + strategy, _, _ = create_test_objects( + num_gpus=2, use_core_strategy=use_core_strategy) # All the devices on a given worker are in sync which in this case is the # number of gpus on each worker. - self.assertEqual(2, distribution.num_replicas_in_sync) + self.assertEqual(2, strategy.num_replicas_in_sync) - def testDeviceAssignmentLocalCPU(self): - distribution = parameter_server_strategy.ParameterServerStrategy( - num_gpus_per_worker=0) + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testDeviceAssignmentLocalCPU(self, use_core_strategy): + strategy, _, _ = create_test_objects( + num_gpus=0, use_core_strategy=use_core_strategy) self._test_device_assignment_local( - distribution, compute_device='CPU', variable_device='CPU', num_gpus=0) + strategy, compute_device='CPU', variable_device='CPU', num_gpus=0) - def testDeviceAssignmentLocalOneGPU(self): - distribution = parameter_server_strategy.ParameterServerStrategy( - num_gpus_per_worker=1) + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testDeviceAssignmentLocalOneGPU(self, use_core_strategy): + strategy, _, _ = create_test_objects( + num_gpus=1, use_core_strategy=use_core_strategy) self._test_device_assignment_local( - distribution, compute_device='GPU', variable_device='GPU', num_gpus=1) + strategy, compute_device='GPU', variable_device='GPU', num_gpus=1) - def testDeviceAssignmentLocalTwoGPUs(self): - distribution = parameter_server_strategy.ParameterServerStrategy( - num_gpus_per_worker=2) + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testDeviceAssignmentLocalTwoGPUs(self, use_core_strategy): + strategy, _, _ = create_test_objects( + num_gpus=2, use_core_strategy=use_core_strategy) self._test_device_assignment_local( - distribution, compute_device='GPU', variable_device='CPU', num_gpus=2) + strategy, compute_device='GPU', variable_device='CPU', num_gpus=2) @combinations.generate( - combinations.combine(mode=['graph'], num_gpus=[0, 1, 2])) - def testDeviceAssignmentDistributed(self, num_gpus): - self._test_device_assignment_distributed('worker', 1, num_gpus) + combinations.combine( + mode=['graph'], num_gpus=[0, 1, 2], use_core_strategy=[True, False])) + def testDeviceAssignmentDistributed(self, num_gpus, use_core_strategy): + self._test_device_assignment_distributed( + 'worker', 1, num_gpus, use_core_strategy=use_core_strategy) @combinations.generate( - combinations.combine(mode=['graph'], num_gpus=[0, 1, 2])) - def testDeviceAssignmentDistributedEnablePartitioner(self, num_gpus): + combinations.combine( + mode=['graph'], num_gpus=[0, 1, 2], use_core_strategy=[True, False])) + def testDeviceAssignmentDistributedEnablePartitioner(self, num_gpus, + use_core_strategy): self._test_device_assignment_distributed_enable_partitioner( - 'worker', 1, num_gpus) + 'worker', 1, num_gpus, use_core_strategy=use_core_strategy) - def testSimpleBetweenGraph(self): - self._run_between_graph_clients(self._test_simple_increment, - self._cluster_spec, context.num_gpus()) + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testSimpleBetweenGraph(self, use_core_strategy): + self._run_between_graph_clients( + self._test_simple_increment, + self._cluster_spec, + context.num_gpus(), + use_core_strategy=use_core_strategy) @combinations.generate( - combinations.combine(mode=['graph'], num_gpus=[0, 1, 2])) - def testLocalSimpleIncrement(self, num_gpus): - self._test_simple_increment(None, 0, num_gpus) + combinations.combine( + mode=['graph'], num_gpus=[0, 1, 2], use_core_strategy=[True, False])) + def testLocalSimpleIncrement(self, num_gpus, use_core_strategy): + self._test_simple_increment(None, 0, num_gpus, use_core_strategy) @combinations.generate( - combinations.combine(mode=['graph'], num_gpus=[0, 1, 2])) - def testMinimizeLossGraphDistributed(self, num_gpus): - self._run_between_graph_clients(self._test_minimize_loss_graph, - self._cluster_spec, num_gpus) + combinations.combine( + mode=['graph'], num_gpus=[0, 1, 2], use_core_strategy=[True, False])) + def testMinimizeLossGraphDistributed(self, num_gpus, use_core_strategy): + self._run_between_graph_clients( + self._test_minimize_loss_graph, + self._cluster_spec, + num_gpus, + use_core_strategy=use_core_strategy) @combinations.generate( - combinations.combine(mode=['graph'], num_gpus=[0, 1, 2])) - def testMinimizeLossGraphLocal(self, num_gpus): - self._test_minimize_loss_graph(None, None, num_gpus) + combinations.combine( + mode=['graph'], num_gpus=[0, 1, 2], use_core_strategy=[True, False])) + def testMinimizeLossGraphLocal(self, num_gpus, use_core_strategy): + self._test_minimize_loss_graph(None, None, num_gpus, use_core_strategy) # TODO(priyag): Refactor this and other multi worker tests. @combinations.generate( - combinations.combine(mode=['graph'], num_gpus=[1, 2], required_gpus=1)) - def testMakeInputFnIteratorDistributed(self, num_gpus): + combinations.combine( + mode=['graph'], + num_gpus=[1, 2], + required_gpus=1, + use_core_strategy=[True, False], + use_dataset=[True, False])) + def testMakeInputFnIteratorDistributed(self, num_gpus, use_core_strategy, + use_dataset): if context.num_gpus() < num_gpus: self.skipTest('Not enough GPUs') - dataset_fn = lambda: dataset_ops.Dataset.range(100) + if use_dataset: + fn = lambda: dataset_ops.Dataset.range(100) + else: + def fn(): + dataset = dataset_ops.Dataset.range(100) + it = dataset.make_one_shot_iterator() + return it.get_next expected_values = [[i+j for j in range(num_gpus)] for i in range(0, 100, num_gpus)] input_fn = self._input_fn_to_test_input_context( - dataset_fn, + fn, expected_num_replicas_in_sync=num_gpus, expected_num_input_pipelines=3, expected_input_pipeline_id=1) # because task_id = 1 - self._test_input_fn_iterator('worker', 1, num_gpus, - input_fn, expected_values) + self._test_input_fn_iterator( + 'worker', + 1, + num_gpus, + input_fn, + expected_values, + test_reinitialize=use_dataset, + use_core_strategy=use_core_strategy) @combinations.generate( - combinations.combine(mode=['graph'], num_gpus=[1, 2], required_gpus=1)) - def testMakeInputFnIteratorLocal(self, num_gpus): + combinations.combine( + mode=['graph'], + num_gpus=[1, 2], + required_gpus=1, + use_core_strategy=[True, False], + use_dataset=[True, False])) + def testMakeInputFnIteratorLocal(self, num_gpus, use_core_strategy, + use_dataset): if context.num_gpus() < num_gpus: self.skipTest('Not enough GPUs') - dataset_fn = lambda: dataset_ops.Dataset.range(100) + if use_dataset: + fn = lambda: dataset_ops.Dataset.range(100) + else: + def fn(): + dataset = dataset_ops.Dataset.range(100) + it = dataset.make_one_shot_iterator() + return it.get_next expected_values = [[i+j for j in range(num_gpus)] for i in range(0, 100, num_gpus)] input_fn = self._input_fn_to_test_input_context( - dataset_fn, + fn, expected_num_replicas_in_sync=num_gpus, expected_num_input_pipelines=1, expected_input_pipeline_id=0) # only one worker and pipeline for local. - self._test_input_fn_iterator(None, None, num_gpus, - input_fn, expected_values) + self._test_input_fn_iterator( + None, + None, + num_gpus, + input_fn, + expected_values, + test_reinitialize=use_dataset, + use_core_strategy=use_core_strategy) - def testGlobalStepUpdate(self): - strategy = parameter_server_strategy.ParameterServerStrategy( - num_gpus_per_worker=context.num_gpus()) + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testGlobalStepUpdate(self, use_core_strategy): + strategy, _, _ = create_test_objects(use_core_strategy=use_core_strategy) self._test_global_step_update(strategy) - def testUpdateConfigProtoMultiWorker(self): - distribution = parameter_server_strategy.ParameterServerStrategy( - num_gpus_per_worker=2) - distribution.configure( + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testUpdateConfigProtoMultiWorker(self, use_core_strategy): + strategy, _, _ = create_test_objects( + num_gpus=2, use_core_strategy=use_core_strategy) + strategy.configure( cluster_spec=self._cluster_spec, task_type='worker', task_id=1) config_proto = config_pb2.ConfigProto(device_filters=['to_be_overridden']) - new_config = distribution.update_config_proto(config_proto) + new_config = strategy.update_config_proto(config_proto) # Verify device filters. self.assertEqual(['/job:worker/task:1', '/job:ps'], @@ -654,16 +792,48 @@ class ParameterServerStrategyTest(ParameterServerStrategyTestBase, # Verify isolate_session_state self.assertFalse(new_config.isolate_session_state) - def testUpdateConfigProtoLocal(self): - distribution = parameter_server_strategy.ParameterServerStrategy( - num_gpus_per_worker=2) + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testUpdateConfigProtoLocal(self, use_core_strategy): + strategy, _, _ = create_test_objects( + num_gpus=2, use_core_strategy=use_core_strategy) config_proto = config_pb2.ConfigProto() - new_config = distribution.update_config_proto(config_proto) + new_config = strategy.update_config_proto(config_proto) # Verify isolate_session_state self.assertTrue(new_config.isolate_session_state) + def testAllReduceSum(self): + distribution = parameter_server_strategy.ParameterServerStrategy( + num_gpus_per_worker=2) + self._test_all_reduce_sum(distribution) + + def testAllReduceSumGradients(self): + distribution = parameter_server_strategy.ParameterServerStrategy( + num_gpus_per_worker=2) + self._test_all_reduce_sum_gradients(distribution) + + def testAllReduceSumGradientTape(self): + distribution = parameter_server_strategy.ParameterServerStrategy( + num_gpus_per_worker=2) + self._test_all_reduce_sum_gradient_tape(distribution) + + def testAllReduceMean(self): + distribution = parameter_server_strategy.ParameterServerStrategy( + num_gpus_per_worker=2) + self._test_all_reduce_mean(distribution) + + def testAllReduceMeanGradients(self): + distribution = parameter_server_strategy.ParameterServerStrategy( + num_gpus_per_worker=2) + self._test_all_reduce_mean_gradients(distribution) + + def testAllReduceMeanGradientTape(self): + distribution = parameter_server_strategy.ParameterServerStrategy( + num_gpus_per_worker=2) + self._test_all_reduce_mean_gradient_tape(distribution) + class ParameterServerStrategyWithChiefTest(ParameterServerStrategyTestBase, parameterized.TestCase): @@ -674,20 +844,31 @@ class ParameterServerStrategyWithChiefTest(ParameterServerStrategyTestBase, num_workers=3, num_ps=2, has_chief=True) cls._default_target = 'grpc://' + cls._cluster_spec[CHIEF][0] - def testSimpleBetweenGraph(self): - self._run_between_graph_clients(self._test_simple_increment, - self._cluster_spec, context.num_gpus()) + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testSimpleBetweenGraph(self, use_core_strategy): + self._run_between_graph_clients( + self._test_simple_increment, + self._cluster_spec, + context.num_gpus(), + use_core_strategy=use_core_strategy) @combinations.generate( - combinations.combine(mode=['graph'], num_gpus=[0, 1, 2])) - def testMinimizeLossGraph(self, num_gpus): - self._run_between_graph_clients(self._test_minimize_loss_graph, - self._cluster_spec, num_gpus) + combinations.combine( + mode=['graph'], num_gpus=[0, 1, 2], use_core_strategy=[True, False])) + def testMinimizeLossGraph(self, num_gpus, use_core_strategy): + self._run_between_graph_clients( + self._test_minimize_loss_graph, + self._cluster_spec, + num_gpus, + use_core_strategy=use_core_strategy) - def testGlobalStepIsWrappedOnTwoGPUs(self): - distribution = parameter_server_strategy.ParameterServerStrategy( - num_gpus_per_worker=2) - with ops.Graph().as_default(), distribution.scope(): + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testGlobalStepIsWrappedOnTwoGPUs(self, use_core_strategy): + strategy, _, _ = create_test_objects( + num_gpus=2, use_core_strategy=use_core_strategy) + with ops.Graph().as_default(), strategy.scope(): created_step = training_util.create_global_step() get_step = training_util.get_global_step() self.assertEqual(created_step, get_step, @@ -696,12 +877,14 @@ class ParameterServerStrategyWithChiefTest(ParameterServerStrategyTestBase, id(get_step), get_step.__class__.__name__))) self.assertIs(values.AggregatingVariable, type(created_step)) self.assertIs(values.AggregatingVariable, type(get_step)) - self.assertIs(distribution, created_step.distribute_strategy) + self.assertIs(strategy, created_step.distribute_strategy) - def testGlobalStepIsNotWrappedOnOneGPU(self): - distribution = parameter_server_strategy.ParameterServerStrategy( - num_gpus_per_worker=1) - with ops.Graph().as_default(), distribution.scope(): + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testGlobalStepIsNotWrappedOnOneGPU(self, use_core_strategy): + strategy, _, _ = create_test_objects( + num_gpus=1, use_core_strategy=use_core_strategy) + with ops.Graph().as_default(), strategy.scope(): created_step = training_util.create_global_step() get_step = training_util.get_global_step() self.assertEqual(created_step, get_step, @@ -710,20 +893,39 @@ class ParameterServerStrategyWithChiefTest(ParameterServerStrategyTestBase, id(get_step), get_step.__class__.__name__))) self.assertIs(resource_variable_ops.ResourceVariable, type(created_step)) self.assertIs(resource_variable_ops.ResourceVariable, type(get_step)) - self.assertIs(distribution, created_step.distribute_strategy) + # All variables have an _distribute_strategy parameter. Only variable + # subclasses in distribution strategy expose it publicly. + self.assertFalse(hasattr(strategy, 'distribute_strategy')) + self.assertIs(strategy, created_step._distribute_strategy) + + @combinations.generate( + combinations.combine(mode=['graph'], use_core_strategy=[True, False])) + def testValueContainer(self, use_core_strategy): + strategy, _, _ = create_test_objects( + num_gpus=2, use_core_strategy=use_core_strategy) + with ops.Graph().as_default(), strategy.scope(): - def testValueContainer(self): - distribution = parameter_server_strategy.ParameterServerStrategy( - num_gpus_per_worker=2) - with ops.Graph().as_default(), distribution.scope(): def f(): with backprop.GradientTape() as tape: v = variable_scope.get_variable('v', initializer=10.0) _ = v * v v, = tape.watched_variables() - w = distribution.extended.value_container(v) + w = strategy.extended.value_container(v) self.assertIs(values.AggregatingVariable, type(w)) - distribution.extended.call_for_each_replica(f) + + strategy.extended.call_for_each_replica(f) + + +class LocalParameterServerStrategyTest(strategy_test_lib.DistributionTestBase, + parameterized.TestCase): + + @combinations.generate(combinations.combine(mode=['graph', 'eager'], + use_core_strategy=[True, False], + required_gpus=2)) + def testNumpyIterator(self, use_core_strategy): + strategy, _, _ = create_test_objects( + num_gpus=2, use_core_strategy=use_core_strategy) + self._test_numpy_iterator(strategy) if __name__ == '__main__': diff --git a/tensorflow/contrib/distribute/python/step_fn.py b/tensorflow/contrib/distribute/python/step_fn.py index faeb96bcb7c516b1e494661ef2cbe8dad476ab55..27aad46b97195aa498d0382f08c04c312cebbe65 100644 --- a/tensorflow/contrib/distribute/python/step_fn.py +++ b/tensorflow/contrib/distribute/python/step_fn.py @@ -19,7 +19,6 @@ from __future__ import division from __future__ import print_function from tensorflow.python.eager import backprop -from tensorflow.python.eager import context from tensorflow.python.training import optimizer as optimizer_lib @@ -33,6 +32,9 @@ class Step(object): def distribution(self): return self._distribution + def initialize(self): + return [] + def __call__(self): """Perform one step of this training algorithm.""" raise NotImplementedError("must be implemented in descendants") @@ -50,12 +52,10 @@ class StandardInputStep(Step): def __init__(self, dataset_fn, distribution): super(StandardInputStep, self).__init__(distribution) - self._distributed_input = distribution.distribute_dataset(dataset_fn) - if context.executing_eagerly(): - self._iterator = self._distributed_input.make_one_shot_iterator() - else: - # TODO(priyag): Expose initializer via some initializer property. - self._iterator = self._distributed_input.make_initializable_iterator() + self._iterator = distribution.make_input_fn_iterator(lambda _: dataset_fn()) + + def initialize(self): + return self._iterator.initialize() class StandardSingleLossStep(StandardInputStep): @@ -99,7 +99,7 @@ class StandardSingleLossStep(StandardInputStep): gradients_fn = backprop.implicit_grad(self._loss_fn) gradients_fn = optimizer_lib.get_filtered_grad_fn(gradients_fn) - grads_and_vars = self.distribution.call_for_each_replica( + grads_and_vars = self.distribution.extended.call_for_each_replica( gradients_fn, args=(ctx, inputs)) # If threads use layers, then we need to run the first step # sequentially, so that layers.build() is not executed in parallel. @@ -109,6 +109,6 @@ class StandardSingleLossStep(StandardInputStep): self.distribution, grads_and_vars) # TODO(priyag): Return the outputs, context, etc as well. - ctx = self.distribution.run_steps_on_dataset( + ctx = self.distribution.extended.experimental_run_steps_on_iterator( step_fn, self._iterator, self._iterations_per_step) return ctx.run_op diff --git a/tensorflow/contrib/distribute/python/step_fn_test.py b/tensorflow/contrib/distribute/python/step_fn_test.py index 1ff9b9ceec13351b098d47ed3ff62f689a625a31..9f48560b2666036e149a63c98b6529fb24cc5067 100644 --- a/tensorflow/contrib/distribute/python/step_fn_test.py +++ b/tensorflow/contrib/distribute/python/step_fn_test.py @@ -45,24 +45,21 @@ class SingleLossStepTest(test.TestCase, parameterized.TestCase): single_loss_step, layer = single_loss_example( optimizer_fn, distribution, use_bias=True, iterations_per_step=2) - self.evaluate(distribution.initialize()) if context.executing_eagerly(): + single_loss_step.initialize() run_step = single_loss_step else: with self.cached_session() as sess: - sess.run(single_loss_step._iterator.initializer) + sess.run(single_loss_step.initialize()) run_step = sess.make_callable(single_loss_step()) self.evaluate(variables.global_variables_initializer()) weights, biases = [], [] for _ in range(5): run_step() - weights.append(self.evaluate(layer.kernel)) biases.append(self.evaluate(layer.bias)) - self.evaluate(distribution.finalize()) - error = abs(numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1) is_not_increasing = all(y <= x for x, y in zip(error, error[1:])) self.assertTrue(is_not_increasing) diff --git a/tensorflow/contrib/distribute/python/strategy_test_lib.py b/tensorflow/contrib/distribute/python/strategy_test_lib.py index 6e5280e35632d3f3cb6a4fe172a15fb7f508354c..90f552eda4c41742f21ca276d8a059b2b102554f 100644 --- a/tensorflow/contrib/distribute/python/strategy_test_lib.py +++ b/tensorflow/contrib/distribute/python/strategy_test_lib.py @@ -18,7 +18,10 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import numpy as np + from tensorflow.core.protobuf import config_pb2 +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.distribute import distribution_strategy_context as ds_context from tensorflow.python.distribute import reduce_util from tensorflow.python.distribute import values @@ -31,6 +34,7 @@ from tensorflow.python.framework import errors from tensorflow.python.framework import ops from tensorflow.python.layers import core from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gradients_impl from tensorflow.python.ops import init_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables @@ -41,25 +45,26 @@ class _TestException(Exception): pass -# May be the argument to either distribution.call_for_each_replica() or +# May be the argument to either distribution.extended.call_for_each_replica() or # get_replica_context().merge_call() def _raise_exception_fn(_=None): raise _TestException() -# Must be the argument to a distribution.call_for_each_replica() call, calls a -# get_replica_context().merge_call() that raises an exception. +# Must be the argument to a distribution.extended.call_for_each_replica() call, +# calls a get_replica_context().merge_call() that raises an exception. def _merge_raises_fn(): ds_context.get_replica_context().merge_call(_raise_exception_fn) # Must be the argument to a get_replica_context().merge_call() call, calls -# dist.call_for_each_replica() with a function that raises an exception. +# dist.extended.call_for_each_replica() with a function that raises an +# exception. def _call_raises_fn(dist): - dist.call_for_each_replica(_raise_exception_fn) + dist.extended.call_for_each_replica(_raise_exception_fn) -# Must be the argument to a distribution.call_for_each_replica() call, +# Must be the argument to a distribution.extended.call_for_each_replica() call, # calls a get_replica_context().merge_call() that calls a # call_for_each_replica() that raises an exception. def _merge_call_raises_fn(): @@ -67,15 +72,16 @@ def _merge_call_raises_fn(): # Must be the argument to a get_replica_context().merge_call() call, calls -# dist.call_for_each_replica() with a function that calls a +# dist.extended.call_for_each_replica() with a function that calls a # get_replica_context().merge_call() that raises an exception. def _call_merge_raises_fn(dist): - dist.call_for_each_replica(_merge_raises_fn) + dist.extended.call_for_each_replica(_merge_raises_fn) -# Must be the argument to a distribution.call_for_each_replica() call, calls a -# get_replica_context().merge_call() that calls a call_for_each_replica() that -# calls a get_replica_context().merge_call() that raises an exception. +# Must be the argument to a distribution.extended.call_for_each_replica() call, +# calls a get_replica_context().merge_call() that calls a +# call_for_each_replica() that calls a get_replica_context().merge_call() that +# raises an exception. def _merge_call_merge_raises_fn(): ds_context.get_replica_context().merge_call(_call_merge_raises_fn) @@ -106,7 +112,7 @@ class DistributionTestBase(test.TestCase): def step(): """Perform one optimization step.""" # Run forward & backward to get gradients, variables list. - g_v = d.call_for_each_replica(grad_fn, args=(one,)) + g_v = d.extended.call_for_each_replica(grad_fn, args=(one,)) # Update the variables using the gradients and the update() function. before_list = [] @@ -118,8 +124,8 @@ class DistributionTestBase(test.TestCase): with ops.control_dependencies([fetched]): g = d.extended.reduce_to( reduce_util.ReduceOp.SUM, g, destinations=v) - with ops.control_dependencies(d.update( - v, update, g, grouped=False)): + with ops.control_dependencies(d.extended.update( + v, update, args=(g,), group=False)): after_list.append(d.extended.read_var(v)) return before_list, after_list @@ -162,7 +168,7 @@ class DistributionTestBase(test.TestCase): def step(): """Perform one optimization step.""" # Run forward & backward to get gradients, variables list. - g_v = d.call_for_each_replica(grad_fn, args=(one,)) + g_v = d.extended.call_for_each_replica(grad_fn, args=(one,)) # Update the variables using the gradients and the update() function. before_list = [] @@ -173,8 +179,8 @@ class DistributionTestBase(test.TestCase): with ops.control_dependencies([fetched]): g = d.extended.reduce_to( reduce_util.ReduceOp.SUM, g, destinations=v) - with ops.control_dependencies(d.update( - v, update, g, grouped=False)): + with ops.control_dependencies(d.extended.update( + v, update, args=(g,), group=False)): after_list.append(d.extended.read_var(v)) return before_list, after_list @@ -202,23 +208,23 @@ class DistributionTestBase(test.TestCase): self.assertFalse(expected_devices[replica_id]) expected_devices[replica_id] = True - d.call_for_each_replica(mark_devices_fn) + d.extended.call_for_each_replica(mark_devices_fn) self.assertAllEqual(expected_devices, [True] * len(d.extended.worker_devices)) def _test_call_and_merge_exceptions(self, dist): with dist.scope(): with self.assertRaises(_TestException): - dist.call_for_each_replica(_raise_exception_fn) + dist.extended.call_for_each_replica(_raise_exception_fn) with self.assertRaises(_TestException): - dist.call_for_each_replica(_merge_raises_fn) + dist.extended.call_for_each_replica(_merge_raises_fn) with self.assertRaises(_TestException): - dist.call_for_each_replica(_merge_call_raises_fn) + dist.extended.call_for_each_replica(_merge_call_raises_fn) with self.assertRaises(_TestException): - dist.call_for_each_replica(_merge_call_merge_raises_fn) + dist.extended.call_for_each_replica(_merge_call_merge_raises_fn) def _input_fn_to_test_input_context(self, - dataset_fn, + dataset_or_callable_fn, expected_num_replicas_in_sync, expected_num_input_pipelines, expected_input_pipeline_id): @@ -242,12 +248,12 @@ class DistributionTestBase(test.TestCase): self.assertEqual(worker_id_counter[0], input_context.input_pipeline_id) worker_id_counter[0] += 1 - return dataset_fn() + return dataset_or_callable_fn() return _input_fn def _test_input_fn_iterator(self, iterator, devices, expected_values, - sess=None): + sess=None, test_reinitialize=True): evaluate = lambda x: sess.run(x) if sess else self.evaluate(x) evaluate(iterator.initialize()) @@ -263,13 +269,14 @@ class DistributionTestBase(test.TestCase): [values.select_replica(r, next_element) for r in range(len(devices))]) # After re-initializing the iterator, should be able to iterate again. - evaluate(iterator.initialize()) + if test_reinitialize: + evaluate(iterator.initialize()) - for expected_value in expected_values: - next_element = iterator.get_next() - computed_value = evaluate( - [values.select_replica(r, next_element) for r in range(len(devices))]) - self.assertEqual(expected_value, computed_value) + for expected_value in expected_values: + next_element = iterator.get_next() + computed_value = evaluate([values.select_replica(r, next_element) + for r in range(len(devices))]) + self.assertEqual(expected_value, computed_value) def _test_global_step_update(self, strategy): with strategy.scope(): @@ -287,8 +294,195 @@ class DistributionTestBase(test.TestCase): value = global_step.read_value() return train_op, value - train_ops, value = strategy.call_for_each_replica(model_fn) + train_ops, value = strategy.extended.call_for_each_replica(model_fn) self.evaluate(strategy.group(train_ops)) global_step_tensors = strategy.unwrap(value) global_step_values = self.evaluate(global_step_tensors) self.assertEqual((1,) * len(global_step_tensors), global_step_values) + + def _test_numpy_iterator(self, strategy): + with strategy.scope(), self.cached_session() as sess: + x = np.asarray([[1, 2], [6, 12], [2, 4], + [5, 10], [3, 6], [4, 8]]) + y = np.asarray([5, 4, 3, 2, 1, 0]) + batch_size = 6 + if not strategy.extended._global_batch_size: # pylint: disable=protected-access + batch_size = batch_size // strategy.num_replicas_in_sync + i = strategy.experimental_make_numpy_iterator( + (x, y), batch_size=batch_size, num_epochs=2, shuffle=None, + session=sess) + self.evaluate(i.initialize()) + + def run_and_concatenate(strategy, i): + x, y = strategy.experimental_run(lambda z: z, i) + x, y = self.evaluate((strategy.unwrap(x), strategy.unwrap(y))) + return np.concatenate(x), np.concatenate(y) + + x_1, y_1 = run_and_concatenate(strategy, i) + self.assertAllEqual(x, x_1) + self.assertAllEqual(y, y_1) + x_2, y_2 = run_and_concatenate(strategy, i) + self.assertAllEqual(x, x_2) + self.assertAllEqual(y, y_2) + with self.assertRaises(errors.OutOfRangeError): + run_and_concatenate(strategy, i) + + +class OneDeviceDistributionTestBase(test.TestCase): + """Some tests that should work with any one-device DistributionStrategy.""" + + def _test_all_reduce_sum(self, strategy): + self._test_collective_comms( + strategy, _all_sum, inputs=(4., [42., 43.]), expected=(4., [42., 43.])) + + def _test_all_reduce_sum_gradients(self, strategy): + self._test_collective_comms_gradients( + strategy, _all_sum, inputs=[4.], expected_grads=[4.]) + + def _test_all_reduce_sum_gradient_tape(self, strategy): + self._test_collective_comms_gradient_tape( + strategy, _all_sum, inputs=[4.], expected_grads=[4.]) + + def _test_all_reduce_mean(self, strategy): + self._test_collective_comms( + strategy, _all_mean, inputs=(2., [21., 22.]), expected=(2., [21., 22.])) + + def _test_all_reduce_mean_gradients(self, strategy): + self._test_collective_comms_gradients( + strategy, _all_mean, inputs=[5.], expected_grads=[5.]) + + def _test_all_reduce_mean_gradient_tape(self, strategy): + self._test_collective_comms_gradient_tape( + strategy, _all_mean, inputs=[5.], expected_grads=[5.]) + + def _test_collective_comms(self, strategy, comm_fn, inputs, expected): + inputs = strategy.make_input_fn_iterator( + lambda _: dataset_ops.Dataset.from_tensors(inputs)) + + self.evaluate(inputs.initialize()) + outputs = self.evaluate( + list(map(strategy.unwrap, strategy.experimental_run(comm_fn, inputs)))) + self.assertAllEqual([expected[0]], outputs[0]) + self.assertAllEqual([expected[1]], outputs[1]) + + def _test_collective_comms_gradients( + self, strategy, comm_fn, inputs, expected_grads): + if context.executing_eagerly(): + self.skipTest("`tf.gradients` is not supported with eager execution.") + + def step(c): + x = constant_op.constant(42.) + y = comm_fn(x) * c + return gradients_impl.gradients(y, [x])[0] + + inputs = strategy.make_input_fn_iterator( + lambda _: dataset_ops.Dataset.from_tensors(inputs)) + + self.evaluate(inputs.initialize()) + self.assertAllEqual( + expected_grads, + self.evaluate(strategy.unwrap(strategy.experimental_run(step, inputs)))) + + def _test_collective_comms_gradient_tape( + self, strategy, comm_fn, inputs, expected_grads): + def step(c): + x = constant_op.constant(42.) + with backprop.GradientTape() as tape: + tape.watch(x) + y = comm_fn(x) * c + return tape.gradient(y, x) + + inputs = strategy.make_input_fn_iterator( + lambda _: dataset_ops.Dataset.from_tensors(inputs)) + + self.evaluate(inputs.initialize()) + self.assertAllEqual( + expected_grads, + self.evaluate(strategy.unwrap(strategy.experimental_run(step, inputs)))) + + +class TwoDeviceDistributionTestBase(test.TestCase): + """Some tests that should work with any two-device DistributionStrategy.""" + + def _test_all_reduce_sum(self, strategy): + self._test_collective_comms( + strategy, _all_sum, + inputs=([1., 3.], [[39., 2.], [3., 41.]]), + expected=(4., [42., 43.])) + + def _test_all_reduce_sum_gradients(self, strategy): + self._test_collective_comms_gradients( + strategy, _all_sum, inputs=[1., 3.], expected_grads=[4., 4.]) + + def _test_all_reduce_sum_gradient_tape(self, strategy): + self._test_collective_comms_gradient_tape( + strategy, _all_sum, inputs=[1., 3.], expected_grads=[4., 4.]) + + def _test_all_reduce_mean(self, strategy): + self._test_collective_comms( + strategy, _all_mean, + inputs=([1., 3.], [[39., 2.], [3., 41.]]), + expected=(2., [21., 21.5])) + + def _test_all_reduce_mean_gradients(self, strategy): + self._test_collective_comms_gradients( + strategy, _all_mean, inputs=[1., 3.], expected_grads=[2., 2.]) + + def _test_all_reduce_mean_gradient_tape(self, strategy): + self._test_collective_comms_gradient_tape( + strategy, _all_mean, inputs=[1., 3.], expected_grads=[2., 2.]) + + def _test_collective_comms(self, strategy, comm_fn, inputs, expected): + inputs = strategy.make_input_fn_iterator( + lambda _: dataset_ops.Dataset.from_tensor_slices(inputs)) + + self.evaluate(inputs.initialize()) + outputs = self.evaluate( + list(map(strategy.unwrap, strategy.experimental_run(comm_fn, inputs)))) + self.assertAllEqual([expected[0], expected[0]], outputs[0]) + self.assertAllEqual([expected[1], expected[1]], outputs[1]) + + def _test_collective_comms_gradients( + self, strategy, comm_fn, inputs, expected_grads): + if context.executing_eagerly(): + self.skipTest("`tf.gradients` is not supported with eager execution.") + + def step(c): + x = constant_op.constant(42.) + y = comm_fn(x) * c + return gradients_impl.gradients(y, [x])[0] + + inputs = strategy.make_input_fn_iterator( + lambda _: dataset_ops.Dataset.from_tensor_slices(inputs)) + + self.evaluate(inputs.initialize()) + self.assertAllEqual( + expected_grads, + self.evaluate(strategy.unwrap(strategy.experimental_run(step, inputs)))) + + def _test_collective_comms_gradient_tape( + self, strategy, comm_fn, inputs, expected_grads): + def step(c): + x = constant_op.constant(42.) + with backprop.GradientTape() as tape: + tape.watch(x) + y = comm_fn(x) * c + return tape.gradient(y, x) + + inputs = strategy.make_input_fn_iterator( + lambda _: dataset_ops.Dataset.from_tensor_slices(inputs)) + + self.evaluate(inputs.initialize()) + self.assertAllEqual( + expected_grads, + self.evaluate(strategy.unwrap(strategy.experimental_run(step, inputs)))) + + +def _all_sum(value): + ctx = ds_context.get_replica_context() + return ctx.all_reduce(reduce_util.ReduceOp.SUM, value) + + +def _all_mean(value): + ctx = ds_context.get_replica_context() + return ctx.all_reduce(reduce_util.ReduceOp.MEAN, value) diff --git a/tensorflow/contrib/distribute/python/tpu_strategy.py b/tensorflow/contrib/distribute/python/tpu_strategy.py index 9cb2e461ab7b47abd90e5326bfe9d250af6ff4fe..2d9d221f427422f8bbeba55c5644658af9a9a620 100644 --- a/tensorflow/contrib/distribute/python/tpu_strategy.py +++ b/tensorflow/contrib/distribute/python/tpu_strategy.py @@ -21,10 +21,13 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import collections import copy -import functools 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 @@ -34,12 +37,15 @@ from tensorflow.python.distribute import cross_device_ops as cross_device_ops_li 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 tpu_cluster_resolver as resolver_lib +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 @@ -48,30 +54,56 @@ 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 -_TPU_INITIALIZE_SYSTEM_COLLECTION = "TPU_STRATEGY_INITIALIZE" - - 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 = resolver_lib.TPUClusterResolver("") + cluster_resolver = TPUClusterResolver("") master = cluster_resolver.master() logging.info("Initializing the TPU system.") - 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: - sess.run([tpu.initialize_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): @@ -96,9 +128,9 @@ def _create_tpu_mirrored_variable( # pylint: disable=missing-docstring *args, **kwargs): # Figure out what collections this variable should be added to. # We'll add the TPUMirroredVariable to those collections instead. - collections = kwargs.pop("collections", None) - if collections is None: - collections = [ops.GraphKeys.GLOBAL_VARIABLES] + 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 @@ -129,7 +161,7 @@ def _create_tpu_mirrored_variable( # pylint: disable=missing-docstring strategy, device_map, value_list, aggregation, logical_device=logical_device) - if not context.executing_eagerly(): + 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 @@ -137,11 +169,11 @@ def _create_tpu_mirrored_variable( # pylint: disable=missing-docstring # "trainable" to False for next_creator() since that causes functions # like implicit_gradients to skip those variables. if kwargs.get("trainable", True): - collections.append(ops.GraphKeys.TRAINABLE_VARIABLES) + 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(collections, result) + g.add_to_collections(var_collections, result) return result @@ -151,7 +183,7 @@ class TPUStrategy(distribute_lib.DistributionStrategy): def __init__(self, tpu_cluster_resolver=None, steps_per_run=None, - **kwargs): + device_assignment=None): """Initializes the TPUStrategy object. Args: @@ -162,27 +194,59 @@ class TPUStrategy(distribute_lib.DistributionStrategy): metrics, summaries etc. This parameter is only used when Distribution Strategy is used with estimator or keras. - **kwargs: Additional experimental flags. Will be removed in future. + 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)) - - self._disable_training_loop_on_host = False - if len(kwargs) > 1: - raise ValueError("TPUStrategy constructor only takes one experimental " - "flag now") - if len(kwargs) == 1: - if "_disable_training_loop_on_host" not in kwargs: - raise ValueError("TPUStrategy constructor does not support arguments: " - "{}".format(kwargs)) - self._disable_training_loop_on_host = ( - kwargs["_disable_training_loop_on_host"]) + 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.""" @@ -190,11 +254,12 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): def __init__(self, container_strategy, tpu_cluster_resolver=None, - steps_per_run=None): + steps_per_run=None, + device_assignment=None): super(TPUExtended, self).__init__(container_strategy) if tpu_cluster_resolver is None: - tpu_cluster_resolver = resolver_lib.TPUClusterResolver("") + tpu_cluster_resolver = TPUClusterResolver("") if steps_per_run is None: # TODO(frankchn): Warn when we are being used by DS/Keras and this is @@ -203,6 +268,21 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): 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. @@ -216,15 +296,14 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): self._tpu_devices = self._tpu_devices[:self._num_replicas_in_sync] self._device_map = values.ReplicaDeviceMap(self._tpu_devices) - # For input: - input_device_map = values.ReplicaDeviceMap(tuple( - self.get_host_cpu_device(hid) for hid in range(self.num_hosts))) - worker_devices = [ - (self.get_host(hid), [self.get_host_cpu_device(hid)]) - for hid in range(self.num_hosts) - ] + # 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( - input_device_map, worker_devices) + 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. @@ -297,14 +376,27 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): 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 _distribute_dataset(self, dataset_fn): - return input_lib.MultiWorkerDataset( - functools.partial(self._call_dataset_fn, dataset_fn), - self._input_workers) + 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 @@ -318,16 +410,6 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): "TPU currently requires fully defined shapes. Either use " "set_shape() on the input tensors or use " "dataset.batch(..., drop_remainder=True).") - types = nest.flatten(multi_worker_iterator.output_types) - - enqueue_ops = [ - self._get_enqueue_op_per_host(host_id, multi_worker_iterator, shapes, - iterations) - for host_id in range(self.num_hosts)] - - def dequeue_fn(): - dequeued = tpu_ops.infeed_dequeue_tuple(dtypes=types, shapes=shapes) - return nest.pack_sequence_as(output_shapes, dequeued) # Wrap `fn` for repeat. if initial_loop_values is None: @@ -335,10 +417,9 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): initial_loop_values = nest.flatten(initial_loop_values) ctx = input_lib.MultiStepContext() - def run_fn(*args, **kwargs): + def run_fn(inputs): """Single step on the TPU device.""" - del args, kwargs - fn_result = fn(ctx, dequeue_fn()) + 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]): @@ -346,9 +427,6 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): else: return fn_result - def iterate_on_tpu(): - return training_loop.repeat(iterations, run_fn, initial_loop_values) - # 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 @@ -358,95 +436,65 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): self._outer_control_flow_context = ( ops.get_default_graph()._get_control_flow_context()) # pylint: disable=protected-access - # pylint: disable=protected-access - if self._container_strategy()._disable_training_loop_on_host: - replicate_inputs = [[]] * self._num_replicas_in_sync - replicate_outputs = tpu.replicate(iterate_on_tpu, replicate_inputs) - else: - def rewrite_fn(*args): - """The rewritten step fn running on TPU.""" - del args - replicate_inputs = [[]] * self._num_replicas_in_sync - 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) + 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, enqueue_ops) + ctx.run_op = control_flow_ops.group(replicate_outputs) - if self._container_strategy()._disable_training_loop_on_host: + 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 (grouped by device) - # [[output0_device0, output1_device0, output2_device0], - # [output0_device1, output1_device1, output2_device1]] + 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]] - last_step_tensor_outputs = [list(x) for x in - zip(*last_step_tensor_outputs)] + 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: - 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 = [] - - # 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 + # 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): @@ -455,27 +503,13 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): with _TPUReplicaContext(self._container_strategy()): return fn(*args, **kwargs) - def _initialize(self): - if context.executing_eagerly(): - # TODO(priyag): Add appopriate call here when eager is supported for TPUs. - raise NotImplementedError("Eager mode not supported in TPUStrategy.") - else: - # TODO(jhseu): We need this hack because DistributionStrategies must be - # pickleable for copy.deepcopy(). Remove when initialize_system goes away. - graph = ops.get_default_graph() - tpu_init = graph.get_collection(_TPU_INITIALIZE_SYSTEM_COLLECTION) - if tpu_init: - return tpu_init - graph.add_to_collection(_TPU_INITIALIZE_SYSTEM_COLLECTION, - tpu.initialize_system()) - return graph.get_collection(_TPU_INITIALIZE_SYSTEM_COLLECTION) - - def _finalize(self): - if context.executing_eagerly(): - # TODO(priyag): Add appopriate call here when eager is supported for TPUs. - raise NotImplementedError("Eager mode not supported in TPUStrategy.") - else: - return [] + 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`.""" @@ -483,6 +517,9 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): 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 @@ -499,9 +536,10 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): # 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(): - kwargs["initial_value"] = array_ops.identity( - value_list[0].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): @@ -535,19 +573,24 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): return cross_device_ops_lib.reduce_non_distributed_value( reduce_op, self._device_map, value, destinations) - # Validate that the destination is same as the host device - # Note we don't do this when in replicate context as the reduction is - # performed on the TPU device itself. devices = cross_device_ops_lib.get_devices_from(destinations) - if len(devices) == 1: - assert device_util.canonicalize(devices[0]) == device_util.canonicalize( - self._host_device) - else: + if len(devices) != 1: raise ValueError("Multiple devices are not supported for TPUStrategy") - output = math_ops.add_n(value) - if reduce_op == reduce_util.ReduceOp.MEAN: - return output * (1. / len(value)) + # 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): @@ -599,15 +642,34 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): @property def num_hosts(self): - return self._tpu_metadata.num_hosts + 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): - return self._tpu_metadata.num_of_cores_per_host + 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): - return self._tpu_metadata.num_cores + 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): @@ -675,6 +737,13 @@ class TPUExtended(distribute_lib.DistributionStrategyExtended): # 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 @@ -682,15 +751,48 @@ class _TPUReplicaContext(distribute_lib.ReplicaContext): """Replication Context class for TPU Strategy.""" # TODO(sourabhbajaj): Call for each replica should be updating this. - def __init__(self, strategy): - # TODO(b/118385803): properly initialize replica_id, instead of always 0 - replica_id = constant_op.constant(0, dtypes.int32) + # 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) + 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) - return (ds.extended.worker_devices[replica_id],) + + 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 diff --git a/tensorflow/contrib/distributions/BUILD b/tensorflow/contrib/distributions/BUILD index 3079175015a9aee1625404902070df8f13b2089c..c2300286d3be4bb757dac588623c47044a1a9db5 100644 --- a/tensorflow/contrib/distributions/BUILD +++ b/tensorflow/contrib/distributions/BUILD @@ -822,7 +822,7 @@ cuda_py_test( cuda_py_test( name = "affine_test", - size = "large", + size = "medium", srcs = ["python/kernel_tests/bijectors/affine_test.py"], additional_deps = [ ":bijectors_py", @@ -837,7 +837,7 @@ cuda_py_test( "//tensorflow/python:math_ops", "//tensorflow/python:platform_test", ], - shard_count = 5, + shard_count = 10, tags = ["noasan"], # times out b/63678675 ) diff --git a/tensorflow/contrib/distributions/python/ops/inverse_gamma.py b/tensorflow/contrib/distributions/python/ops/inverse_gamma.py index 452628257ea96713453bf2aa32b5baa9d6d0cb86..1006dfac49f36baa7cf5136f6f2982e3fd965298 100644 --- a/tensorflow/contrib/distributions/python/ops/inverse_gamma.py +++ b/tensorflow/contrib/distributions/python/ops/inverse_gamma.py @@ -249,9 +249,9 @@ class InverseGamma(distribution.Distribution): `self.allow_nan_stats` is `False`, an exception will be raised rather than returning `NaN`.""") def _variance(self): - var = (math_ops.square(self.rate) - / math_ops.square(self.concentration - 1.) - / (self.concentration - 2.)) + var = ( + math_ops.square(self.rate) / math_ops.squared_difference( + self.concentration, 1.) / (self.concentration - 2.)) if self.allow_nan_stats: nan = array_ops.fill( self.batch_shape_tensor(), diff --git a/tensorflow/contrib/eager/python/examples/BUILD b/tensorflow/contrib/eager/python/examples/BUILD index 97c299a911c9180bf69faa0fa46527e80eada790..3e0881754c750f4d36e2e4dd8b80835b031c658c 100644 --- a/tensorflow/contrib/eager/python/examples/BUILD +++ b/tensorflow/contrib/eager/python/examples/BUILD @@ -6,16 +6,16 @@ package(default_visibility = ["//tensorflow:internal"]) py_library( name = "examples_pip", deps = [ - "//tensorflow/contrib/eager/python/examples/densenet", - "//tensorflow/contrib/eager/python/examples/gan:mnist", + "//tensorflow/contrib/eager/python/examples/densenet:densenet_lib", + "//tensorflow/contrib/eager/python/examples/gan:mnist_lib", "//tensorflow/contrib/eager/python/examples/l2hmc", "//tensorflow/contrib/eager/python/examples/l2hmc:neural_nets", - "//tensorflow/contrib/eager/python/examples/linear_regression", + "//tensorflow/contrib/eager/python/examples/linear_regression:linear_regression_lib", "//tensorflow/contrib/eager/python/examples/resnet50", "//tensorflow/contrib/eager/python/examples/revnet", "//tensorflow/contrib/eager/python/examples/revnet:config", - "//tensorflow/contrib/eager/python/examples/rnn_colorbot", - "//tensorflow/contrib/eager/python/examples/rnn_ptb", + "//tensorflow/contrib/eager/python/examples/rnn_colorbot:rnn_colorbot_lib", + "//tensorflow/contrib/eager/python/examples/rnn_ptb:rnn_ptb_lib", "//tensorflow/contrib/eager/python/examples/spinn:data", ], ) diff --git a/tensorflow/contrib/eager/python/examples/densenet/BUILD b/tensorflow/contrib/eager/python/examples/densenet/BUILD index e2154fcc5fcf774dcd52285d9442dfd5073a4992..fbb5daf230bb79f08a3d071062ddc0e8507ab324 100644 --- a/tensorflow/contrib/eager/python/examples/densenet/BUILD +++ b/tensorflow/contrib/eager/python/examples/densenet/BUILD @@ -9,6 +9,13 @@ py_binary( name = "densenet", srcs = ["densenet.py"], srcs_version = "PY2AND3", + deps = [":densenet_lib"], +) + +py_library( + name = "densenet_lib", + srcs = ["densenet.py"], + srcs_version = "PY2AND3", deps = [ "//tensorflow:tensorflow_py", "//tensorflow/contrib/eager/python:tfe", @@ -17,33 +24,37 @@ py_binary( cuda_py_test( name = "densenet_test", - size = "large", + size = "medium", srcs = ["densenet_test.py"], additional_deps = [ ":densenet", "//tensorflow/contrib/eager/python:tfe", "//tensorflow:tensorflow_py", ], + shard_count = 4, tags = [ "no_pip", "optonly", + "oss_serial", ], ) cuda_py_test( name = "densenet_graph_test", - size = "large", + size = "medium", srcs = ["densenet_graph_test.py"], additional_deps = [ ":densenet", "//third_party/py/numpy", "//tensorflow:tensorflow_py", ], + shard_count = 4, tags = [ "no_pip", "noasan", "nomsan", "notsan", "optonly", + "oss_serial", ], ) diff --git a/tensorflow/contrib/eager/python/examples/gan/BUILD b/tensorflow/contrib/eager/python/examples/gan/BUILD index d64c8eb9ce122fa277567b2fbc632abfbc72df64..d99a519112787bad664232983208279cfb4d0036 100644 --- a/tensorflow/contrib/eager/python/examples/gan/BUILD +++ b/tensorflow/contrib/eager/python/examples/gan/BUILD @@ -9,6 +9,13 @@ py_binary( name = "mnist", srcs = ["mnist.py"], srcs_version = "PY2AND3", + deps = [":mnist_lib"], +) + +py_library( + name = "mnist_lib", + srcs = ["mnist.py"], + srcs_version = "PY2AND3", deps = [ "//tensorflow:tensorflow_py", "//tensorflow/contrib/eager/python:tfe", @@ -20,7 +27,7 @@ cuda_py_test( name = "mnist_test", srcs = ["mnist_test.py"], additional_deps = [ - ":mnist", + ":mnist_lib", "//tensorflow/contrib/eager/python:tfe", "//tensorflow:tensorflow_py", ], @@ -30,7 +37,7 @@ cuda_py_test( name = "mnist_graph_test", srcs = ["mnist_graph_test.py"], additional_deps = [ - ":mnist", + ":mnist_lib", "//third_party/py/numpy", "//tensorflow:tensorflow_py", ], diff --git a/tensorflow/contrib/eager/python/examples/generative_examples/cvae.ipynb b/tensorflow/contrib/eager/python/examples/generative_examples/cvae.ipynb index 1a08cc0fd06516be4af5c2b0b46a3ffcf9101e95..e1a02db76f705414a34d232022f50124a5a6a3ed 100644 --- a/tensorflow/contrib/eager/python/examples/generative_examples/cvae.ipynb +++ b/tensorflow/contrib/eager/python/examples/generative_examples/cvae.ipynb @@ -13,11 +13,13 @@ "\n", "# Convolutional VAE: An example with tf.keras and eager\n", "\n", + "This example has moved:\n", + "\n", "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\u003ctd\u003e\n", - "\u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/contrib/eager/python/examples/generative_examples/cvae.ipynb\"\u003e\n", + "\u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/r2/tutorials/generative/cvae.ipynb\"\u003e\n", " \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e \n", "\u003c/td\u003e\u003ctd\u003e\n", - "\u003ca target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/examples/generative_examples/cvae.ipynb\"\u003e\u003cimg width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\u003c/td\u003e\u003c/table\u003e" + "\u003ca target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/r2/tutorials/generative/cvae.ipynb\"\u003e\u003cimg width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\u003c/td\u003e\u003c/table\u003e" ] }, { @@ -28,604 +30,14 @@ }, "source": [ "![evolution of output during training](https://tensorflow.org/images/autoencoders/cvae.gif)\n", - "\n", - "This notebook demonstrates how to generate images of handwritten digits using [tf.keras](https://www.tensorflow.org/programmers_guide/keras) and [eager execution](https://www.tensorflow.org/programmers_guide/eager) by training a Variational Autoencoder. (VAE, [[1]](https://arxiv.org/abs/1312.6114), [[2]](https://arxiv.org/abs/1401.4082)).\n", "\n" ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "P-JuIu2N_SQf" - }, - "outputs": [], - "source": [ - "# to generate gifs\n", - "!pip install imageio" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "e1_Y75QXJS6h" - }, - "source": [ - "## Import TensorFlow and enable Eager execution" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "YfIk2es3hJEd" - }, - "outputs": [], - "source": [ - "from __future__ import absolute_import, division, print_function\n", - "\n", - "# Import TensorFlow \u003e= 1.9 and enable eager execution\n", - "import tensorflow as tf\n", - "tfe = tf.contrib.eager\n", - "tf.enable_eager_execution()\n", - "\n", - "import os\n", - "import time\n", - "import numpy as np\n", - "import glob\n", - "import matplotlib.pyplot as plt\n", - "import PIL\n", - "import imageio\n", - "from IPython import display" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "iYn4MdZnKCey" - }, - "source": [ - "## Load the MNIST dataset\n", - "Each MNIST image is originally a vector of 784 integers, each of which is between 0-255 and represents the intensity of a pixel. We model each pixel with a Bernoulli distribution in our model, and we statically binarize the dataset." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "a4fYMGxGhrna" - }, - "outputs": [], - "source": [ - "(train_images, _), (test_images, _) = tf.keras.datasets.mnist.load_data()" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "NFC2ghIdiZYE" - }, - "outputs": [], - "source": [ - "train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')\n", - "test_images = test_images.reshape(test_images.shape[0], 28, 28, 1).astype('float32')\n", - "\n", - "# Normalizing the images to the range of [0., 1.]\n", - "train_images /= 255.\n", - "test_images /= 255.\n", - "\n", - "# Binarization\n", - "train_images[train_images \u003e= .5] = 1.\n", - "train_images[train_images \u003c .5] = 0.\n", - "test_images[test_images \u003e= .5] = 1.\n", - "test_images[test_images \u003c .5] = 0." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "S4PIDhoDLbsZ" - }, - "outputs": [], - "source": [ - "TRAIN_BUF = 60000\n", - "BATCH_SIZE = 100\n", - "\n", - "TEST_BUF = 10000" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "PIGN6ouoQxt3" - }, - "source": [ - "## Use *tf.data* to create batches and shuffle the dataset" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "-yKCCQOoJ7cn" - }, - "outputs": [], - "source": [ - "train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(TRAIN_BUF).batch(BATCH_SIZE)\n", - "test_dataset = tf.data.Dataset.from_tensor_slices(test_images).shuffle(TEST_BUF).batch(BATCH_SIZE)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "THY-sZMiQ4UV" - }, - "source": [ - "## Wire up the generative and inference network with *tf.keras.Sequential*\n", - "\n", - "In our VAE example, we use two small ConvNets for the generative and inference network. Since these neural nets are small, we use `tf.keras.Sequential` to simplify our code. Let $x$ and $z$ denote the observation and latent variable respectively in the following descriptions. \n", - "\n", - "### Generative Network\n", - "This defines the generative model which takes a latent encoding as input, and outputs the parameters for a conditional distribution of the observation, i.e. $p(x|z)$. Additionally, we use a unit Gaussian prior $p(z)$ for the latent variable.\n", - "\n", - "### Inference Network\n", - "This defines an approximate posterior distribution $q(z|x)$, which takes as input an observation and outputs a set of parameters for the conditional distribution of the latent representation. In this example, we simply model this distribution as a diagonal Gaussian. In this case, the inference network outputs the mean and log-variance parameters of a factorized Gaussian (log-variance instead of the variance directly is for numerical stability).\n", - "\n", - "### Reparameterization Trick\n", - "During optimization, we can sample from $q(z|x)$ by first sampling from a unit Gaussian, and then multiplying by the standard deviation and adding the mean. This ensures the gradients could pass through the sample to the inference network parameters.\n", - "\n", - "### Network architecture\n", - "For the inference network, we use two convolutional layers followed by a fully-connected layer. In the generative network, we mirror this architecture by using a fully-connected layer followed by three convolution transpose layers (a.k.a. deconvolutional layers in some contexts). Note, it's common practice to avoid using batch normalization when training VAEs, since the additional stochasticity due to using mini-batches may aggravate instability on top of the stochasticity from sampling." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "VGLbvBEmjK0a" - }, - "outputs": [], - "source": [ - "class CVAE(tf.keras.Model):\n", - " def __init__(self, latent_dim):\n", - " super(CVAE, self).__init__()\n", - " self.latent_dim = latent_dim\n", - " self.inference_net = tf.keras.Sequential(\n", - " [\n", - " tf.keras.layers.InputLayer(input_shape=(28, 28, 1)),\n", - " tf.keras.layers.Conv2D(\n", - " filters=32, kernel_size=3, strides=(2, 2), activation=tf.nn.relu),\n", - " tf.keras.layers.Conv2D(\n", - " filters=64, kernel_size=3, strides=(2, 2), activation=tf.nn.relu),\n", - " tf.keras.layers.Flatten(),\n", - " # No activation\n", - " tf.keras.layers.Dense(latent_dim + latent_dim),\n", - " ]\n", - " )\n", - "\n", - " self.generative_net = tf.keras.Sequential(\n", - " [\n", - " tf.keras.layers.InputLayer(input_shape=(latent_dim,)),\n", - " tf.keras.layers.Dense(units=7*7*32, activation=tf.nn.relu),\n", - " tf.keras.layers.Reshape(target_shape=(7, 7, 32)),\n", - " tf.keras.layers.Conv2DTranspose(\n", - " filters=64,\n", - " kernel_size=3,\n", - " strides=(2, 2),\n", - " padding=\"SAME\",\n", - " activation=tf.nn.relu),\n", - " tf.keras.layers.Conv2DTranspose(\n", - " filters=32,\n", - " kernel_size=3,\n", - " strides=(2, 2),\n", - " padding=\"SAME\",\n", - " activation=tf.nn.relu),\n", - " # No activation\n", - " tf.keras.layers.Conv2DTranspose(\n", - " filters=1, kernel_size=3, strides=(1, 1), padding=\"SAME\"),\n", - " ]\n", - " )\n", - "\n", - " def sample(self, eps=None):\n", - " if eps is None:\n", - " eps = tf.random_normal(shape=(100, self.latent_dim))\n", - " return self.decode(eps, apply_sigmoid=True)\n", - "\n", - " def encode(self, x):\n", - " mean, logvar = tf.split(self.inference_net(x), num_or_size_splits=2, axis=1)\n", - " return mean, logvar\n", - "\n", - " def reparameterize(self, mean, logvar):\n", - " eps = tf.random_normal(shape=mean.shape)\n", - " return eps * tf.exp(logvar * .5) + mean\n", - "\n", - " def decode(self, z, apply_sigmoid=False):\n", - " logits = self.generative_net(z)\n", - " if apply_sigmoid:\n", - " probs = tf.sigmoid(logits)\n", - " return probs\n", - "\n", - " return logits" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "0FMYgY_mPfTi" - }, - "source": [ - "## Define the loss function and the optimizer\n", - "\n", - "VAEs train by maximizing the evidence lower bound (ELBO) on the marginal log-likelihood:\n", - "\n", - "$$\\log p(x) \\ge \\text{ELBO} = \\mathbb{E}_{q(z|x)}\\left[\\log \\frac{p(x, z)}{q(z|x)}\\right].$$\n", - "\n", - "In practice, we optimize the single sample Monte Carlo estimate of this expectation:\n", - "\n", - "$$\\log p(x| z) + \\log p(z) - \\log q(z|x),$$\n", - "where $z$ is sampled from $q(z|x)$.\n", - "\n", - "**Note**: we could also analytically compute the KL term, but here we incorporate all three terms in the Monte Carlo estimator for simplicity." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "iWCn_PVdEJZ7" - }, - "outputs": [], - "source": [ - "def log_normal_pdf(sample, mean, logvar, raxis=1):\n", - " log2pi = tf.log(2. * np.pi)\n", - " return tf.reduce_sum(\n", - " -.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),\n", - " axis=raxis)\n", - "\n", - "def compute_loss(model, x):\n", - " mean, logvar = model.encode(x)\n", - " z = model.reparameterize(mean, logvar)\n", - " x_logit = model.decode(z)\n", - "\n", - " cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=x)\n", - " logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])\n", - " logpz = log_normal_pdf(z, 0., 0.)\n", - " logqz_x = log_normal_pdf(z, mean, logvar)\n", - " return -tf.reduce_mean(logpx_z + logpz - logqz_x)\n", - "\n", - "def compute_gradients(model, x):\n", - " with tf.GradientTape() as tape:\n", - " loss = compute_loss(model, x)\n", - " return tape.gradient(loss, model.trainable_variables), loss\n", - "\n", - "optimizer = tf.train.AdamOptimizer(1e-4)\n", - "def apply_gradients(optimizer, gradients, variables, global_step=None):\n", - " optimizer.apply_gradients(zip(gradients, variables), global_step=global_step)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "Rw1fkAczTQYh" - }, - "source": [ - "## Training\n", - "\n", - "* We start by iterating over the dataset\n", - "* During each iteration, we pass the image to the encoder to obtain a set of mean and log-variance parameters of the approximate posterior $q(z|x)$\n", - "* We then apply the *reparameterization trick* to sample from $q(z|x)$\n", - "* Finally, we pass the reparameterized samples to the decoder to obtain the logits of the generative distribution $p(x|z)$\n", - "* **Note:** Since we use the dataset loaded by keras with 60k datapoints in the training set and 10k datapoints in the test set, our resulting ELBO on the test set is slightly higher than reported results in the literature which uses dynamic binarization of Larochelle's MNIST.\n", - "\n", - "## Generate Images\n", - "\n", - "* After training, it is time to generate some images\n", - "* We start by sampling a set of latent vectors from the unit Gaussian prior distribution $p(z)$\n", - "* The generator will then convert the latent sample $z$ to logits of the observation, giving a distribution $p(x|z)$\n", - "* Here we plot the probabilities of Bernoulli distributions\n" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "NS2GWywBbAWo" - }, - "outputs": [], - "source": [ - "epochs = 100\n", - "latent_dim = 50\n", - "num_examples_to_generate = 16\n", - "\n", - "# keeping the random vector constant for generation (prediction) so\n", - "# it will be easier to see the improvement.\n", - "random_vector_for_generation = tf.random_normal(\n", - " shape=[num_examples_to_generate, latent_dim])\n", - "model = CVAE(latent_dim)" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "RmdVsmvhPxyy" - }, - "outputs": [], - "source": [ - "def generate_and_save_images(model, epoch, test_input):\n", - " predictions = model.sample(test_input)\n", - " fig = plt.figure(figsize=(4,4))\n", - "\n", - " for i in range(predictions.shape[0]):\n", - " plt.subplot(4, 4, i+1)\n", - " plt.imshow(predictions[i, :, :, 0], cmap='gray')\n", - " plt.axis('off')\n", - "\n", - " # tight_layout minimizes the overlap between 2 sub-plots\n", - " plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))\n", - " plt.show()" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "2M7LmLtGEMQJ" - }, - "outputs": [], - "source": [ - "generate_and_save_images(model, 0, random_vector_for_generation)\n", - "\n", - "for epoch in range(1, epochs + 1):\n", - " start_time = time.time()\n", - " for train_x in train_dataset:\n", - " gradients, loss = compute_gradients(model, train_x)\n", - " apply_gradients(optimizer, gradients, model.trainable_variables)\n", - " end_time = time.time()\n", - "\n", - " if epoch % 1 == 0:\n", - " loss = tfe.metrics.Mean()\n", - " for test_x in test_dataset:\n", - " loss(compute_loss(model, test_x))\n", - " elbo = -loss.result()\n", - " display.clear_output(wait=False)\n", - " print('Epoch: {}, Test set ELBO: {}, '\n", - " 'time elapse for current epoch {}'.format(epoch,\n", - " elbo,\n", - " end_time - start_time))\n", - " generate_and_save_images(\n", - " model, epoch, random_vector_for_generation)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "P4M_vIbUi7c0" - }, - "source": [ - "### Display an image using the epoch number" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "WfO5wCdclHGL" - }, - "outputs": [], - "source": [ - "def display_image(epoch_no):\n", - " return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "5x3q9_Oe5q0A" - }, - "outputs": [], - "source": [ - "display_image(epochs) # Display images" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "NywiH3nL8guF" - }, - "source": [ - "### Generate a GIF of all the saved images." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "IGKQgENQ8lEI" - }, - "outputs": [], - "source": [ - "with imageio.get_writer('cvae.gif', mode='I') as writer:\n", - " filenames = glob.glob('image*.png')\n", - " filenames = sorted(filenames)\n", - " last = -1\n", - " for i,filename in enumerate(filenames):\n", - " frame = 2*(i**0.5)\n", - " if round(frame) \u003e round(last):\n", - " last = frame\n", - " else:\n", - " continue\n", - " image = imageio.imread(filename)\n", - " writer.append_data(image)\n", - " image = imageio.imread(filename)\n", - " writer.append_data(image)\n", - " \n", - "# this is a hack to display the gif inside the notebook\n", - "os.system('cp cvae.gif cvae.gif.png')" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "uV0yiKpzNP1b" - }, - "outputs": [], - "source": [ - "display.Image(filename=\"cvae.gif.png\")" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "yQXO_dlXkKsT" - }, - "source": [ - "To downlod the animation from Colab uncomment the code below:" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "4fSJS3m5HLFM" - }, - "outputs": [], - "source": [ - "#from google.colab import files\n", - "#files.download('cvae.gif')" - ] } ], "metadata": { "accelerator": "GPU", "colab": { "collapsed_sections": [], - "default_view": {}, "name": "cvae.ipynb", "private_outputs": true, "provenance": [ @@ -635,8 +47,7 @@ } ], "toc_visible": true, - "version": "0.3.2", - "views": {} + "version": "0.3.2" }, "kernelspec": { "display_name": "Python 3", diff --git a/tensorflow/contrib/eager/python/examples/generative_examples/dcgan.ipynb b/tensorflow/contrib/eager/python/examples/generative_examples/dcgan.ipynb index 78fcd397087fd1fd64aebed7ac3b5c6b2f45c450..53767058838459e56215d286e9f8f8eb66287147 100644 --- a/tensorflow/contrib/eager/python/examples/generative_examples/dcgan.ipynb +++ b/tensorflow/contrib/eager/python/examples/generative_examples/dcgan.ipynb @@ -1,26 +1,11 @@ { - "nbformat": 4, - "nbformat_minor": 0, - "metadata": { - "colab": { - "name": "dcgan.ipynb", - "version": "0.3.2", - "provenance": [], - "collapsed_sections": [] - }, - "kernelspec": { - "name": "python2", - "display_name": "Python 2" - }, - "accelerator": "GPU" - }, "cells": [ { + "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "0TD5ZrvEMbhZ" }, - "cell_type": "markdown", "source": [ "**Copyright 2018 The TensorFlow Authors**.\n", "\n", @@ -28,851 +13,39 @@ "\n", "# Generating Handwritten Digits with DCGAN\n", "\n", - "
\n", - "\n", - " Run in Google Colab \n", - "\n", - "View source on GitHub
" - ] - }, - { - "metadata": { - "colab_type": "text", - "id": "ITZuApL56Mny" - }, - "cell_type": "markdown", - "source": [ - "This tutorial demonstrates how to generate images of handwritten digits using a Deep Convolutional Generative Adversarial Network ([DCGAN](https://arxiv.org/pdf/1511.06434.pdf)). The code is written in [tf.keras](https://www.tensorflow.org/programmers_guide/keras) with [eager execution](https://www.tensorflow.org/programmers_guide/eager) enabled. " - ] - }, - { - "metadata": { - "colab_type": "toc", - "id": "x2McrO9bMyLN" - }, - "cell_type": "markdown", - "source": [ - ">[Generating Handwritten Digits with DCGAN](#scrollTo=0TD5ZrvEMbhZ)\n", - "\n", - ">>[What are GANs?](#scrollTo=2MbKJY38Puy9)\n", - "\n", - ">>>[Import TensorFlow and enable eager execution](#scrollTo=e1_Y75QXJS6h)\n", - "\n", - ">>>[Load the dataset](#scrollTo=iYn4MdZnKCey)\n", - "\n", - ">>>[Use tf.data to create batches and shuffle the dataset](#scrollTo=PIGN6ouoQxt3)\n", - "\n", - ">>[Create the models](#scrollTo=THY-sZMiQ4UV)\n", - "\n", - ">>>[The Generator Model](#scrollTo=-tEyxE-GMC48)\n", - "\n", - ">>>[The Discriminator model](#scrollTo=D0IKnaCtg6WE)\n", - "\n", - ">>[Define the loss functions and the optimizer](#scrollTo=0FMYgY_mPfTi)\n", - "\n", - ">>>[Generator loss](#scrollTo=Jd-3GCUEiKtv)\n", - "\n", - ">>>[Discriminator loss](#scrollTo=PKY_iPSPNWoj)\n", - "\n", - ">>[Set up GANs for Training](#scrollTo=Rw1fkAczTQYh)\n", - "\n", - ">>[Train the GANs](#scrollTo=dZrd4CdjR-Fp)\n", - "\n", - ">>[Generated images](#scrollTo=P4M_vIbUi7c0)\n", + "This example has moved.\n", "\n", - ">>[Learn more about GANs](#scrollTo=k6qC-SbjK0yW)\n", - "\n" + "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\u003ctd\u003e\n", + "\u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/r2/tutorials/generative/dcgan.ipynb\"\u003e\n", + " \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e \n", + "\u003c/td\u003e\u003ctd\u003e\n", + "\u003ca target=\"_blank\" href=\"https://github.com/tensorflow/blob/master/site/en/r2/tutorials/generative/dcgan.ipynb\"\u003e\u003cimg width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\u003c/td\u003e\u003c/table\u003e" ] }, { + "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "2MbKJY38Puy9" }, - "cell_type": "markdown", "source": [ - "## What are GANs?\n", - "GANs, or [Generative Adversarial Networks](https://arxiv.org/abs/1406.2661), are a framework for estimating generative models. Two models are trained simultaneously by an adversarial process: a Generator, which is responsible for generating data (say, images), and a Discriminator, which is responsible for estimating the probability that an image was drawn from the training data (the image is real), or was produced by the Generator (the image is fake). During training, the Generator becomes progressively better at generating images, until the Discriminator is no longer able to distinguish real images from fake. \n", - "\n", - "![alt text](https://github.com/margaretmz/tensorflow/blob/margaret-dcgan/tensorflow/contrib/eager/python/examples/generative_examples/gans_diagram.png?raw=1)\n", - "\n", - "We will demonstrate this process end-to-end on MNIST. Below is an animation that shows a series of images produced by the Generator as it was trained for 50 epochs. Overtime, the generated images become increasingly difficult to distinguish from the training set.\n", - "\n", - "To learn more about GANs, we recommend MIT's [Intro to Deep Learning](http://introtodeeplearning.com/) course, which includes a lecture on Deep Generative Models ([video](https://youtu.be/JVb54xhEw6Y) | [slides](http://introtodeeplearning.com/materials/2018_6S191_Lecture4.pdf)). Now, let's head to the code!\n", - "\n", "![sample output](https://tensorflow.org/images/gan/dcgan.gif)" ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "collapsed_sections": [], + "name": "dcgan.ipynb", + "provenance": [], + "version": "0.3.2" }, - { - "metadata": { - "colab_type": "code", - "id": "u_2z-B3piVsw", - "colab": {} - }, - "cell_type": "code", - "source": [ - "# Install imgeio in order to generate an animated gif showing the image generating process\n", - "!pip install imageio" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "e1_Y75QXJS6h" - }, - "cell_type": "markdown", - "source": [ - "### Import TensorFlow and enable eager execution" - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "YfIk2es3hJEd", - "colab": {} - }, - "cell_type": "code", - "source": [ - "import tensorflow as tf\n", - "tf.enable_eager_execution()\n", - "\n", - "import glob\n", - "import imageio\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import os\n", - "import PIL\n", - "import time\n", - "\n", - "from IPython import display" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "iYn4MdZnKCey" - }, - "cell_type": "markdown", - "source": [ - "### Load the dataset\n", - "\n", - "We are going to use the MNIST dataset to train the generator and the discriminator. The generator will generate handwritten digits resembling the MNIST data." - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "a4fYMGxGhrna", - "colab": {} - }, - "cell_type": "code", - "source": [ - "(train_images, train_labels), (_, _) = tf.keras.datasets.mnist.load_data()" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "code", - "id": "NFC2ghIdiZYE", - "colab": {} - }, - "cell_type": "code", - "source": [ - "train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')\n", - "train_images = (train_images - 127.5) / 127.5 # Normalize the images to [-1, 1]" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "code", - "id": "S4PIDhoDLbsZ", - "colab": {} - }, - "cell_type": "code", - "source": [ - "BUFFER_SIZE = 60000\n", - "BATCH_SIZE = 256" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "PIGN6ouoQxt3" - }, - "cell_type": "markdown", - "source": [ - "### Use tf.data to create batches and shuffle the dataset" - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "-yKCCQOoJ7cn", - "colab": {} - }, - "cell_type": "code", - "source": [ - "train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "THY-sZMiQ4UV" - }, - "cell_type": "markdown", - "source": [ - "## Create the models\n", - "\n", - "We will use tf.keras [Sequential API](https://www.tensorflow.org/guide/keras#sequential_model) to define the generator and discriminator models." - ] - }, - { - "metadata": { - "colab_type": "text", - "id": "-tEyxE-GMC48" - }, - "cell_type": "markdown", - "source": [ - "### The Generator Model\n", - "\n", - "The generator is responsible for creating convincing images that are good enough to fool the discriminator. The network architecture for the generator consists of [Conv2DTranspose](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Conv2DTranspose) (Upsampling) layers. We start with a fully connected layer and upsample the image two times in order to reach the desired image size of 28x28x1. We increase the width and height, and reduce the depth as we move through the layers in the network. We use [Leaky ReLU](https://www.tensorflow.org/api_docs/python/tf/keras/layers/LeakyReLU) activation for each layer except for the last one where we use a tanh activation." - ] - }, - { - "metadata": { - "id": "6bpTcDqoLWjY", - "colab_type": "code", - "colab": {} - }, - "cell_type": "code", - "source": [ - "def make_generator_model():\n", - " model = tf.keras.Sequential()\n", - " model.add(tf.keras.layers.Dense(7*7*256, use_bias=False, input_shape=(100,)))\n", - " model.add(tf.keras.layers.BatchNormalization())\n", - " model.add(tf.keras.layers.LeakyReLU())\n", - " \n", - " model.add(tf.keras.layers.Reshape((7, 7, 256)))\n", - " assert model.output_shape == (None, 7, 7, 256) # Note: None is the batch size\n", - " \n", - " model.add(tf.keras.layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False))\n", - " assert model.output_shape == (None, 7, 7, 128) \n", - " model.add(tf.keras.layers.BatchNormalization())\n", - " model.add(tf.keras.layers.LeakyReLU())\n", - "\n", - " model.add(tf.keras.layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False))\n", - " assert model.output_shape == (None, 14, 14, 64) \n", - " model.add(tf.keras.layers.BatchNormalization())\n", - " model.add(tf.keras.layers.LeakyReLU())\n", - "\n", - " model.add(tf.keras.layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh'))\n", - " assert model.output_shape == (None, 28, 28, 1)\n", - " \n", - " return model" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "D0IKnaCtg6WE" - }, - "cell_type": "markdown", - "source": [ - "### The Discriminator model\n", - "\n", - "The discriminator is responsible for distinguishing fake images from real images. It's similar to a regular CNN-based image classifier." - ] - }, - { - "metadata": { - "id": "dw2tPLmk2pEP", - "colab_type": "code", - "colab": {} - }, - "cell_type": "code", - "source": [ - "def make_discriminator_model():\n", - " model = tf.keras.Sequential()\n", - " model.add(tf.keras.layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same'))\n", - " model.add(tf.keras.layers.LeakyReLU())\n", - " model.add(tf.keras.layers.Dropout(0.3))\n", - " \n", - " model.add(tf.keras.layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))\n", - " model.add(tf.keras.layers.LeakyReLU())\n", - " model.add(tf.keras.layers.Dropout(0.3))\n", - " \n", - " model.add(tf.keras.layers.Flatten())\n", - " model.add(tf.keras.layers.Dense(1))\n", - " \n", - " return model" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "code", - "id": "gDkA05NE6QMs", - "colab": {} - }, - "cell_type": "code", - "source": [ - "generator = make_generator_model()\n", - "discriminator = make_discriminator_model()" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "0FMYgY_mPfTi" - }, - "cell_type": "markdown", - "source": [ - "## Define the loss functions and the optimizer\n", - "\n", - "Let's define the loss functions and the optimizers for the generator and the discriminator.\n" - ] - }, - { - "metadata": { - "colab_type": "text", - "id": "Jd-3GCUEiKtv" - }, - "cell_type": "markdown", - "source": [ - "### Generator loss\n", - "The generator loss is a sigmoid cross entropy loss of the generated images and an array of ones, since the generator is trying to generate fake images that resemble the real images." - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "90BIcCKcDMxz", - "colab": {} - }, - "cell_type": "code", - "source": [ - "def generator_loss(generated_output):\n", - " return tf.losses.sigmoid_cross_entropy(tf.ones_like(generated_output), generated_output)" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "PKY_iPSPNWoj" - }, - "cell_type": "markdown", - "source": [ - "### Discriminator loss\n", - "\n", - "The discriminator loss function takes two inputs: real images, and generated images. Here is how to calculate the discriminator loss:\n", - "1. Calculate real_loss which is a sigmoid cross entropy loss of the real images and an array of ones (since these are the real images).\n", - "2. Calculate generated_loss which is a sigmoid cross entropy loss of the generated images and an array of zeros (since these are the fake images).\n", - "3. Calculate the total_loss as the sum of real_loss and generated_loss." - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "wkMNfBWlT-PV", - "colab": {} - }, - "cell_type": "code", - "source": [ - "def discriminator_loss(real_output, generated_output):\n", - " # [1,1,...,1] with real output since it is true and we want our generated examples to look like it\n", - " real_loss = tf.losses.sigmoid_cross_entropy(multi_class_labels=tf.ones_like(real_output), logits=real_output)\n", - "\n", - " # [0,0,...,0] with generated images since they are fake\n", - " generated_loss = tf.losses.sigmoid_cross_entropy(multi_class_labels=tf.zeros_like(generated_output), logits=generated_output)\n", - "\n", - " total_loss = real_loss + generated_loss\n", - "\n", - " return total_loss" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "MgIc7i0th_Iu" - }, - "cell_type": "markdown", - "source": [ - "The discriminator and the generator optimizers are different since we will train two networks separately." - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "iWCn_PVdEJZ7", - "colab": {} - }, - "cell_type": "code", - "source": [ - "generator_optimizer = tf.train.AdamOptimizer(1e-4)\n", - "discriminator_optimizer = tf.train.AdamOptimizer(1e-4)" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "mWtinsGDPJlV" - }, - "cell_type": "markdown", - "source": [ - "**Checkpoints (Object-based saving)**" - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "CA1w-7s2POEy", - "colab": {} - }, - "cell_type": "code", - "source": [ - "checkpoint_dir = './training_checkpoints'\n", - "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\n", - "checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,\n", - " discriminator_optimizer=discriminator_optimizer,\n", - " generator=generator,\n", - " discriminator=discriminator)" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "Rw1fkAczTQYh" - }, - "cell_type": "markdown", - "source": [ - "## Set up GANs for Training\n", - "\n" - ] - }, - { - "metadata": { - "colab_type": "text", - "id": "5QC5BABamh_c" - }, - "cell_type": "markdown", - "source": [ - "Now it's time to put together the generator and discriminator to set up the Generative Adversarial Networks, as you see in the diagam at the beginning of the tutorial." - ] - }, - { - "metadata": { - "colab_type": "text", - "id": "Ff6oN6PZX27n" - }, - "cell_type": "markdown", - "source": [ - "**Define training parameters**" - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "NS2GWywBbAWo", - "colab": {} - }, - "cell_type": "code", - "source": [ - "EPOCHS = 50\n", - "noise_dim = 100\n", - "num_examples_to_generate = 16\n", - "\n", - "# We'll re-use this random vector used to seed the generator so\n", - "# it will be easier to see the improvement over time.\n", - "random_vector_for_generation = tf.random_normal([num_examples_to_generate,\n", - " noise_dim])" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "jylSonrqSWfi" - }, - "cell_type": "markdown", - "source": [ - "**Define training method**\n", - "\n", - "We start by iterating over the dataset. The generator is given a random vector as an input which is processed to output an image looking like a handwritten digit. The discriminator is then shown the real MNIST images as well as the generated images.\n", - "\n", - "Next, we calculate the generator and the discriminator loss. Then, we calculate the gradients of loss with respect to both the generator and the discriminator variables." - ] - }, - { - "metadata": { - "id": "3t5ibNo05jCB", - "colab_type": "code", - "colab": {} - }, - "cell_type": "code", - "source": [ - "def train_step(images):\n", - " # generating noise from a normal distribution\n", - " noise = tf.random_normal([BATCH_SIZE, noise_dim])\n", - " \n", - " with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n", - " generated_images = generator(noise, training=True)\n", - " \n", - " real_output = discriminator(images, training=True)\n", - " generated_output = discriminator(generated_images, training=True)\n", - " \n", - " gen_loss = generator_loss(generated_output)\n", - " disc_loss = discriminator_loss(real_output, generated_output)\n", - " \n", - " gradients_of_generator = gen_tape.gradient(gen_loss, generator.variables)\n", - " gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.variables)\n", - " \n", - " generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.variables))\n", - " discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.variables))" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "6TSZgwc2BUQ-" - }, - "cell_type": "markdown", - "source": [ - "\n", - "This model takes about ~30 seconds per epoch to train on a single Tesla K80 on Colab, as of October 2018. \n", - "\n", - "Eager execution can be slower than executing the equivalent graph as it can't benefit from whole-program optimizations on the graph, and also incurs overheads of interpreting Python code. By using [tf.contrib.eager.defun](https://www.tensorflow.org/api_docs/python/tf/contrib/eager/defun) to create graph functions, we get a ~20 secs/epoch performance boost (from ~50 secs/epoch down to ~30 secs/epoch). This way we get the best of both eager execution (easier for debugging) and graph mode (better performance)." - ] - }, - { - "metadata": { - "id": "Iwya07_j5p2A", - "colab_type": "code", - "colab": {} - }, - "cell_type": "code", - "source": [ - "train_step = tf.contrib.eager.defun(train_step)" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "code", - "id": "2M7LmLtGEMQJ", - "colab": {} - }, - "cell_type": "code", - "source": [ - "def train(dataset, epochs): \n", - " for epoch in range(epochs):\n", - " start = time.time()\n", - " \n", - " for images in dataset:\n", - " train_step(images)\n", - "\n", - " display.clear_output(wait=True)\n", - " generate_and_save_images(generator,\n", - " epoch + 1,\n", - " random_vector_for_generation)\n", - " \n", - " # saving (checkpoint) the model every 15 epochs\n", - " if (epoch + 1) % 15 == 0:\n", - " checkpoint.save(file_prefix = checkpoint_prefix)\n", - " \n", - " print ('Time taken for epoch {} is {} sec'.format(epoch + 1,\n", - " time.time()-start))\n", - " # generating after the final epoch\n", - " display.clear_output(wait=True)\n", - " generate_and_save_images(generator,\n", - " epochs,\n", - " random_vector_for_generation)" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "2aFF7Hk3XdeW" - }, - "cell_type": "markdown", - "source": [ - "**Generate and save images**\n", - "\n" - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "RmdVsmvhPxyy", - "colab": {} - }, - "cell_type": "code", - "source": [ - "def generate_and_save_images(model, epoch, test_input):\n", - " # make sure the training parameter is set to False because we\n", - " # don't want to train the batchnorm layer when doing inference.\n", - " predictions = model(test_input, training=False)\n", - "\n", - " fig = plt.figure(figsize=(4,4))\n", - " \n", - " for i in range(predictions.shape[0]):\n", - " plt.subplot(4, 4, i+1)\n", - " plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='gray')\n", - " plt.axis('off')\n", - " \n", - " plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))\n", - " plt.show()" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "dZrd4CdjR-Fp" - }, - "cell_type": "markdown", - "source": [ - "## Train the GANs\n", - "We will call the train() method defined above to train the generator and discriminator simultaneously. Note, training GANs can be tricky. It's important that the generator and discriminator do not overpower each other (e.g., that they train at a similar rate).\n", - "\n", - "At the beginning of the training, the generated images look like random noise. As training progresses, you can see the generated digits look increasingly real. After 50 epochs, they look very much like the MNIST digits." - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "Ly3UN0SLLY2l", - "colab": {} - }, - "cell_type": "code", - "source": [ - "%%time\n", - "train(train_dataset, EPOCHS)" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "rfM4YcPVPkNO" - }, - "cell_type": "markdown", - "source": [ - "**Restore the latest checkpoint**" - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "XhXsd0srPo8c", - "colab": {} - }, - "cell_type": "code", - "source": [ - "# restoring the latest checkpoint in checkpoint_dir\n", - "checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "P4M_vIbUi7c0" - }, - "cell_type": "markdown", - "source": [ - "## Generated images \n" - ] - }, - { - "metadata": { - "colab_type": "text", - "id": "mLskt7EfXAjr" - }, - "cell_type": "markdown", - "source": [ - "\n", - "After training, its time to generate some images! \n", - "The last step is to plot the generated images and voila!\n" - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "WfO5wCdclHGL", - "colab": {} - }, - "cell_type": "code", - "source": [ - "# Display a single image using the epoch number\n", - "def display_image(epoch_no):\n", - " return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "code", - "id": "5x3q9_Oe5q0A", - "colab": {} - }, - "cell_type": "code", - "source": [ - "display_image(EPOCHS)" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "NywiH3nL8guF" - }, - "cell_type": "markdown", - "source": [ - "**Generate a GIF of all the saved images**\n", - "\n", - "We will use imageio to create an animated gif using all the images saved during training." - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "IGKQgENQ8lEI", - "colab": {} - }, - "cell_type": "code", - "source": [ - "with imageio.get_writer('dcgan.gif', mode='I') as writer:\n", - " filenames = glob.glob('image*.png')\n", - " filenames = sorted(filenames)\n", - " last = -1\n", - " for i,filename in enumerate(filenames):\n", - " frame = 2*(i**0.5)\n", - " if round(frame) > round(last):\n", - " last = frame\n", - " else:\n", - " continue\n", - " image = imageio.imread(filename)\n", - " writer.append_data(image)\n", - " image = imageio.imread(filename)\n", - " writer.append_data(image)\n", - " \n", - "# this is a hack to display the gif inside the notebook\n", - "os.system('cp dcgan.gif dcgan.gif.png')" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "cGhC3-fMWSwl" - }, - "cell_type": "markdown", - "source": [ - "Display the animated gif with all the mages generated during the training of GANs." - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "uV0yiKpzNP1b", - "colab": {} - }, - "cell_type": "code", - "source": [ - "display.Image(filename=\"dcgan.gif.png\")" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "6EEG-wePkmJQ" - }, - "cell_type": "markdown", - "source": [ - "**Download the animated gif**\n", - "\n", - "Uncomment the code below to download an animated gif from Colab." - ] - }, - { - "metadata": { - "colab_type": "code", - "id": "4UJjSnIMOzOJ", - "colab": {} - }, - "cell_type": "code", - "source": [ - "#from google.colab import files\n", - "#files.download('dcgan.gif')" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "colab_type": "text", - "id": "k6qC-SbjK0yW" - }, - "cell_type": "markdown", - "source": [ - "## Learn more about GANs\n" - ] - }, - { - "metadata": { - "colab_type": "text", - "id": "xjjkT9KAK6H7" - }, - "cell_type": "markdown", - "source": [ - "We hope this tutorial was helpful! As a next step, you might like to experiment with a different dataset, for example the Large-scale Celeb Faces Attributes (CelebA) dataset [available on Kaggle](https://www.kaggle.com/jessicali9530/celeba-dataset/home).\n", - "\n", - "To learn more about GANs:\n", - "\n", - "* Check out MIT's lecture (linked above), or [this](http://cs231n.stanford.edu/slides/2018/cs231n_2018_lecture12.pdf) lecture form Stanford's CS231n. \n", - "\n", - "* We also recommend the [CVPR 2018 Tutorial on GANs](https://sites.google.com/view/cvpr2018tutorialongans/), and the [NIPS 2016 Tutorial: Generative Adversarial Networks](https://arxiv.org/abs/1701.00160).\n" - ] + "kernelspec": { + "display_name": "Python 2", + "name": "python2" } - ] -} \ No newline at end of file + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/tensorflow/contrib/eager/python/examples/generative_examples/gans_diagram.png b/tensorflow/contrib/eager/python/examples/generative_examples/gans_diagram.png deleted file mode 100644 index b715bd83ef117641c6429e0ac173dbe9b8d5fd88..0000000000000000000000000000000000000000 Binary files a/tensorflow/contrib/eager/python/examples/generative_examples/gans_diagram.png and /dev/null differ diff --git a/tensorflow/contrib/eager/python/examples/generative_examples/image_captioning_with_attention.ipynb b/tensorflow/contrib/eager/python/examples/generative_examples/image_captioning_with_attention.ipynb index 12c5eff2b4aa901bdab52bf545e95b1e4dce7468..979772acd3f823a8cc53ab5e026946ad3bb19353 100644 --- a/tensorflow/contrib/eager/python/examples/generative_examples/image_captioning_with_attention.ipynb +++ b/tensorflow/contrib/eager/python/examples/generative_examples/image_captioning_with_attention.ipynb @@ -1,1174 +1,71 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "K2s1A9eLRPEj" - }, - "source": [ - "##### Copyright 2018 The TensorFlow Authors.\n", - "\n", - "Licensed under the Apache License, Version 2.0 (the \"License\").\n" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "Cffg2i257iMS" - }, - "source": [ - "# Image Captioning with Attention\n", - "\n", - "
\n", - "\n", - " Run in Google Colab \n", - "\n", - "View source on GitHub
" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "QASbY_HGo4Lq" - }, - "source": [ - "Image captioning is the task of generating a caption for an image. Given an image like this:\n", - "\n", - "![Man Surfing](https://tensorflow.org/images/surf.jpg) \n", - "\n", - "[Image Source](https://commons.wikimedia.org/wiki/Surfing#/media/File:Surfing_in_Hawaii.jpg), License: Public Domain\n", - "\n", - "Our goal is to generate a caption, such as \"a surfer riding on a wave\". Here, we'll use an attention-based model. This enables us to see which parts of the image the model focuses on as it generates a caption.\n", - "\n", - "![Prediction](https://tensorflow.org/images/imcap_prediction.png)\n", - "\n", - "This model architecture below is similar to [Show, Attend and Tell: Neural Image Caption Generation with Visual Attention](https://arxiv.org/abs/1502.03044). \n", - "\n", - "The code uses [tf.keras](https://www.tensorflow.org/programmers_guide/keras) and [eager execution](https://www.tensorflow.org/programmers_guide/eager), which you can learn more about in the linked guides.\n", - "\n", - "This notebook is an end-to-end example. If you run it, it will download the [MS-COCO](http://cocodataset.org/#home) dataset, preprocess and cache a subset of the images using Inception V3, train an encoder-decoder model, and use it to generate captions on new images.\n", - "\n", - "The code requires TensorFlow version >=1.9. If you're running this in [Colab]()\n", - "\n", - "In this example, we're training on a relatively small amount of data as an example. On a single P100 GPU, this example will take about ~2 hours to train. We train on the first 30,000 captions (corresponding to about ~20,000 images depending on shuffling, as there are multiple captions per image in the dataset)\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "U8l4RJ0XRPEm" - }, - "outputs": [], - "source": [ - "# Import TensorFlow and enable eager execution\n", - "# This code requires TensorFlow version >=1.9\n", - "import tensorflow as tf\n", - "tf.enable_eager_execution()\n", - "\n", - "# We'll generate plots of attention in order to see which parts of an image\n", - "# our model focuses on during captioning\n", - "import matplotlib.pyplot as plt\n", - "\n", - "# Scikit-learn includes many helpful utilities\n", - "from sklearn.model_selection import train_test_split\n", - "from sklearn.utils import shuffle\n", - "\n", - "import re\n", - "import numpy as np\n", - "import os\n", - "import time\n", - "import json\n", - "from glob import glob\n", - "from PIL import Image\n", - "import pickle" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "b6qbGw8MRPE5" - }, - "source": [ - "## Download and prepare the MS-COCO dataset\n", - "\n", - "We will use the [MS-COCO dataset](http://cocodataset.org/#home) to train our model. This dataset contains >82,000 images, each of which has been annotated with at least 5 different captions. The code below will download and extract the dataset automatically. \n", - "\n", - "**Caution: large download ahead**. We'll use the training set, it's a 13GB file." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "krQuPYTtRPE7" - }, - "outputs": [], - "source": [ - "annotation_zip = tf.keras.utils.get_file('captions.zip', \n", - " cache_subdir=os.path.abspath('.'),\n", - " origin = 'http://images.cocodataset.org/annotations/annotations_trainval2014.zip',\n", - " extract = True)\n", - "annotation_file = os.path.dirname(annotation_zip)+'/annotations/captions_train2014.json'\n", - "\n", - "name_of_zip = 'train2014.zip'\n", - "if not os.path.exists(os.path.abspath('.') + '/' + name_of_zip):\n", - " image_zip = tf.keras.utils.get_file(name_of_zip, \n", - " cache_subdir=os.path.abspath('.'),\n", - " origin = 'http://images.cocodataset.org/zips/train2014.zip',\n", - " extract = True)\n", - " PATH = os.path.dirname(image_zip)+'/train2014/'\n", - "else:\n", - " PATH = os.path.abspath('.')+'/train2014/'" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "aANEzb5WwSzg" - }, - "source": [ - "## Optionally, limit the size of the training set for faster training\n", - "For this example, we'll select a subset of 30,000 captions and use these and the corresponding images to train our model. As always, captioning quality will improve if you choose to use more data." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "4G3b8x8_RPFD" - }, - "outputs": [], - "source": [ - "# read the json file\n", - "with open(annotation_file, 'r') as f:\n", - " annotations = json.load(f)\n", - "\n", - "# storing the captions and the image name in vectors\n", - "all_captions = []\n", - "all_img_name_vector = []\n", - "\n", - "for annot in annotations['annotations']:\n", - " caption = ' ' + annot['caption'] + ' '\n", - " image_id = annot['image_id']\n", - " full_coco_image_path = PATH + 'COCO_train2014_' + '%012d.jpg' % (image_id)\n", - " \n", - " all_img_name_vector.append(full_coco_image_path)\n", - " all_captions.append(caption)\n", - "\n", - "# shuffling the captions and image_names together\n", - "# setting a random state\n", - "train_captions, img_name_vector = shuffle(all_captions,\n", - " all_img_name_vector,\n", - " random_state=1)\n", - "\n", - "# selecting the first 30000 captions from the shuffled set\n", - "num_examples = 30000\n", - "train_captions = train_captions[:num_examples]\n", - "img_name_vector = img_name_vector[:num_examples]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "mPBMgK34RPFL" - }, - "outputs": [], - "source": [ - "len(train_captions), len(all_captions)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "8cSW4u-ORPFQ" - }, - "source": [ - "## Preprocess the images using InceptionV3\n", - "Next, we will use InceptionV3 (pretrained on Imagenet) to classify each image. We will extract features from the last convolutional layer. \n", - "\n", - "First, we will need to convert the images into the format inceptionV3 expects by:\n", - "* Resizing the image to (299, 299)\n", - "* Using the [preprocess_input](https://www.tensorflow.org/api_docs/python/tf/keras/applications/inception_v3/preprocess_input) method to place the pixels in the range of -1 to 1 (to match the format of the images used to train InceptionV3)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "zXR0217aRPFR" - }, - "outputs": [], - "source": [ - "def load_image(image_path):\n", - " img = tf.read_file(image_path)\n", - " img = tf.image.decode_jpeg(img, channels=3)\n", - " img = tf.image.resize_images(img, (299, 299))\n", - " img = tf.keras.applications.inception_v3.preprocess_input(img)\n", - " return img, image_path" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "MDvIu4sXRPFV" - }, - "source": [ - "## Initialize InceptionV3 and load the pretrained Imagenet weights\n", - "\n", - "To do so, we'll create a tf.keras model where the output layer is the last convolutional layer in the InceptionV3 architecture. \n", - "* Each image is forwarded through the network and the vector that we get at the end is stored in a dictionary (image_name --> feature_vector). \n", - "* We use the last convolutional layer because we are using attention in this example. The shape of the output of this layer is ```8x8x2048```. \n", - "* We avoid doing this during training so it does not become a bottleneck. \n", - "* After all the images are passed through the network, we pickle the dictionary and save it to disk." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "RD3vW4SsRPFW" - }, - "outputs": [], - "source": [ - "image_model = tf.keras.applications.InceptionV3(include_top=False, \n", - " weights='imagenet')\n", - "new_input = image_model.input\n", - "hidden_layer = image_model.layers[-1].output\n", - "\n", - "image_features_extract_model = tf.keras.Model(new_input, hidden_layer)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "rERqlR3WRPGO" - }, - "source": [ - "## Caching the features extracted from InceptionV3\n", - "\n", - "We will pre-process each image with InceptionV3 and cache the output to disk. Caching the output in RAM would be faster but memory intensive, requiring 8 \\* 8 \\* 2048 floats per image. At the time of writing, this would exceed the memory limitations of Colab (although these may change, an instance appears to have about 12GB of memory currently). \n", - "\n", - "Performance could be improved with a more sophisticated caching strategy (e.g., by sharding the images to reduce random access disk I/O) at the cost of more code.\n", - "\n", - "This will take about 10 minutes to run in Colab with a GPU. If you'd like to see a progress bar, you could: install [tqdm](https://github.com/tqdm/tqdm) (```!pip install tqdm```), then change this line: \n", - "\n", - "```for img, path in image_dataset:``` \n", - "\n", - "to:\n", - "\n", - "```for img, path in tqdm(image_dataset):```." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "Dx_fvbVgRPGQ" - }, - "outputs": [], - "source": [ - "# getting the unique images\n", - "encode_train = sorted(set(img_name_vector))\n", - "\n", - "# feel free to change the batch_size according to your system configuration\n", - "image_dataset = tf.data.Dataset.from_tensor_slices(\n", - " encode_train).map(load_image).batch(16)\n", - "\n", - "for img, path in image_dataset:\n", - " batch_features = image_features_extract_model(img)\n", - " batch_features = tf.reshape(batch_features, \n", - " (batch_features.shape[0], -1, batch_features.shape[3]))\n", - "\n", - " for bf, p in zip(batch_features, path):\n", - " path_of_feature = p.numpy().decode(\"utf-8\")\n", - " np.save(path_of_feature, bf.numpy())" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "nyqH3zFwRPFi" - }, - "source": [ - "## Preprocess and tokenize the captions\n", - "\n", - "* First, we'll tokenize the captions (e.g., by splitting on spaces). This will give us a vocabulary of all the unique words in the data (e.g., \"surfing\", \"football\", etc).\n", - "* Next, we'll limit the vocabulary size to the top 5,000 words to save memory. We'll replace all other words with the token \"UNK\" (for unknown).\n", - "* Finally, we create a word --> index mapping and vice-versa.\n", - "* We will then pad all sequences to the be same length as the longest one. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "HZfK8RhQRPFj" - }, - "outputs": [], - "source": [ - "# This will find the maximum length of any caption in our dataset\n", - "def calc_max_length(tensor):\n", - " return max(len(t) for t in tensor)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "oJGE34aiRPFo" - }, - "outputs": [], - "source": [ - "# The steps above is a general process of dealing with text processing\n", - "\n", - "# choosing the top 5000 words from the vocabulary\n", - "top_k = 5000\n", - "tokenizer = tf.keras.preprocessing.text.Tokenizer(num_words=top_k, \n", - " oov_token=\"\", \n", - " filters='!\"#$%&()*+.,-/:;=?@[\\]^_`{|}~ ')\n", - "tokenizer.fit_on_texts(train_captions)\n", - "train_seqs = tokenizer.texts_to_sequences(train_captions)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "8Q44tNQVRPFt" - }, - "outputs": [], - "source": [ - "tokenizer.word_index[''] = 0" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "0fpJb5ojRPFv" - }, - "outputs": [], - "source": [ - "# creating the tokenized vectors\n", - "train_seqs = tokenizer.texts_to_sequences(train_captions)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "AidglIZVRPF4" - }, - "outputs": [], - "source": [ - "# padding each vector to the max_length of the captions\n", - "# if the max_length parameter is not provided, pad_sequences calculates that automatically\n", - "cap_vector = tf.keras.preprocessing.sequence.pad_sequences(train_seqs, padding='post')" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "gL0wkttkRPGA" - }, - "outputs": [], - "source": [ - "# calculating the max_length \n", - "# used to store the attention weights\n", - "max_length = calc_max_length(train_seqs)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "M3CD75nDpvTI" - }, - "source": [ - "## Split the data into training and testing" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "iS7DDMszRPGF" - }, - "outputs": [], - "source": [ - "# Create training and validation sets using 80-20 split\n", - "img_name_train, img_name_val, cap_train, cap_val = train_test_split(img_name_vector, \n", - " cap_vector, \n", - " test_size=0.2, \n", - " random_state=0)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "XmViPkRFRPGH" - }, - "outputs": [], - "source": [ - "len(img_name_train), len(cap_train), len(img_name_val), len(cap_val)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "uEWM9xrYcg45" - }, - "source": [ - "## Our images and captions are ready! Next, let's create a tf.data dataset to use for training our model.\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "Q3TnZ1ToRPGV" - }, - "outputs": [], - "source": [ - "# feel free to change these parameters according to your system's configuration\n", - "\n", - "BATCH_SIZE = 64\n", - "BUFFER_SIZE = 1000\n", - "embedding_dim = 256\n", - "units = 512\n", - "vocab_size = len(tokenizer.word_index)\n", - "# shape of the vector extracted from InceptionV3 is (64, 2048)\n", - "# these two variables represent that\n", - "features_shape = 2048\n", - "attention_features_shape = 64" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "SmZS2N0bXG3T" - }, - "outputs": [], - "source": [ - "# loading the numpy files \n", - "def map_func(img_name, cap):\n", - " img_tensor = np.load(img_name.decode('utf-8')+'.npy')\n", - " return img_tensor, cap" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "FDF_Nm3tRPGZ" - }, - "outputs": [], - "source": [ - "dataset = tf.data.Dataset.from_tensor_slices((img_name_train, cap_train))\n", - "\n", - "# using map to load the numpy files in parallel\n", - "# NOTE: Be sure to set num_parallel_calls to the number of CPU cores you have\n", - "# https://www.tensorflow.org/api_docs/python/tf/py_func\n", - "dataset = dataset.map(lambda item1, item2: tf.py_func(\n", - " map_func, [item1, item2], [tf.float32, tf.int32]), num_parallel_calls=8)\n", - "\n", - "# shuffling and batching\n", - "dataset = dataset.shuffle(BUFFER_SIZE)\n", - "# https://www.tensorflow.org/api_docs/python/tf/contrib/data/batch_and_drop_remainder\n", - "dataset = dataset.batch(BATCH_SIZE)\n", - "dataset = dataset.prefetch(1)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "nrvoDphgRPGd" - }, - "source": [ - "## Model\n", - "\n", - "Fun fact, the decoder below is identical to the one in the example for [Neural Machine Translation with Attention]( https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/eager/python/examples/nmt_with_attention/nmt_with_attention.ipynb).\n", - "\n", - "The model architecture is inspired by the [Show, Attend and Tell](https://arxiv.org/pdf/1502.03044.pdf) paper.\n", - "\n", - "* In this example, we extract the features from the lower convolutional layer of InceptionV3 giving us a vector of shape (8, 8, 2048). \n", - "* We squash that to a shape of (64, 2048).\n", - "* This vector is then passed through the CNN Encoder(which consists of a single Fully connected layer).\n", - "* The RNN(here GRU) attends over the image to predict the next word." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "AAppCGLKRPGd" - }, - "outputs": [], - "source": [ - "def gru(units):\n", - " # If you have a GPU, we recommend using the CuDNNGRU layer (it provides a \n", - " # significant speedup).\n", - " if tf.test.is_gpu_available():\n", - " return tf.keras.layers.CuDNNGRU(units, \n", - " return_sequences=True, \n", - " return_state=True, \n", - " recurrent_initializer='glorot_uniform')\n", - " else:\n", - " return tf.keras.layers.GRU(units, \n", - " return_sequences=True, \n", - " return_state=True, \n", - " recurrent_activation='sigmoid', \n", - " recurrent_initializer='glorot_uniform')" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "ja2LFTMSdeV3" - }, - "outputs": [], - "source": [ - "class BahdanauAttention(tf.keras.Model):\n", - " def __init__(self, units):\n", - " super(BahdanauAttention, self).__init__()\n", - " self.W1 = tf.keras.layers.Dense(units)\n", - " self.W2 = tf.keras.layers.Dense(units)\n", - " self.V = tf.keras.layers.Dense(1)\n", - " \n", - " def call(self, features, hidden):\n", - " # features(CNN_encoder output) shape == (batch_size, 64, embedding_dim)\n", - " \n", - " # hidden shape == (batch_size, hidden_size)\n", - " # hidden_with_time_axis shape == (batch_size, 1, hidden_size)\n", - " hidden_with_time_axis = tf.expand_dims(hidden, 1)\n", - " \n", - " # score shape == (batch_size, 64, hidden_size)\n", - " score = tf.nn.tanh(self.W1(features) + self.W2(hidden_with_time_axis))\n", - " \n", - " # attention_weights shape == (batch_size, 64, 1)\n", - " # we get 1 at the last axis because we are applying score to self.V\n", - " attention_weights = tf.nn.softmax(self.V(score), axis=1)\n", - " \n", - " # context_vector shape after sum == (batch_size, hidden_size)\n", - " context_vector = attention_weights * features\n", - " context_vector = tf.reduce_sum(context_vector, axis=1)\n", - " \n", - " return context_vector, attention_weights" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "AZ7R1RxHRPGf" - }, - "outputs": [], - "source": [ - "class CNN_Encoder(tf.keras.Model):\n", - " # Since we have already extracted the features and dumped it using pickle\n", - " # This encoder passes those features through a Fully connected layer\n", - " def __init__(self, embedding_dim):\n", - " super(CNN_Encoder, self).__init__()\n", - " # shape after fc == (batch_size, 64, embedding_dim)\n", - " self.fc = tf.keras.layers.Dense(embedding_dim)\n", - " \n", - " def call(self, x):\n", - " x = self.fc(x)\n", - " x = tf.nn.relu(x)\n", - " return x" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "V9UbGQmERPGi" - }, - "outputs": [], - "source": [ - "class RNN_Decoder(tf.keras.Model):\n", - " def __init__(self, embedding_dim, units, vocab_size):\n", - " super(RNN_Decoder, self).__init__()\n", - " self.units = units\n", - "\n", - " self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)\n", - " self.gru = gru(self.units)\n", - " self.fc1 = tf.keras.layers.Dense(self.units)\n", - " self.fc2 = tf.keras.layers.Dense(vocab_size)\n", - " \n", - " self.attention = BahdanauAttention(self.units)\n", - " \n", - " def call(self, x, features, hidden):\n", - " # defining attention as a separate model\n", - " context_vector, attention_weights = self.attention(features, hidden)\n", - " \n", - " # x shape after passing through embedding == (batch_size, 1, embedding_dim)\n", - " x = self.embedding(x)\n", - " \n", - " # x shape after concatenation == (batch_size, 1, embedding_dim + hidden_size)\n", - " x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)\n", - " \n", - " # passing the concatenated vector to the GRU\n", - " output, state = self.gru(x)\n", - " \n", - " # shape == (batch_size, max_length, hidden_size)\n", - " x = self.fc1(output)\n", - " \n", - " # x shape == (batch_size * max_length, hidden_size)\n", - " x = tf.reshape(x, (-1, x.shape[2]))\n", - " \n", - " # output shape == (batch_size * max_length, vocab)\n", - " x = self.fc2(x)\n", - "\n", - " return x, state, attention_weights\n", - "\n", - " def reset_state(self, batch_size):\n", - " return tf.zeros((batch_size, self.units))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "Qs_Sr03wRPGk" - }, - "outputs": [], - "source": [ - "encoder = CNN_Encoder(embedding_dim)\n", - "decoder = RNN_Decoder(embedding_dim, units, vocab_size)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "-bYN7xA0RPGl" - }, - "outputs": [], - "source": [ - "optimizer = tf.train.AdamOptimizer()\n", - "\n", - "# We are masking the loss calculated for padding\n", - "def loss_function(real, pred):\n", - " mask = 1 - np.equal(real, 0)\n", - " loss_ = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=real, logits=pred) * mask\n", - " return tf.reduce_mean(loss_)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "PHod7t72RPGn" - }, - "source": [ - "## Training\n", - "\n", - "* We extract the features stored in the respective `.npy` files and then pass those features through the encoder.\n", - "* The encoder output, hidden state(initialized to 0) and the decoder input (which is the start token) is passed to the decoder.\n", - "* The decoder returns the predictions and the decoder hidden state.\n", - "* The decoder hidden state is then passed back into the model and the predictions are used to calculate the loss.\n", - "* Use teacher forcing to decide the next input to the decoder.\n", - "* Teacher forcing is the technique where the target word is passed as the next input to the decoder.\n", - "* The final step is to calculate the gradients and apply it to the optimizer and backpropagate.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "Vt4WZ5mhJE-E" - }, - "outputs": [], - "source": [ - "# adding this in a separate cell because if you run the training cell \n", - "# many times, the loss_plot array will be reset\n", - "loss_plot = []" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "UlA4VIQpRPGo" - }, - "outputs": [], - "source": [ - "EPOCHS = 20\n", - "\n", - "for epoch in range(EPOCHS):\n", - " start = time.time()\n", - " total_loss = 0\n", - " \n", - " for (batch, (img_tensor, target)) in enumerate(dataset):\n", - " loss = 0\n", - " \n", - " # initializing the hidden state for each batch\n", - " # because the captions are not related from image to image\n", - " hidden = decoder.reset_state(batch_size=target.shape[0])\n", - "\n", - " dec_input = tf.expand_dims([tokenizer.word_index['']] * BATCH_SIZE, 1)\n", - " \n", - " with tf.GradientTape() as tape:\n", - " features = encoder(img_tensor)\n", - " \n", - " for i in range(1, target.shape[1]):\n", - " # passing the features through the decoder\n", - " predictions, hidden, _ = decoder(dec_input, features, hidden)\n", - "\n", - " loss += loss_function(target[:, i], predictions)\n", - " \n", - " # using teacher forcing\n", - " dec_input = tf.expand_dims(target[:, i], 1)\n", - " \n", - " total_loss += (loss / int(target.shape[1]))\n", - " \n", - " variables = encoder.variables + decoder.variables\n", - " \n", - " gradients = tape.gradient(loss, variables) \n", - " \n", - " optimizer.apply_gradients(zip(gradients, variables), tf.train.get_or_create_global_step())\n", - " \n", - " if batch % 100 == 0:\n", - " print ('Epoch {} Batch {} Loss {:.4f}'.format(epoch + 1, \n", - " batch, \n", - " loss.numpy() / int(target.shape[1])))\n", - " # storing the epoch end loss value to plot later\n", - " loss_plot.append(total_loss / len(cap_vector))\n", - " \n", - " print ('Epoch {} Loss {:.6f}'.format(epoch + 1, \n", - " total_loss/len(cap_vector)))\n", - " print ('Time taken for 1 epoch {} sec\\n'.format(time.time() - start))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "1Wm83G-ZBPcC" - }, - "outputs": [], - "source": [ - "plt.plot(loss_plot)\n", - "plt.xlabel('Epochs')\n", - "plt.ylabel('Loss')\n", - "plt.title('Loss Plot')\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "xGvOcLQKghXN" - }, - "source": [ - "## Caption!\n", - "\n", - "* The evaluate function is similar to the training loop, except we don't use teacher forcing here. The input to the decoder at each time step is its previous predictions along with the hidden state and the encoder output.\n", - "* Stop predicting when the model predicts the end token.\n", - "* And store the attention weights for every time step." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "RCWpDtyNRPGs" - }, - "outputs": [], - "source": [ - "def evaluate(image):\n", - " attention_plot = np.zeros((max_length, attention_features_shape))\n", - "\n", - " hidden = decoder.reset_state(batch_size=1)\n", - "\n", - " temp_input = tf.expand_dims(load_image(image)[0], 0)\n", - " img_tensor_val = image_features_extract_model(temp_input)\n", - " img_tensor_val = tf.reshape(img_tensor_val, (img_tensor_val.shape[0], -1, img_tensor_val.shape[3]))\n", - "\n", - " features = encoder(img_tensor_val)\n", - "\n", - " dec_input = tf.expand_dims([tokenizer.word_index['']], 0)\n", - " result = []\n", - "\n", - " for i in range(max_length):\n", - " predictions, hidden, attention_weights = decoder(dec_input, features, hidden)\n", - "\n", - " attention_plot[i] = tf.reshape(attention_weights, (-1, )).numpy()\n", - "\n", - " predicted_id = tf.argmax(predictions[0]).numpy()\n", - " result.append(tokenizer.index_word[predicted_id])\n", - "\n", - " if tokenizer.index_word[predicted_id] == '':\n", - " return result, attention_plot\n", - "\n", - " dec_input = tf.expand_dims([predicted_id], 0)\n", - "\n", - " attention_plot = attention_plot[:len(result), :]\n", - " return result, attention_plot" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - }, - "colab_type": "code", - "id": "fD_y7PD6RPGt" - }, - "outputs": [], - "source": [ - "def plot_attention(image, result, attention_plot):\n", - " temp_image = np.array(Image.open(image))\n", - "\n", - " fig = plt.figure(figsize=(10, 10))\n", - " \n", - " len_result = len(result)\n", - " for l in range(len_result):\n", - " temp_att = np.resize(attention_plot[l], (8, 8))\n", - " ax = fig.add_subplot(len_result//2, len_result//2, l+1)\n", - " ax.set_title(result[l])\n", - " img = ax.imshow(temp_image)\n", - " ax.imshow(temp_att, cmap='gray', alpha=0.6, extent=img.get_extent())\n", - "\n", - " plt.tight_layout()\n", - " plt.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "colab_type": "text", + "id": "K2s1A9eLRPEj" + }, + "source": [ + "##### Copyright 2018 The TensorFlow Authors.\n", + "\n", + "Licensed under the Apache License, Version 2.0 (the \"License\").\n" + ] }, - "colab_type": "code", - "id": "io7ws3ReRPGv" - }, - "outputs": [], - "source": [ - "# captions on the validation set\n", - "rid = np.random.randint(0, len(img_name_val))\n", - "image = img_name_val[rid]\n", - "real_caption = ' '.join([tokenizer.index_word[i] for i in cap_val[rid] if i not in [0]])\n", - "result, attention_plot = evaluate(image)\n", - "\n", - "print ('Real Caption:', real_caption)\n", - "print ('Prediction Caption:', ' '.join(result))\n", - "plot_attention(image, result, attention_plot)\n", - "# opening the image\n", - "Image.open(img_name_val[rid])" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "Rprk3HEvZuxb" - }, - "source": [ - "## Try it on your own images\n", - "For fun, below we've provided a method you can use to caption your own images with the model we've just trained. Keep in mind, it was trained on a relatively small amount of data, and your images may be different from the training data (so be prepared for weird results!)\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } + { + "cell_type": "markdown", + "metadata": { + "colab_type": "text", + "id": "Cffg2i257iMS" + }, + "source": [ + "# Image Captioning with Attention\n", + "\n", + "This example has moved:\n", + "\n", + "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\u003ctd\u003e\n", + "\u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/r2/tutorials/generative/image_captioning.ipynb\"\u003e\n", + " \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e \n", + "\u003c/td\u003e\u003ctd\u003e\n", + "\u003ca target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/r2/tutorials/generative/image_captioning.ipynb\"\u003e\u003cimg width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\u003c/td\u003e\u003c/table\u003e" + ] }, - "colab_type": "code", - "id": "9Psd1quzaAWg" - }, - "outputs": [], - "source": [ - "image_url = 'https://tensorflow.org/images/surf.jpg'\n", - "image_extension = image_url[-4:]\n", - "image_path = tf.keras.utils.get_file('image'+image_extension, \n", - " origin=image_url)\n", - "\n", - "result, attention_plot = evaluate(image_path)\n", - "print ('Prediction Caption:', ' '.join(result))\n", - "plot_attention(image_path, result, attention_plot)\n", - "# opening the image\n", - "Image.open(image_path)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "VJZXyJco6uLO" - }, - "source": [ - "# Next steps\n", - "\n", - "Congrats! You've just trained an image captioning model with attention. Next, we recommend taking a look at this example [Neural Machine Translation with Attention]( https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/eager/python/examples/nmt_with_attention/nmt_with_attention.ipynb). It uses a similar architecture to translate between Spanish and English sentences. You can also experiment with training the code in this notebook on a different dataset." - ] - } - ], - "metadata": { - "accelerator": "GPU", - "colab": { - "collapsed_sections": [], - "default_view": {}, - "name": "image_captioning_with_attention.ipynb", - "private_outputs": true, - "provenance": [ { - "file_id": "1HI8OK2sMjcx9CTWVn0122QAHOuXaOaMg", - "timestamp": 1530222436922 + "cell_type": "markdown", + "metadata": { + "colab_type": "text", + "id": "QASbY_HGo4Lq" + }, + "source": [ + "![Man Surfing](https://tensorflow.org/images/surf.jpg) \n", + "\n", + "[Image Source](https://commons.wikimedia.org/wiki/Surfing#/media/File:Surfing_in_Hawaii.jpg), License: Public Domain\n", + "\n", + "![Prediction](https://tensorflow.org/images/imcap_prediction.png)\n" + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "collapsed_sections": [], + "name": "image_captioning_with_attention.ipynb", + "private_outputs": true, + "provenance": [ + { + "file_id": "1HI8OK2sMjcx9CTWVn0122QAHOuXaOaMg", + "timestamp": 1530222436922 + } + ], + "toc_visible": true, + "version": "0.3.2" + }, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" } - ], - "toc_visible": true, - "version": "0.3.2", - "views": {} - }, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.3" - } - }, - "nbformat": 4, - "nbformat_minor": 2 + "nbformat": 4, + "nbformat_minor": 0 } diff --git a/tensorflow/contrib/eager/python/examples/generative_examples/text_generation.ipynb b/tensorflow/contrib/eager/python/examples/generative_examples/text_generation.ipynb index bda9e77085e45ae31a228142135425e22a1c6780..c945c753b3ba36d16aa6985d23a5849f8f552304 100644 --- a/tensorflow/contrib/eager/python/examples/generative_examples/text_generation.ipynb +++ b/tensorflow/contrib/eager/python/examples/generative_examples/text_generation.ipynb @@ -13,633 +13,13 @@ "\n", "# Text Generation using a RNN\n", "\n", + "This example has moved.\n", + "\n", "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\u003ctd\u003e\n", - "\u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/contrib/eager/python/examples/generative_examples/text_generation.ipynb\"\u003e\n", + "\u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/sequences/text_generation.ipynb\"\u003e\n", " \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e \n", "\u003c/td\u003e\u003ctd\u003e\n", - "\u003ca target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/examples/generative_examples/text_generation.ipynb\"\u003e\u003cimg width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on Github\u003c/a\u003e\u003c/td\u003e\u003c/table\u003e" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "BwpJ5IffzRG6" - }, - "source": [ - "This notebook demonstrates how to generate text using an RNN using [tf.keras](https://www.tensorflow.org/programmers_guide/keras) and [eager execution](https://www.tensorflow.org/programmers_guide/eager). If you like, you can write a similar [model](https://github.com/fchollet/deep-learning-with-python-notebooks/blob/master/8.1-text-generation-with-lstm.ipynb) using less code. Here, we show a lower-level impementation that's useful to understand as prework before diving in to deeper examples in a similar, like [Neural Machine Translation with Attention](https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/contrib/eager/python/examples/nmt_with_attention/nmt_with_attention.ipynb).\n", - "\n", - "This notebook is an end-to-end example. When you run it, it will download a dataset of Shakespeare's writing. We'll use a collection of plays, borrowed from Andrej Karpathy's excellent [The Unreasonable Effectiveness of Recurrent Neural Networks](http://karpathy.github.io/2015/05/21/rnn-effectiveness/). The notebook will train a model, and use it to generate sample output.\n", - " \n", - "Here is the output(with start string='w') after training a single layer GRU for 30 epochs with the default settings below:\n", - "\n", - "```\n", - "were to the death of him\n", - "And nothing of the field in the view of hell,\n", - "When I said, banish him, I will not burn thee that would live.\n", - "\n", - "HENRY BOLINGBROKE:\n", - "My gracious uncle--\n", - "\n", - "DUKE OF YORK:\n", - "As much disgraced to the court, the gods them speak,\n", - "And now in peace himself excuse thee in the world.\n", - "\n", - "HORTENSIO:\n", - "Madam, 'tis not the cause of the counterfeit of the earth,\n", - "And leave me to the sun that set them on the earth\n", - "And leave the world and are revenged for thee.\n", - "\n", - "GLOUCESTER:\n", - "I would they were talking with the very name of means\n", - "To make a puppet of a guest, and therefore, good Grumio,\n", - "Nor arm'd to prison, o' the clouds, of the whole field,\n", - "With the admire\n", - "With the feeding of thy chair, and we have heard it so,\n", - "I thank you, sir, he is a visor friendship with your silly your bed.\n", - "\n", - "SAMPSON:\n", - "I do desire to live, I pray: some stand of the minds, make thee remedies\n", - "With the enemies of my soul.\n", - "\n", - "MENENIUS:\n", - "I'll keep the cause of my mistress.\n", - "\n", - "POLIXENES:\n", - "My brother Marcius!\n", - "\n", - "Second Servant:\n", - "Will't ple\n", - "```\n", - "\n", - "Of course, while some of the sentences are grammatical, most do not make sense. But, consider:\n", - "\n", - "* Our model is character based (when we began training, it did not yet know how to spell a valid English word, or that words were even a unit of text).\n", - "\n", - "* The structure of the output resembles a play (blocks begin with a speaker name, in all caps similar to the original text). Sentences generally end with a period. If you look at the text from a distance (or don't read the invididual words too closely, it appears as if it's an excerpt from a play).\n", - "\n", - "As a next step, you can experiment training the model on a different dataset - any large text file(ASCII) will do, and you can modify a single line of code below to make that change. Have fun!\n" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "R3p22DBDsaCA" - }, - "source": [ - "## Install unidecode library\n", - "A helpful library to convert unicode to ASCII." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "wZ6LOM12wKGH" - }, - "outputs": [], - "source": [ - "!pip install unidecode" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "WGyKZj3bzf9p" - }, - "source": [ - "## Import tensorflow and enable eager execution." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "yG_n40gFzf9s" - }, - "outputs": [], - "source": [ - "# Import TensorFlow \u003e= 1.10 and enable eager execution\n", - "import tensorflow as tf\n", - "\n", - "# Note: Once you enable eager execution, it cannot be disabled. \n", - "tf.enable_eager_execution()\n", - "\n", - "import numpy as np\n", - "import os\n", - "import re\n", - "import random\n", - "import unidecode\n", - "import time" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "EHDoRoc5PKWz" - }, - "source": [ - "## Download the dataset\n", - "\n", - "In this example, we will use the [shakespeare dataset](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt). You can use any other dataset that you like.\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "pD_55cOxLkAb" - }, - "outputs": [], - "source": [ - "path_to_file = tf.keras.utils.get_file('shakespeare.txt', 'https://storage.googleapis.com/download.tensorflow.org/data/shakespeare.txt')" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "UHjdCjDuSvX_" - }, - "source": [ - "## Read the dataset\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "-E5JvY3wzf94" - }, - "outputs": [], - "source": [ - "text = unidecode.unidecode(open(path_to_file).read())\n", - "# length of text is the number of characters in it\n", - "print (len(text))" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "Il9ww98izf-D" - }, - "source": [ - "Creating dictionaries to map from characters to their indices and vice-versa, which will be used to vectorize the inputs" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "IalZLbvOzf-F" - }, - "outputs": [], - "source": [ - "# unique contains all the unique characters in the file\n", - "unique = sorted(set(text))\n", - "\n", - "# creating a mapping from unique characters to indices\n", - "char2idx = {u:i for i, u in enumerate(unique)}\n", - "idx2char = {i:u for i, u in enumerate(unique)}" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "1v_qUYfAzf-I" - }, - "outputs": [], - "source": [ - "# setting the maximum length sentence we want for a single input in characters\n", - "max_length = 100\n", - "\n", - "# length of the vocabulary in chars\n", - "vocab_size = len(unique)\n", - "\n", - "# the embedding dimension \n", - "embedding_dim = 256\n", - "\n", - "# number of RNN (here GRU) units\n", - "units = 1024\n", - "\n", - "# batch size \n", - "BATCH_SIZE = 64\n", - "\n", - "# buffer size to shuffle our dataset\n", - "BUFFER_SIZE = 10000" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "LFjSVAlWzf-N" - }, - "source": [ - "## Creating the input and output tensors\n", - "\n", - "Vectorizing the input and the target text because our model cannot understand strings only numbers.\n", - "\n", - "But first, we need to create the input and output vectors.\n", - "Remember the max_length we set above, we will use it here. We are creating **max_length** chunks of input, where each input vector is all the characters in that chunk except the last and the target vector is all the characters in that chunk except the first.\n", - "\n", - "For example, consider that the string = 'tensorflow' and the max_length is 9\n", - "\n", - "So, the `input = 'tensorflo'` and `output = 'ensorflow'`\n", - "\n", - "After creating the vectors, we convert each character into numbers using the **char2idx** dictionary we created above." - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "0UHJDA39zf-O" - }, - "outputs": [], - "source": [ - "input_text = []\n", - "target_text = []\n", - "\n", - "for f in range(0, len(text)-max_length, max_length):\n", - " inps = text[f:f+max_length]\n", - " targ = text[f+1:f+1+max_length]\n", - "\n", - " input_text.append([char2idx[i] for i in inps])\n", - " target_text.append([char2idx[t] for t in targ])\n", - " \n", - "print (np.array(input_text).shape)\n", - "print (np.array(target_text).shape)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "MJdfPmdqzf-R" - }, - "source": [ - "## Creating batches and shuffling them using tf.data" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "p2pGotuNzf-S" - }, - "outputs": [], - "source": [ - "dataset = tf.data.Dataset.from_tensor_slices((input_text, target_text)).shuffle(BUFFER_SIZE)\n", - "dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "m8gPwEjRzf-Z" - }, - "source": [ - "## Creating the model\n", - "\n", - "We use the Model Subclassing API which gives us full flexibility to create the model and change it however we like. We use 3 layers to define our model.\n", - "\n", - "* Embedding layer\n", - "* GRU layer (you can use an LSTM layer here)\n", - "* Fully connected layer" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "P3KTiiInzf-a" - }, - "outputs": [], - "source": [ - "class Model(tf.keras.Model):\n", - " def __init__(self, vocab_size, embedding_dim, units, batch_size):\n", - " super(Model, self).__init__()\n", - " self.units = units\n", - " self.batch_sz = batch_size\n", - "\n", - " self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)\n", - "\n", - " if tf.test.is_gpu_available():\n", - " self.gru = tf.keras.layers.CuDNNGRU(self.units, \n", - " return_sequences=True, \n", - " return_state=True, \n", - " recurrent_initializer='glorot_uniform')\n", - " else:\n", - " self.gru = tf.keras.layers.GRU(self.units, \n", - " return_sequences=True, \n", - " return_state=True, \n", - " recurrent_activation='sigmoid', \n", - " recurrent_initializer='glorot_uniform')\n", - "\n", - " self.fc = tf.keras.layers.Dense(vocab_size)\n", - " \n", - " def call(self, x, hidden):\n", - " x = self.embedding(x)\n", - "\n", - " # output shape == (batch_size, max_length, hidden_size) \n", - " # states shape == (batch_size, hidden_size)\n", - "\n", - " # states variable to preserve the state of the model\n", - " # this will be used to pass at every step to the model while training\n", - " output, states = self.gru(x, initial_state=hidden)\n", - "\n", - "\n", - " # reshaping the output so that we can pass it to the Dense layer\n", - " # after reshaping the shape is (batch_size * max_length, hidden_size)\n", - " output = tf.reshape(output, (-1, output.shape[2]))\n", - "\n", - " # The dense layer will output predictions for every time_steps(max_length)\n", - " # output shape after the dense layer == (max_length * batch_size, vocab_size)\n", - " x = self.fc(output)\n", - "\n", - " return x, states" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "trpqTWyvk0nr" - }, - "source": [ - "## Call the model and set the optimizer and the loss function" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "7t2XrzEOzf-e" - }, - "outputs": [], - "source": [ - "model = Model(vocab_size, embedding_dim, units, BATCH_SIZE)" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "dkjWIATszf-h" - }, - "outputs": [], - "source": [ - "optimizer = tf.train.AdamOptimizer()\n", - "\n", - "# using sparse_softmax_cross_entropy so that we don't have to create one-hot vectors\n", - "def loss_function(real, preds):\n", - " return tf.losses.sparse_softmax_cross_entropy(labels=real, logits=preds)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "3K6s6F79P7za" - }, - "source": [ - "## Checkpoints (Object-based saving)" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "oAGisDdfP9rL" - }, - "outputs": [], - "source": [ - "checkpoint_dir = './training_checkpoints'\n", - "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\n", - "checkpoint = tf.train.Checkpoint(optimizer=optimizer,\n", - " model=model)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "lPrP0XMUzf-p" - }, - "source": [ - "## Train the model\n", - "\n", - "Here we will use a custom training loop with the help of GradientTape()\n", - "\n", - "* We initialize the hidden state of the model with zeros and shape == (batch_size, number of rnn units). We do this by calling the function defined while creating the model.\n", - "\n", - "* Next, we iterate over the dataset(batch by batch) and calculate the **predictions and the hidden states** associated with that input.\n", - "\n", - "* There are a lot of interesting things happening here.\n", - " * The model gets hidden state(initialized with 0), lets call that **H0** and the first batch of input, lets call that **I0**.\n", - " * The model then returns the predictions **P1** and **H1**.\n", - " * For the next batch of input, the model receives **I1** and **H1**.\n", - " * The interesting thing here is that we pass **H1** to the model with **I1** which is how the model learns. The context learned from batch to batch is contained in the **hidden state**.\n", - " * We continue doing this until the dataset is exhausted and then we start a new epoch and repeat this.\n", - "\n", - "* After calculating the predictions, we calculate the **loss** using the loss function defined above. Then we calculate the gradients of the loss with respect to the model variables(input)\n", - "\n", - "* Finally, we take a step in that direction with the help of the optimizer using the apply_gradients function.\n", - "\n", - "Note:- If you are running this notebook in Colab which has a **Tesla K80 GPU** it takes about 23 seconds per epoch.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "d4tSNwymzf-q" - }, - "outputs": [], - "source": [ - "# Training step\n", - "\n", - "EPOCHS = 20\n", - "\n", - "for epoch in range(EPOCHS):\n", - " start = time.time()\n", - " \n", - " # initializing the hidden state at the start of every epoch\n", - " hidden = model.reset_states()\n", - " \n", - " for (batch, (inp, target)) in enumerate(dataset):\n", - " with tf.GradientTape() as tape:\n", - " # feeding the hidden state back into the model\n", - " # This is the interesting step\n", - " predictions, hidden = model(inp, hidden)\n", - " \n", - " # reshaping the target because that's how the \n", - " # loss function expects it\n", - " target = tf.reshape(target, (-1,))\n", - " loss = loss_function(target, predictions)\n", - " \n", - " grads = tape.gradient(loss, model.variables)\n", - " optimizer.apply_gradients(zip(grads, model.variables))\n", - "\n", - " if batch % 100 == 0:\n", - " print ('Epoch {} Batch {} Loss {:.4f}'.format(epoch+1,\n", - " batch,\n", - " loss))\n", - " # saving (checkpoint) the model every 5 epochs\n", - " if (epoch + 1) % 5 == 0:\n", - " checkpoint.save(file_prefix = checkpoint_prefix)\n", - "\n", - " print ('Epoch {} Loss {:.4f}'.format(epoch+1, loss))\n", - " print('Time taken for 1 epoch {} sec\\n'.format(time.time() - start))" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "01AR9vpNQMFF" - }, - "source": [ - "## Restore the latest checkpoint" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "tyvpYomYQQkF" - }, - "outputs": [], - "source": [ - "# restoring the latest checkpoint in checkpoint_dir\n", - "checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "DjGz1tDkzf-u" - }, - "source": [ - "## Predicting using our trained model\n", - "\n", - "The below code block is used to generated the text\n", - "\n", - "* We start by choosing a start string and initializing the hidden state and setting the number of characters we want to generate.\n", - "\n", - "* We get predictions using the start_string and the hidden state\n", - "\n", - "* Then we use argmax to calculate the index of the predicted word. **We use this predicted word as our next input to the model**\n", - "\n", - "* **The hidden state returned by the model is fed back into the model so that it now has more context rather than just one word.** After we predict the next word, the modified hidden states are again fed back into the model, which is how it learns as it gets more context from the previously predicted words.\n", - "\n", - "* If you see the predictions, the model knows when to capitalize, make paragraphs and the text follows a shakespeare style of writing which is pretty awesome!" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "WvuwZBX5Ogfd" - }, - "outputs": [], - "source": [ - "# Evaluation step(generating text using the model learned)\n", - "\n", - "# number of characters to generate\n", - "num_generate = 1000\n", - "\n", - "# You can change the start string to experiment\n", - "start_string = 'Q'\n", - "# converting our start string to numbers(vectorizing!) \n", - "input_eval = [char2idx[s] for s in start_string]\n", - "input_eval = tf.expand_dims(input_eval, 0)\n", - "\n", - "# empty string to store our results\n", - "text_generated = ''\n", - "\n", - "# hidden state shape == (batch_size, number of rnn units); here batch size == 1\n", - "hidden = [tf.zeros((1, units))]\n", - "for i in range(num_generate):\n", - " predictions, hidden = model(input_eval, hidden)\n", - "\n", - " # using argmax to predict the word returned by the model\n", - " predicted_id = tf.argmax(predictions[-1]).numpy()\n", - " \n", - " # We pass the predicted word as the next input to the model\n", - " # along with the previous hidden state\n", - " input_eval = tf.expand_dims([predicted_id], 0)\n", - " \n", - " text_generated += idx2char[predicted_id]\n", - "\n", - "print (start_string + text_generated)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "colab_type": "text", - "id": "AM2Uma_-yVIq" - }, - "source": [ - "## Next steps\n", - "\n", - "* Change the start string to a different character, or the start of a sentence.\n", - "* Experiment with training on a different, or with different parameters. [Project Gutenberg](http://www.gutenberg.org/ebooks/100), for example, contains a large collection of books.\n", - "* Add another RNN layer.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 0, - "metadata": { - "colab": {}, - "colab_type": "code", - "id": "gtEd86sX5cB2" - }, - "outputs": [], - "source": [ - "" + "\u003ca target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/tutorials/sequences/text_generation.ipynb\"\u003e\u003cimg width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\u003c/td\u003e\u003c/table\u003e" ] } ], diff --git a/tensorflow/contrib/eager/python/examples/l2hmc/BUILD b/tensorflow/contrib/eager/python/examples/l2hmc/BUILD index 7bdf9053de749af9d09b12ba7b848e21c1fdb8f0..35d509904211d98f124d2555fc48166e75cb0dd9 100644 --- a/tensorflow/contrib/eager/python/examples/l2hmc/BUILD +++ b/tensorflow/contrib/eager/python/examples/l2hmc/BUILD @@ -28,7 +28,7 @@ py_library( cuda_py_test( name = "l2hmc_test", - size = "large", + size = "medium", srcs = ["l2hmc_test.py"], additional_deps = [ ":l2hmc", @@ -36,4 +36,8 @@ cuda_py_test( "//tensorflow/contrib/eager/python:tfe", "//third_party/py/numpy", ], + shard_count = 4, + tags = [ + "oss_serial", + ], ) diff --git a/tensorflow/contrib/eager/python/examples/linear_regression/BUILD b/tensorflow/contrib/eager/python/examples/linear_regression/BUILD index 74ce9e84f013d79b3a33ffa79993980b561e366d..30afef83bc5c6c164c8456ed472f4d6064068a25 100644 --- a/tensorflow/contrib/eager/python/examples/linear_regression/BUILD +++ b/tensorflow/contrib/eager/python/examples/linear_regression/BUILD @@ -9,6 +9,13 @@ py_binary( name = "linear_regression", srcs = ["linear_regression.py"], srcs_version = "PY2AND3", + deps = [":linear_regression_lib"], +) + +py_library( + name = "linear_regression_lib", + srcs = ["linear_regression.py"], + srcs_version = "PY2AND3", deps = [ "//tensorflow:tensorflow_py", "//tensorflow/contrib/eager/python:tfe", @@ -20,10 +27,13 @@ cuda_py_test( size = "small", srcs = ["linear_regression_test.py"], additional_deps = [ - ":linear_regression", + ":linear_regression_lib", "//tensorflow:tensorflow_py", ], - tags = ["no_windows"], # TODO: needs investigation on Windows + tags = [ + "no_windows", # TODO: needs investigation on Windows + "oss_serial", + ], ) cuda_py_test( @@ -31,7 +41,7 @@ cuda_py_test( size = "small", srcs = ["linear_regression_graph_test.py"], additional_deps = [ - ":linear_regression", + ":linear_regression_lib", "//tensorflow:tensorflow_py", ], ) diff --git a/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression.py b/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression.py index 099b712fc06d1d3eb9ab4095f8db7283690bda76..206ef9409df7b1dc21de42ba919d2ba97f334a8c 100644 --- a/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression.py +++ b/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression.py @@ -56,7 +56,7 @@ class LinearModel(tf.keras.Model): def mean_square_loss(model, xs, ys): - return tf.reduce_mean(tf.square(tf.subtract(model(xs), ys))) + return tf.reduce_mean(tf.squared_difference(model(xs), ys)) def fit(model, dataset, optimizer, verbose=False, logdir=None): 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 66d52a74943d0d81fde05ce51b019558b327978d..436e887736158ec1ba8e46eac8de4ac7b8e6be01 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 @@ -1,11 +1,28 @@ { + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "name": "nmt_with_attention.ipynb", + "version": "0.3.2", + "provenance": [], + "private_outputs": true, + "collapsed_sections": [], + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "accelerator": "GPU" + }, "cells": [ { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "AOpGoE2T-YXS" }, + "cell_type": "markdown", "source": [ "##### Copyright 2018 The TensorFlow Authors.\n", "\n", @@ -13,19 +30,19 @@ "\n", "# Neural Machine Translation with Attention\n", "\n", - "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\u003ctd\u003e\n", - "\u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/contrib/eager/python/examples/nmt_with_attention/nmt_with_attention.ipynb\"\u003e\n", - " \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e \n", - "\u003c/td\u003e\u003ctd\u003e\n", - "\u003ca target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/examples/nmt_with_attention/nmt_with_attention.ipynb\"\u003e\u003cimg width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\u003c/td\u003e\u003c/table\u003e" + "
\n", + "\n", + " Run in Google Colab \n", + "\n", + "View source on GitHub
" ] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "CiwtNgENbx2g" }, + "cell_type": "markdown", "source": [ "This notebook trains a sequence to sequence (seq2seq) model for Spanish to English translation using [tf.keras](https://www.tensorflow.org/programmers_guide/keras) and [eager execution](https://www.tensorflow.org/programmers_guide/eager). This is an advanced example that assumes some knowledge of sequence to sequence models.\n", "\n", @@ -33,24 +50,22 @@ "\n", "The translation quality is reasonable for a toy example, but the generated attention plot is perhaps more interesting. This shows which parts of the input sentence has the model's attention while translating:\n", "\n", - "\u003cimg src=\"https://tensorflow.org/images/spanish-english.png\" alt=\"spanish-english attention plot\"\u003e\n", + "\"spanish-english\n", "\n", "Note: This example takes approximately 10 mintues to run on a single P100 GPU." ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "tnxXKDjq3jEL" + "id": "tnxXKDjq3jEL", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "from __future__ import absolute_import, division, print_function\n", "\n", - "# Import TensorFlow \u003e= 1.10 and enable eager execution\n", + "# Import TensorFlow >= 1.10 and enable eager execution\n", "import tensorflow as tf\n", "\n", "tf.enable_eager_execution()\n", @@ -65,14 +80,16 @@ "import time\n", "\n", "print(tf.__version__)" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "wfodePkj3jEa" }, + "cell_type": "markdown", "source": [ "## Download and prepare the dataset\n", "\n", @@ -91,14 +108,12 @@ ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "kRVATYOgJs1b" + "id": "kRVATYOgJs1b", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "# Download the file\n", "path_to_zip = tf.keras.utils.get_file(\n", @@ -106,17 +121,17 @@ " extract=True)\n", "\n", "path_to_file = os.path.dirname(path_to_zip)+\"/spa-eng/spa.txt\"" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "rd0jw-eC3jEh" + "id": "rd0jw-eC3jEh", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "# Converts the unicode file to ascii\n", "def unicode_to_ascii(s):\n", @@ -128,7 +143,7 @@ " w = unicode_to_ascii(w.lower().strip())\n", " \n", " # creating a space between a word and the punctuation following it\n", - " # eg: \"he is a boy.\" =\u003e \"he is a boy .\" \n", + " # eg: \"he is a boy.\" => \"he is a boy .\" \n", " # Reference:- https://stackoverflow.com/questions/3645931/python-padding-punctuation-with-white-spaces-keeping-punctuation\n", " w = re.sub(r\"([?.!,¿])\", r\" \\1 \", w)\n", " w = re.sub(r'[\" \"]+', \" \", w)\n", @@ -140,19 +155,19 @@ " \n", " # adding a start and an end token to the sentence\n", " # so that the model know when to start and stop predicting.\n", - " w = '\u003cstart\u003e ' + w + ' \u003cend\u003e'\n", + " w = ' ' + w + ' '\n", " return w" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "OHn4Dct23jEm" + "id": "OHn4Dct23jEm", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "# 1. Remove the accents\n", "# 2. Clean the sentences\n", @@ -163,20 +178,20 @@ " word_pairs = [[preprocess_sentence(w) for w in l.split('\\t')] for l in lines[:num_examples]]\n", " \n", " return word_pairs" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "9xbqO7Iie9bb" + "id": "9xbqO7Iie9bb", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ - "# This class creates a word -\u003e index mapping (e.g,. \"dad\" -\u003e 5) and vice-versa \n", - "# (e.g., 5 -\u003e \"dad\") for each language,\n", + "# This class creates a word -> index mapping (e.g,. \"dad\" -> 5) and vice-versa \n", + "# (e.g., 5 -> \"dad\") for each language,\n", "class LanguageIndex():\n", " def __init__(self, lang):\n", " self.lang = lang\n", @@ -192,23 +207,23 @@ " \n", " self.vocab = sorted(self.vocab)\n", " \n", - " self.word2idx['\u003cpad\u003e'] = 0\n", + " self.word2idx[''] = 0\n", " for index, word in enumerate(self.vocab):\n", " self.word2idx[word] = index + 1\n", " \n", " for word, index in self.word2idx.items():\n", " self.idx2word[index] = word" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "eAY9k49G3jE_" + "id": "eAY9k49G3jE_", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "def max_length(tensor):\n", " return max(len(t) for t in tensor)\n", @@ -244,71 +259,71 @@ " padding='post')\n", " \n", " return input_tensor, target_tensor, inp_lang, targ_lang, max_length_inp, max_length_tar" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "GOi42V79Ydlr" }, + "cell_type": "markdown", "source": [ "### Limit the size of the dataset to experiment faster (optional)\n", "\n", - "Training on the complete dataset of \u003e100,000 sentences will take a long time. To train faster, we can limit the size of the dataset to 30,000 sentences (of course, translation quality degrades with less data):" + "Training on the complete dataset of >100,000 sentences will take a long time. To train faster, we can limit the size of the dataset to 30,000 sentences (of course, translation quality degrades with less data):" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "cnxC7q-j3jFD" + "id": "cnxC7q-j3jFD", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "# Try experimenting with the size of that dataset\n", "num_examples = 30000\n", "input_tensor, target_tensor, inp_lang, targ_lang, max_length_inp, max_length_targ = load_dataset(path_to_file, num_examples)" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "4QILQkOs3jFG" + "id": "4QILQkOs3jFG", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "# Creating training and validation sets using an 80-20 split\n", "input_tensor_train, input_tensor_val, target_tensor_train, target_tensor_val = train_test_split(input_tensor, target_tensor, test_size=0.2)\n", "\n", "# Show length\n", "len(input_tensor_train), len(target_tensor_train), len(input_tensor_val), len(target_tensor_val)" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "rgCLkfv5uO3d" }, + "cell_type": "markdown", "source": [ "### Create a tf.data dataset" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "TqHsArVZ3jFS" + "id": "TqHsArVZ3jFS", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "BUFFER_SIZE = len(input_tensor_train)\n", "BATCH_SIZE = 64\n", @@ -320,27 +335,29 @@ "\n", "dataset = tf.data.Dataset.from_tensor_slices((input_tensor_train, target_tensor_train)).shuffle(BUFFER_SIZE)\n", "dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "TNfHIF71ulLu" }, + "cell_type": "markdown", "source": [ "## Write the encoder and decoder model\n", "\n", - "Here, we'll implement an encoder-decoder model with attention which you can read about in the TensorFlow [Neural Machine Translation (seq2seq) tutorial](https://www.tensorflow.org/tutorials/seq2seq). This example uses a more recent set of APIs. This notebook implements the [attention equations](https://www.tensorflow.org/tutorials/seq2seq#background_on_the_attention_mechanism) from the seq2seq tutorial. The following diagram shows that each input words is assigned a weight by the attention mechanism which is then used by the decoder to predict the next word in the sentence.\n", + "Here, we'll implement an encoder-decoder model with attention which you can read about in the TensorFlow [Neural Machine Translation (seq2seq) tutorial](https://github.com/tensorflow/nmt). This example uses a more recent set of APIs. This notebook implements the [attention equations](https://github.com/tensorflow/nmt#background-on-the-attention-mechanism) from the seq2seq tutorial. The following diagram shows that each input word is assigned a weight by the attention mechanism which is then used by the decoder to predict the next word in the sentence.\n", "\n", - "\u003cimg src=\"https://www.tensorflow.org/images/seq2seq/attention_mechanism.jpg\" width=\"500\" alt=\"attention mechanism\"\u003e\n", + "\"attention\n", "\n", "The input is put through an encoder model which gives us the encoder output of shape *(batch_size, max_length, hidden_size)* and the encoder hidden state of shape *(batch_size, hidden_size)*. \n", "\n", "Here are the equations that are implemented:\n", "\n", - "\u003cimg src=\"https://www.tensorflow.org/images/seq2seq/attention_equation_0.jpg\" alt=\"attention equation 0\" width=\"800\"\u003e\n", - "\u003cimg src=\"https://www.tensorflow.org/images/seq2seq/attention_equation_1.jpg\" alt=\"attention equation 1\" width=\"800\"\u003e\n", + "\"attention\n", + "\"attention\n", "\n", "We're using *Bahdanau attention*. Lets decide on notation before writing the simplified form:\n", "\n", @@ -362,14 +379,12 @@ ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "avyJ_4VIUoHb" + "id": "avyJ_4VIUoHb", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "def gru(units):\n", " # If you have a GPU, we recommend using CuDNNGRU(provides a 3x speedup than GRU)\n", @@ -385,17 +400,17 @@ " return_state=True, \n", " recurrent_activation='sigmoid', \n", " recurrent_initializer='glorot_uniform')" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "nZ2rI24i3jFg" + "id": "nZ2rI24i3jFg", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "class Encoder(tf.keras.Model):\n", " def __init__(self, vocab_size, embedding_dim, enc_units, batch_sz):\n", @@ -412,17 +427,17 @@ " \n", " def initialize_hidden_state(self):\n", " return tf.zeros((self.batch_sz, self.enc_units))" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "yJ_B3mhW3jFk" + "id": "yJ_B3mhW3jFk", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "class Decoder(tf.keras.Model):\n", " def __init__(self, vocab_size, embedding_dim, dec_units, batch_sz):\n", @@ -476,41 +491,41 @@ " \n", " def initialize_hidden_state(self):\n", " return tf.zeros((self.batch_sz, self.dec_units))" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "P5UY8wko3jFp" + "id": "P5UY8wko3jFp", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "encoder = Encoder(vocab_inp_size, embedding_dim, units, BATCH_SIZE)\n", "decoder = Decoder(vocab_tar_size, embedding_dim, units, BATCH_SIZE)" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "_ch_71VbIRfK" }, + "cell_type": "markdown", "source": [ "## Define the optimizer and the loss function" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "WmTHr5iV3jFr" + "id": "WmTHr5iV3jFr", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "optimizer = tf.train.AdamOptimizer()\n", "\n", @@ -519,41 +534,43 @@ " mask = 1 - np.equal(real, 0)\n", " loss_ = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=real, logits=pred) * mask\n", " return tf.reduce_mean(loss_)" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "DMVWzzsfNl4e" }, + "cell_type": "markdown", "source": [ "## Checkpoints (Object-based saving)" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "Zj8bXQTgNwrF" + "id": "Zj8bXQTgNwrF", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "checkpoint_dir = './training_checkpoints'\n", "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\n", "checkpoint = tf.train.Checkpoint(optimizer=optimizer,\n", " encoder=encoder,\n", " decoder=decoder)" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "hpObfY22IddU" }, + "cell_type": "markdown", "source": [ "## Training\n", "\n", @@ -567,14 +584,12 @@ ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "ddefjBMa3jF0" + "id": "ddefjBMa3jF0", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "EPOCHS = 10\n", "\n", @@ -592,7 +607,7 @@ " \n", " dec_hidden = enc_hidden\n", " \n", - " dec_input = tf.expand_dims([targ_lang.word2idx['\u003cstart\u003e']] * BATCH_SIZE, 1) \n", + " dec_input = tf.expand_dims([targ_lang.word2idx['']] * BATCH_SIZE, 1) \n", " \n", " # Teacher forcing - feeding the target as the next input\n", " for t in range(1, targ.shape[1]):\n", @@ -625,14 +640,16 @@ " print('Epoch {} Loss {:.4f}'.format(epoch + 1,\n", " total_loss / N_BATCH))\n", " print('Time taken for 1 epoch {} sec\\n'.format(time.time() - start))" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "mU3Ce8M6I3rz" }, + "cell_type": "markdown", "source": [ "## Translate\n", "\n", @@ -644,14 +661,12 @@ ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "EbQpyYs13jF_" + "id": "EbQpyYs13jF_", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "def evaluate(sentence, encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ):\n", " attention_plot = np.zeros((max_length_targ, max_length_inp))\n", @@ -668,7 +683,7 @@ " enc_out, enc_hidden = encoder(inputs, hidden)\n", "\n", " dec_hidden = enc_hidden\n", - " dec_input = tf.expand_dims([targ_lang.word2idx['\u003cstart\u003e']], 0)\n", + " dec_input = tf.expand_dims([targ_lang.word2idx['']], 0)\n", "\n", " for t in range(max_length_targ):\n", " predictions, dec_hidden, attention_weights = decoder(dec_input, dec_hidden, enc_out)\n", @@ -681,24 +696,24 @@ "\n", " result += targ_lang.idx2word[predicted_id] + ' '\n", "\n", - " if targ_lang.idx2word[predicted_id] == '\u003cend\u003e':\n", + " if targ_lang.idx2word[predicted_id] == '':\n", " return result, sentence, attention_plot\n", " \n", " # the predicted ID is fed back into the model\n", " dec_input = tf.expand_dims([predicted_id], 0)\n", "\n", " return result, sentence, attention_plot" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "s5hQWlbN3jGF" + "id": "s5hQWlbN3jGF", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "# function for plotting the attention weights\n", "def plot_attention(attention, sentence, predicted_sentence):\n", @@ -712,17 +727,17 @@ " ax.set_yticklabels([''] + predicted_sentence, fontdict=fontdict)\n", "\n", " plt.show()" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "sl9zUHzg3jGI" + "id": "sl9zUHzg3jGI", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "def translate(sentence, encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ):\n", " result, sentence, attention_plot = evaluate(sentence, encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)\n", @@ -732,91 +747,93 @@ " \n", " attention_plot = attention_plot[:len(result.split(' ')), :len(sentence.split(' '))]\n", " plot_attention(attention_plot, sentence.split(' '), result.split(' '))" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "n250XbnjOaqP" }, + "cell_type": "markdown", "source": [ "## Restore the latest checkpoint and test" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "UJpT9D5_OgP6" + "id": "UJpT9D5_OgP6", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "# restoring the latest checkpoint in checkpoint_dir\n", "checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "WrAM0FDomq3E" + "id": "WrAM0FDomq3E", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "translate(u'hace mucho frio aqui.', encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "zSx2iM36EZQZ" + "id": "zSx2iM36EZQZ", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "translate(u'esta es mi vida.', encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "A3LLCx3ZE0Ls" + "id": "A3LLCx3ZE0Ls", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "translate(u'todavia estan en casa?', encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "DUQVLVqUE1YW" + "id": "DUQVLVqUE1YW", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "# wrong translation\n", "translate(u'trata de averiguarlo.', encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "RTe5P5ioMJwN" }, + "cell_type": "markdown", "source": [ "## Next steps\n", "\n", @@ -824,31 +841,5 @@ "* Experiment with training on a larger dataset, or using more epochs\n" ] } - ], - "metadata": { - "accelerator": "GPU", - "colab": { - "collapsed_sections": [], - "name": "nmt_with_attention.ipynb", - "private_outputs": true, - "provenance": [ - { - "file_id": "1C4fpM7_7IL8ZzF7Gc5abywqQjeQNS2-U", - "timestamp": 1527858391290 - }, - { - "file_id": "1pExo6aUuw0S6MISFWoinfJv0Ftm9V4qv", - "timestamp": 1527776041613 - } - ], - "toc_visible": true, - "version": "0.3.2" - }, - "kernelspec": { - "display_name": "Python 3", - "name": "python3" - } - }, - "nbformat": 4, - "nbformat_minor": 0 -} + ] +} \ No newline at end of file diff --git a/tensorflow/contrib/eager/python/examples/resnet50/BUILD b/tensorflow/contrib/eager/python/examples/resnet50/BUILD index f3135a9668fc0dc7faa93a5f119b53f3efd34c6e..f2851d97223e483da11120f1fe3f0a2f641dfb81 100644 --- a/tensorflow/contrib/eager/python/examples/resnet50/BUILD +++ b/tensorflow/contrib/eager/python/examples/resnet50/BUILD @@ -27,7 +27,7 @@ py_library( cuda_py_test( name = "resnet50_test", - size = "large", + size = "medium", srcs = ["resnet50_test.py"], additional_deps = [ ":resnet50", @@ -35,17 +35,19 @@ cuda_py_test( "//tensorflow/contrib/eager/python:tfe", "//tensorflow:tensorflow_py", ], + shard_count = 4, tags = [ "noasan", # Fix b/118130911 "nomsan", # Fix b/118130911 "notsan", # Fix b/118130911 "optonly", + "oss_serial", ], ) cuda_py_test( name = "resnet50_graph_test", - size = "large", + size = "medium", srcs = ["resnet50_graph_test.py"], additional_deps = [ ":resnet50", @@ -53,10 +55,12 @@ cuda_py_test( "//third_party/py/numpy", "//tensorflow:tensorflow_py", ], + shard_count = 4, tags = [ "noasan", "nomsan", "notsan", "optonly", + "oss_serial", ], ) diff --git a/tensorflow/contrib/eager/python/examples/revnet/BUILD b/tensorflow/contrib/eager/python/examples/revnet/BUILD index 4f0d46b1bae3760a63b2abe871034bdedf258f07..cb207b8ddf3641a68a114386f6a95a26ce2b74d6 100644 --- a/tensorflow/contrib/eager/python/examples/revnet/BUILD +++ b/tensorflow/contrib/eager/python/examples/revnet/BUILD @@ -67,30 +67,36 @@ py_library( # Tests cuda_py_test( name = "ops_test", - size = "large", + size = "medium", srcs = ["ops_test.py"], additional_deps = [ ":ops", "//tensorflow:tensorflow_py", ], + shard_count = 4, + tags = [ + "oss_serial", + ], ) cuda_py_test( name = "blocks_test", - size = "large", + size = "medium", srcs = ["blocks_test.py"], additional_deps = [ ":blocks", "//tensorflow:tensorflow_py", ], + shard_count = 4, tags = [ + "no_oss", # b/123045964 "optonly", ], ) cuda_py_test( name = "revnet_test", - size = "large", + size = "medium", srcs = ["revnet_test.py"], additional_deps = [ ":blocks_test", @@ -98,9 +104,11 @@ cuda_py_test( ":revnet", "//tensorflow:tensorflow_py", ], + shard_count = 4, tags = [ "no_pip", # depends on blocks_test, which is not available in pip package "optonly", + "oss_serial", ], ) @@ -127,6 +135,13 @@ py_binary( name = "main", srcs = ["main.py"], srcs_version = "PY2AND3", + deps = [":main_lib"], +) + +py_library( + name = "main_lib", + srcs = ["main.py"], + srcs_version = "PY2AND3", deps = [ ":cifar_input", ":config", @@ -141,7 +156,7 @@ py_binary( srcs_version = "PY2AND3", deps = [ ":cifar_input", - ":main", + ":main_lib", ":revnet", "//tensorflow:tensorflow_py", ], @@ -153,7 +168,7 @@ py_library( srcs_version = "PY2AND3", deps = [ ":cifar_input", - ":main", + ":main_lib", ":revnet", "//tensorflow:tensorflow_py", ], @@ -165,7 +180,7 @@ py_library( srcs_version = "PY2AND3", deps = [ ":cifar_input", - ":main", + ":main_lib", ":revnet", "//tensorflow:tensorflow_py", ], diff --git a/tensorflow/contrib/eager/python/examples/rnn_colorbot/BUILD b/tensorflow/contrib/eager/python/examples/rnn_colorbot/BUILD index d500b632ebb97fd12ded3a215b0f1a686194874f..f4dbe7ac16f734f7bee045bc71e9559b630adf81 100644 --- a/tensorflow/contrib/eager/python/examples/rnn_colorbot/BUILD +++ b/tensorflow/contrib/eager/python/examples/rnn_colorbot/BUILD @@ -9,6 +9,13 @@ py_binary( name = "rnn_colorbot", srcs = ["rnn_colorbot.py"], srcs_version = "PY2AND3", + deps = [":rnn_colorbot_lib"], +) + +py_library( + name = "rnn_colorbot_lib", + srcs = ["rnn_colorbot.py"], + srcs_version = "PY2AND3", deps = [ "//tensorflow:tensorflow_py", "//tensorflow/contrib/eager/python:tfe", @@ -21,8 +28,11 @@ cuda_py_test( name = "rnn_colorbot_test", srcs = ["rnn_colorbot_test.py"], additional_deps = [ - ":rnn_colorbot", + ":rnn_colorbot_lib", "//tensorflow/contrib/eager/python:tfe", "//tensorflow:tensorflow_py", ], + tags = [ + "oss_serial", + ], ) diff --git a/tensorflow/contrib/eager/python/examples/rnn_colorbot/rnn_colorbot.py b/tensorflow/contrib/eager/python/examples/rnn_colorbot/rnn_colorbot.py index 74ebb1ec77131a560b1ebfd062c690920c35e261..1c718a5ce3d8e1541656d92fd5e8dad6c6683c4c 100644 --- a/tensorflow/contrib/eager/python/examples/rnn_colorbot/rnn_colorbot.py +++ b/tensorflow/contrib/eager/python/examples/rnn_colorbot/rnn_colorbot.py @@ -207,7 +207,7 @@ class RNNColorbot(tf.keras.Model): def loss(labels, predictions): """Computes mean squared loss.""" - return tf.reduce_mean(tf.square(predictions - labels)) + return tf.reduce_mean(tf.squared_difference(predictions, labels)) def test(model, eval_data): diff --git a/tensorflow/contrib/eager/python/examples/rnn_ptb/BUILD b/tensorflow/contrib/eager/python/examples/rnn_ptb/BUILD index 2cc2fcbfeb21ee6218d7912d9a93ea2f7b2ea226..43a6ca526d3a0aecda2c8df865a0487ac28758ab 100644 --- a/tensorflow/contrib/eager/python/examples/rnn_ptb/BUILD +++ b/tensorflow/contrib/eager/python/examples/rnn_ptb/BUILD @@ -9,6 +9,13 @@ py_binary( name = "rnn_ptb", srcs = ["rnn_ptb.py"], srcs_version = "PY2AND3", + deps = [":rnn_ptb_lib"], +) + +py_library( + name = "rnn_ptb_lib", + srcs = ["rnn_ptb.py"], + srcs_version = "PY2AND3", deps = [ "//tensorflow:tensorflow_py", "//tensorflow/contrib/cudnn_rnn:cudnn_rnn_py", @@ -21,18 +28,22 @@ cuda_py_test( name = "rnn_ptb_test", srcs = ["rnn_ptb_test.py"], additional_deps = [ - ":rnn_ptb", + ":rnn_ptb_lib", "//tensorflow/contrib/eager/python:tfe", "//tensorflow:tensorflow_py", ], + tags = ["no_oss"], # b/123045964 ) cuda_py_test( name = "rnn_ptb_graph_test", srcs = ["rnn_ptb_graph_test.py"], additional_deps = [ - ":rnn_ptb", + ":rnn_ptb_lib", "//third_party/py/numpy", "//tensorflow:tensorflow_py", ], + tags = [ + "oss_serial", + ], ) diff --git a/tensorflow/contrib/eager/python/examples/rnn_ptb/rnn_ptb.py b/tensorflow/contrib/eager/python/examples/rnn_ptb/rnn_ptb.py index 15776c694e92825895437a4c1547699f6d9269fb..9b5a2c947b153308c83f1a922d06c034ec5f9ddf 100644 --- a/tensorflow/contrib/eager/python/examples/rnn_ptb/rnn_ptb.py +++ b/tensorflow/contrib/eager/python/examples/rnn_ptb/rnn_ptb.py @@ -128,7 +128,7 @@ class PTBModel(tf.keras.Model): self.linear = layers.Dense( vocab_size, kernel_initializer=tf.random_uniform_initializer(-0.1, 0.1)) - self._output_shape = [-1, embedding_dim] + self._output_shape = [-1, hidden_dim] def call(self, input_seq, training): """Run the forward pass of PTBModel. diff --git a/tensorflow/contrib/eager/python/examples/spinn/BUILD b/tensorflow/contrib/eager/python/examples/spinn/BUILD index 5966f1d4873e8e77b3ad5914da7bfc7e69d4e341..9b0fbaa6793e28d327745767e6ccd3085211ff7d 100644 --- a/tensorflow/contrib/eager/python/examples/spinn/BUILD +++ b/tensorflow/contrib/eager/python/examples/spinn/BUILD @@ -42,5 +42,6 @@ cuda_py_test( "no-internal-py3", # flaky "no_cuda_on_cpu_tap", "no_pip", # because spinn.py is under third_party/. + "oss_serial", ], ) diff --git a/tensorflow/contrib/eager/python/metrics_impl.py b/tensorflow/contrib/eager/python/metrics_impl.py index 566246de4957c1dc5919c10e22146706f9e50be8..c8d9266672a8b87d32338ea7c4f74fb40d41c767 100644 --- a/tensorflow/contrib/eager/python/metrics_impl.py +++ b/tensorflow/contrib/eager/python/metrics_impl.py @@ -37,7 +37,7 @@ from tensorflow.python.training.checkpointable import base as checkpointable _to_replace = re.compile("[^A-Za-z0-9.]") -class Metric(checkpointable.CheckpointableBase): +class Metric(checkpointable.Checkpointable): """A metric holds state for aggregating statistics over an evaluation run. Example use with eager execution: diff --git a/tensorflow/contrib/eager/python/tfe.py b/tensorflow/contrib/eager/python/tfe.py index 31481d7685c79b76c40b1f8041441a0e71d3b00e..12bbdc08cffe3bcb922e7d75c04566f7741fb7f5 100644 --- a/tensorflow/contrib/eager/python/tfe.py +++ b/tensorflow/contrib/eager/python/tfe.py @@ -62,7 +62,6 @@ To use, at program startup, call `tf.enable_eager_execution()`. @@Checkpoint @@Checkpointable -@@CheckpointableSaver @@executing_eagerly @@in_eager_mode @@ -138,8 +137,7 @@ from tensorflow.python.ops.resource_variable_ops import ResourceVariable as Vari from tensorflow.python.ops.variable_scope import EagerVariableStore from tensorflow.python.ops import script_ops from tensorflow.python.ops import template -from tensorflow.python.training.checkpointable.tracking import Checkpointable -from tensorflow.python.training.checkpointable.util import CheckpointableSaver +from tensorflow.python.training.checkpointable.tracking import AutoCheckpointable as Checkpointable from tensorflow.python.training.checkpointable.util import Checkpoint from tensorflow.python.util.all_util import remove_undocumented diff --git a/tensorflow/contrib/eager/python/tfe_test.py b/tensorflow/contrib/eager/python/tfe_test.py index 8c35dddb5a515aa09cc70c173a9f0605e8567e82..6881fabdc09e3275c29f3013283999c96e283770 100644 --- a/tensorflow/contrib/eager/python/tfe_test.py +++ b/tensorflow/contrib/eager/python/tfe_test.py @@ -20,6 +20,7 @@ from __future__ import print_function import tempfile from tensorflow.contrib.eager.python import tfe +from tensorflow.python.eager import context from tensorflow.python.framework import constant_op from tensorflow.python.framework import errors from tensorflow.python.framework import ops @@ -40,6 +41,9 @@ class TFETest(test_util.TensorFlowTestCase): self.assertAllEqual([[4.]], y.numpy()) def testInstantError(self): + if context.num_gpus(): + # TODO(nareshmodi): make this test better + self.skipTest("Gather doesn't do index checking on GPUs") with self.assertRaisesRegexp(errors.InvalidArgumentError, r'indices = 7 is not in \[0, 3\)'): array_ops.gather([0, 1, 2], 7) diff --git a/tensorflow/contrib/factorization/BUILD b/tensorflow/contrib/factorization/BUILD index e344d7a23b55134612aab430b50cf065bd1095e4..48a6ef4dca0ca7682f7b99b66177679f29ad9ec9 100644 --- a/tensorflow/contrib/factorization/BUILD +++ b/tensorflow/contrib/factorization/BUILD @@ -28,7 +28,6 @@ tf_custom_op_py_library( "python/ops/wals.py", ], dso = [ - ":python/ops/_clustering_ops.so", ":python/ops/_factorization_ops.so", ], kernels = [ @@ -38,12 +37,12 @@ tf_custom_op_py_library( srcs_version = "PY2AND3", deps = [ ":factorization_ops_test_utils_py", - ":gen_clustering_ops", ":gen_factorization_ops", "//tensorflow/contrib/framework:framework_py", "//tensorflow/contrib/util:util_py", "//tensorflow/python:array_ops", "//tensorflow/python:check_ops", + "//tensorflow/python:clustering_ops_gen", "//tensorflow/python:control_flow_ops", "//tensorflow/python:data_flow_ops", "//tensorflow/python:embedding_ops", @@ -77,17 +76,6 @@ py_library( ], ) -# Ops -tf_custom_op_library( - name = "python/ops/_clustering_ops.so", - srcs = [ - "ops/clustering_ops.cc", - ], - deps = [ - "//tensorflow/contrib/factorization/kernels:clustering_ops", - ], -) - tf_custom_op_library( name = "python/ops/_factorization_ops.so", srcs = [ @@ -100,26 +88,16 @@ tf_custom_op_library( ) tf_gen_op_libs([ - "clustering_ops", "factorization_ops", ]) cc_library( name = "all_ops", deps = [ - ":clustering_ops_op_lib", ":factorization_ops_op_lib", ], ) -tf_gen_op_wrapper_py( - name = "gen_clustering_ops", - out = "python/ops/gen_clustering_ops.py", - deps = [ - ":clustering_ops_op_lib", - ], -) - tf_gen_op_wrapper_py( name = "gen_factorization_ops", out = "python/ops/gen_factorization_ops.py", @@ -131,7 +109,7 @@ tf_gen_op_wrapper_py( # Ops tests tf_py_test( name = "gmm_test", - size = "large", + size = "medium", srcs = [ "python/ops/gmm_test.py", ], @@ -152,6 +130,7 @@ tf_py_test( "//tensorflow/python:random_seed", "//tensorflow/python:training", ], + shard_count = 4, tags = [ "no_pip", # b/38283730 "notsan", # Flaky: b/30756419 @@ -249,7 +228,7 @@ py_test( tf_py_test( name = "wals_test", - size = "large", + size = "medium", srcs = ["python/ops/wals_test.py"], additional_deps = [ ":factorization_py", @@ -272,8 +251,8 @@ tf_py_test( "//tensorflow/python:training", "//tensorflow/python:variables", ], + shard_count = 4, tags = [ - "manual", "noasan", # times out b/63678675 "nomsan", ], diff --git a/tensorflow/contrib/factorization/kernels/BUILD b/tensorflow/contrib/factorization/kernels/BUILD index ea8b9a17a27093cb57564861815edd6ecb18a014..23d7e088d067effa446e4bcdc9609db612066568 100644 --- a/tensorflow/contrib/factorization/kernels/BUILD +++ b/tensorflow/contrib/factorization/kernels/BUILD @@ -11,7 +11,6 @@ load("//tensorflow:tensorflow.bzl", "tf_cc_test") cc_library( name = "all_kernels", deps = [ - ":clustering_ops", ":masked_matmul_ops", ":wals_solver_ops", "@protobuf_archive//:protobuf_headers", @@ -29,17 +28,6 @@ cc_library( alwayslink = 1, ) -cc_library( - name = "clustering_ops", - srcs = ["clustering_ops.cc"], - deps = [ - "//tensorflow/core:framework_headers_lib", - "//third_party/eigen3", - "@protobuf_archive//:protobuf_headers", - ], - alwayslink = 1, -) - cc_library( name = "masked_matmul_ops", srcs = ["masked_matmul_ops.cc"], @@ -51,19 +39,3 @@ cc_library( ], alwayslink = 1, ) - -tf_cc_test( - name = "clustering_ops_test", - srcs = ["clustering_ops_test.cc"], - deps = [ - ":clustering_ops", - "//tensorflow/contrib/factorization:clustering_ops_op_lib", - "//tensorflow/core:core_cpu", - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - "//tensorflow/core:testlib", - ], -) diff --git a/tensorflow/contrib/factorization/kernels/masked_matmul_ops.cc b/tensorflow/contrib/factorization/kernels/masked_matmul_ops.cc index a8c5d0763c28ba2b54f217405f0da65533f26b91..68078ba8bbb07b4344c19d554012d214229f9c4f 100644 --- a/tensorflow/contrib/factorization/kernels/masked_matmul_ops.cc +++ b/tensorflow/contrib/factorization/kernels/masked_matmul_ops.cc @@ -19,12 +19,12 @@ #include #include +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/bounds_check.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/threadpool.h" diff --git a/tensorflow/contrib/factorization/ops/clustering_ops.cc b/tensorflow/contrib/factorization/ops/clustering_ops.cc deleted file mode 100644 index 2686702c1d5768f661dac610c96089eb02e360d7..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/factorization/ops/clustering_ops.cc +++ /dev/null @@ -1,91 +0,0 @@ -// Copyright 2016 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" - -namespace tensorflow { - -REGISTER_OP("KmeansPlusPlusInitialization") - .Input("points: float32") - .Input("num_to_sample: int64") - .Input("seed: int64") - .Input("num_retries_per_sample: int64") - .Output("samples: float32") - .SetShapeFn(shape_inference::UnknownShape) - .Doc(R"( -Selects num_to_sample rows of input using the KMeans++ criterion. - -Rows of points are assumed to be input points. One row is selected at random. -Subsequent rows are sampled with probability proportional to the squared L2 -distance from the nearest row selected thus far till num_to_sample rows have -been sampled. - -points: Matrix of shape (n, d). Rows are assumed to be input points. -num_to_sample: Scalar. The number of rows to sample. This value must not be - larger than n. -seed: Scalar. Seed for initializing the random number generator. -num_retries_per_sample: Scalar. For each row that is sampled, this parameter - specifies the number of additional points to draw from the current - distribution before selecting the best. If a negative value is specified, a - heuristic is used to sample O(log(num_to_sample)) additional points. -samples: Matrix of shape (num_to_sample, d). The sampled rows. -)"); - -REGISTER_OP("KMC2ChainInitialization") - .Input("distances: float32") - .Input("seed: int64") - .Output("index: int64") - .SetShapeFn(shape_inference::ScalarShape) - .Doc(R"( -Returns the index of a data point that should be added to the seed set. - -Entries in distances are assumed to be squared distances of candidate points to -the already sampled centers in the seed set. The op constructs one Markov chain -of the k-MC^2 algorithm and returns the index of one candidate point to be added -as an additional cluster center. - -distances: Vector with squared distances to the closest previously sampled - cluster center for each candidate point. -seed: Scalar. Seed for initializing the random number generator. -index: Scalar with the index of the sampled point. -)"); - -REGISTER_OP("NearestNeighbors") - .Input("points: float32") - .Input("centers: float32") - .Input("k: int64") - .Output("nearest_center_indices: int64") - .Output("nearest_center_distances: float32") - .SetShapeFn(shape_inference::UnknownShape) - .Doc(R"( -Selects the k nearest centers for each point. - -Rows of points are assumed to be input points. Rows of centers are assumed to be -the list of candidate centers. For each point, the k centers that have least L2 -distance to it are computed. - -points: Matrix of shape (n, d). Rows are assumed to be input points. -centers: Matrix of shape (m, d). Rows are assumed to be centers. -k: Scalar. Number of nearest centers to return for each point. If k is larger - than m, then only m centers are returned. -nearest_center_indices: Matrix of shape (n, min(m, k)). Each row contains the - indices of the centers closest to the corresponding point, ordered by - increasing distance. -nearest_center_distances: Matrix of shape (n, min(m, k)). Each row contains the - squared L2 distance to the corresponding center in nearest_center_indices. -)"); - -} // namespace tensorflow diff --git a/tensorflow/contrib/factorization/python/ops/clustering_ops.py b/tensorflow/contrib/factorization/python/ops/clustering_ops.py index 84e80791f4991ad2b67d0a00ee1e00cf0d0daadc..d48b89cbacce34781819010addbcbd0ba66f9873 100644 --- a/tensorflow/contrib/factorization/python/ops/clustering_ops.py +++ b/tensorflow/contrib/factorization/python/ops/clustering_ops.py @@ -18,28 +18,23 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.factorization.python.ops import gen_clustering_ops -# go/tf-wildcard-import -# pylint: disable=wildcard-import -from tensorflow.contrib.factorization.python.ops.gen_clustering_ops import * -# pylint: enable=wildcard-import -from tensorflow.contrib.util import loader from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes 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 gen_clustering_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_impl from tensorflow.python.ops import random_ops from tensorflow.python.ops import state_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops.embedding_ops import embedding_lookup -from tensorflow.python.platform import resource_loader - -_clustering_ops = loader.load_op_library( - resource_loader.get_path_to_datafile('_clustering_ops.so')) +# go/tf-wildcard-import +# pylint: disable=wildcard-import +from tensorflow.python.ops.gen_clustering_ops import * +# pylint: enable=wildcard-import # Euclidean distance between vectors U and V is defined as \\(||U - V||_F\\) # which is the square root of the sum of the absolute squares of the elements diff --git a/tensorflow/contrib/factorization/python/ops/gmm_ops.py b/tensorflow/contrib/factorization/python/ops/gmm_ops.py index d365ad111760247fc18b730657390f07ba6b865e..9f0664dfe5ba7a098b6976388d1cf737dafb4842 100644 --- a/tensorflow/contrib/factorization/python/ops/gmm_ops.py +++ b/tensorflow/contrib/factorization/python/ops/gmm_ops.py @@ -314,8 +314,7 @@ class GmmAlgorithm(object): # reparametrization of variance parameters. det_expanded = math_ops.reduce_sum( math_ops.log(self._covs + 1e-3), 1, keepdims=True) - diff = shard - self._means - x2 = math_ops.square(diff) + x2 = math_ops.squared_difference(shard, self._means) cov_expanded = array_ops.expand_dims(1.0 / (self._covs + 1e-3), 2) # num_classes X num_examples x2_cov = math_ops.matmul(x2, cov_expanded) diff --git a/tensorflow/contrib/feature_column/BUILD b/tensorflow/contrib/feature_column/BUILD index 4e29e2559986012d8eeeaec807f14181226363aa..8fc5f1cfe7800653ef1e43c6d40d1a66e34f2106 100644 --- a/tensorflow/contrib/feature_column/BUILD +++ b/tensorflow/contrib/feature_column/BUILD @@ -96,7 +96,6 @@ tf_py_test( name = "sequence_feature_column_v2_test", srcs = ["python/feature_column/sequence_feature_column_v2_test.py"], additional_deps = [ - ":sequence_feature_column", ":sequence_feature_column_v2", "@absl_py//absl/testing:parameterized", "//third_party/py/numpy", @@ -113,3 +112,20 @@ tf_py_test( ], tags = ["no_pip"], ) + +py_test( + name = "sequence_feature_column_v2_integration_test", + srcs = ["python/feature_column/sequence_feature_column_v2_integration_test.py"], + srcs_version = "PY2AND3", + tags = ["no_pip"], + deps = [ + ":sequence_feature_column_v2", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_ops", + "//tensorflow/python:parsing_ops", + "//tensorflow/python:training", + "//tensorflow/python:util", + "//tensorflow/python/feature_column:feature_column_py", + "//tensorflow/python/keras:layers", + ], +) diff --git a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2.py b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2.py index 83b93ec332044f754f9dcde8d7c5c19b26e53a4a..2f4bda194a41242167e0abfcaeac5044f6026f85 100644 --- a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2.py +++ b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2.py @@ -27,6 +27,7 @@ import collections from tensorflow.python.feature_column import feature_column as fc_old from tensorflow.python.feature_column import feature_column_lib as fc +from tensorflow.python.feature_column import feature_column_v2 as fc_v2 from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape @@ -34,107 +35,115 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import check_ops from tensorflow.python.ops import parsing_ops from tensorflow.python.ops import sparse_ops -from tensorflow.python.ops import variable_scope # pylint: disable=protected-access -def sequence_input_layer( - features, - feature_columns, - weight_collections=None, - trainable=True): - """"Builds input layer for sequence input. +class SequenceFeatures(fc_v2._BaseFeaturesLayer): + """A layer for sequence input. - All `feature_columns` must be sequence dense columns with the same - `sequence_length`. The output of this method can be fed into sequence - networks, such as RNN. + All `feature_columns` must be sequence dense columns with the same + `sequence_length`. The output of this method can be fed into sequence + networks, such as RNN. - The output of this method is a 3D `Tensor` of shape `[batch_size, T, D]`. - `T` is the maximum sequence length for this batch, which could differ from - batch to batch. + The output of this method is a 3D `Tensor` of shape `[batch_size, T, D]`. + `T` is the maximum sequence length for this batch, which could differ from + batch to batch. - If multiple `feature_columns` are given with `Di` `num_elements` each, their - outputs are concatenated. So, the final `Tensor` has shape - `[batch_size, T, D0 + D1 + ... + Dn]`. + If multiple `feature_columns` are given with `Di` `num_elements` each, their + outputs are concatenated. So, the final `Tensor` has shape + `[batch_size, T, D0 + D1 + ... + Dn]`. - Example: + Example: - ```python - rating = sequence_numeric_column('rating') - watches = sequence_categorical_column_with_identity( - 'watches', num_buckets=1000) - watches_embedding = embedding_column(watches, dimension=10) - columns = [rating, watches] + ```python + rating = sequence_numeric_column('rating') + watches = sequence_categorical_column_with_identity( + 'watches', num_buckets=1000) + watches_embedding = embedding_column(watches, dimension=10) + columns = [rating, watches] - features = tf.parse_example(..., features=make_parse_example_spec(columns)) - input_layer, sequence_length = sequence_input_layer(features, columns) + features = tf.parse_example(..., features=make_parse_example_spec(columns)) + sequence_input_layer = SequenceFeatures(columns) + sequence_input, sequence_length = sequence_input_layer(features) + sequence_length_mask = tf.sequence_mask(sequence_length) - rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size) - outputs, state = tf.nn.dynamic_rnn( - rnn_cell, inputs=input_layer, sequence_length=sequence_length) - ``` + rnn_cell = tf.keras.layers.SimpleRNNCell(hidden_size) + rnn_layer = tf.keras.layers.RNN(rnn_cell) + outputs, state = rnn_layer(sequence_input, mask=sequence_length_mask) + ``` + """ - Args: - features: A dict mapping keys to tensors. - feature_columns: An iterable of dense sequence columns. Valid columns are - - `embedding_column` that wraps a `sequence_categorical_column_with_*` - - `sequence_numeric_column`. - weight_collections: A list of collection names to which the Variable will be - added. Note that variables will also be added to collections - `tf.GraphKeys.GLOBAL_VARIABLES` and `ops.GraphKeys.MODEL_VARIABLES`. - trainable: If `True` also add the variable to the graph collection - `GraphKeys.TRAINABLE_VARIABLES`. + def __init__( + self, + feature_columns, + trainable=True, + name=None, + **kwargs): + """"Constructs a SequenceFeatures layer. - Returns: - An `(input_layer, sequence_length)` tuple where: - - input_layer: A float `Tensor` of shape `[batch_size, T, D]`. - `T` is the maximum sequence length for this batch, which could differ - from batch to batch. `D` is the sum of `num_elements` for all - `feature_columns`. - - sequence_length: An int `Tensor` of shape `[batch_size]`. The sequence - length for each example. + Args: + feature_columns: An iterable of dense sequence columns. Valid columns are + - `embedding_column` that wraps a `sequence_categorical_column_with_*` + - `sequence_numeric_column`. + trainable: Boolean, whether the layer's variables will be updated via + gradient descent during training. + name: Name to give to the SequenceFeatures. + **kwargs: Keyword arguments to construct a layer. + + Raises: + ValueError: If any of the `feature_columns` is not a + `SequenceDenseColumn`. + """ + super(SequenceFeatures, self).__init__( + feature_columns=feature_columns, + trainable=trainable, + name=name, + expected_column_type=fc_v2.SequenceDenseColumn, + **kwargs) - Raises: - ValueError: If any of the `feature_columns` is the wrong type. - """ - feature_columns = fc_old._normalize_feature_columns(feature_columns) - for c in feature_columns: - if not isinstance(c, fc_old._SequenceDenseColumn): - raise ValueError( - 'All feature_columns must be of type _SequenceDenseColumn. ' - 'You can wrap a sequence_categorical_column with an embedding_column ' - 'or indicator_column. ' - 'Given (type {}): {}'.format(type(c), c)) - - with variable_scope.variable_scope( - None, default_name='sequence_input_layer', values=features.values()): - builder = fc_old._LazyBuilder(features) + def _target_shape(self, input_shape, total_elements): + return (input_shape[0], input_shape[1], total_elements) + + def call(self, features): + """Returns sequence input corresponding to the `feature_columns`. + + Args: + features: A dict mapping keys to tensors. + + Returns: + An `(input_layer, sequence_length)` tuple where: + - input_layer: A float `Tensor` of shape `[batch_size, T, D]`. + `T` is the maximum sequence length for this batch, which could differ + from batch to batch. `D` is the sum of `num_elements` for all + `feature_columns`. + - sequence_length: An int `Tensor` of shape `[batch_size]`. The sequence + length for each example. + + Raises: + ValueError: If features are not a dictionary. + """ + if not isinstance(features, dict): + raise ValueError('We expected a dictionary here. Instead we got: ', + features) + transformation_cache = fc.FeatureTransformationCache(features) output_tensors = [] sequence_lengths = [] - ordered_columns = [] - - for column in sorted(feature_columns, key=lambda x: x.name): - ordered_columns.append(column) - with variable_scope.variable_scope( - None, default_name=column._var_scope_name): - dense_tensor, sequence_length = column._get_sequence_dense_tensor( - builder, - weight_collections=weight_collections, - trainable=trainable) + + for column in self._feature_columns: + with ops.name_scope(column.name): + dense_tensor, sequence_length = column.get_sequence_dense_tensor( + transformation_cache, self._state_manager) # Flattens the final dimension to produce a 3D Tensor. - num_elements = column._variable_shape.num_elements() - shape = array_ops.shape(dense_tensor) - target_shape = [shape[0], shape[1], num_elements] - output_tensors.append( - array_ops.reshape(dense_tensor, shape=target_shape)) + output_tensors.append(self._process_dense_tensor(column, dense_tensor)) sequence_lengths.append(sequence_length) - fc_old._verify_static_batch_size_equality(output_tensors, ordered_columns) - fc_old._verify_static_batch_size_equality(sequence_lengths, ordered_columns) + # Check and process sequence lengths. + fc_v2._verify_static_batch_size_equality(sequence_lengths, + self._feature_columns) sequence_length = _assert_all_equal_and_return(sequence_lengths) - return array_ops.concat(output_tensors, -1), sequence_length + return self._verify_and_concat_tensors(output_tensors), sequence_length def concatenate_context_input(context_input, sequence_input): @@ -203,12 +212,13 @@ def sequence_categorical_column_with_identity( columns = [watches_embedding] features = tf.parse_example(..., features=make_parse_example_spec(columns)) - sequence_feature_layer = SequenceFeatureLayer(columns) - input_layer, sequence_length = sequence_feature_layer(features) + sequence_feature_layer = SequenceFeatures(columns) + sequence_input, sequence_length = sequence_feature_layer(features) + sequence_length_mask = tf.sequence_mask(sequence_length) - rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size) - outputs, state = tf.nn.dynamic_rnn( - rnn_cell, inputs=input_layer, sequence_length=sequence_length) + rnn_cell = tf.keras.layers.SimpleRNNCell(hidden_size) + rnn_layer = tf.keras.layers.RNN(rnn_cell) + outputs, state = rnn_layer(sequence_input, mask=sequence_length_mask) ``` Args: @@ -250,12 +260,13 @@ def sequence_categorical_column_with_hash_bucket( columns = [tokens_embedding] features = tf.parse_example(..., features=make_parse_example_spec(columns)) - sequence_feature_layer = SequenceFeatureLayer(columns) - input_layer, sequence_length = sequence_feature_layer(features) + sequence_feature_layer = SequenceFeatures(columns) + sequence_input, sequence_length = sequence_feature_layer(features) + sequence_length_mask = tf.sequence_mask(sequence_length) - rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size) - outputs, state = tf.nn.dynamic_rnn( - rnn_cell, inputs=input_layer, sequence_length=sequence_length) + rnn_cell = tf.keras.layers.SimpleRNNCell(hidden_size) + rnn_layer = tf.keras.layers.RNN(rnn_cell) + outputs, state = rnn_layer(sequence_input, mask=sequence_length_mask) ``` Args: @@ -296,12 +307,13 @@ def sequence_categorical_column_with_vocabulary_file( columns = [states_embedding] features = tf.parse_example(..., features=make_parse_example_spec(columns)) - sequence_feature_layer = SequenceFeatureLayer(columns) - input_layer, sequence_length = sequence_feature_layer(features) + sequence_feature_layer = SequenceFeatures(columns) + sequence_input, sequence_length = sequence_feature_layer(features) + sequence_length_mask = tf.sequence_mask(sequence_length) - rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size) - outputs, state = tf.nn.dynamic_rnn( - rnn_cell, inputs=input_layer, sequence_length=sequence_length) + rnn_cell = tf.keras.layers.SimpleRNNCell(hidden_size) + rnn_layer = tf.keras.layers.RNN(rnn_cell) + outputs, state = rnn_layer(sequence_input, mask=sequence_length_mask) ``` Args: @@ -358,12 +370,13 @@ def sequence_categorical_column_with_vocabulary_list( columns = [colors_embedding] features = tf.parse_example(..., features=make_parse_example_spec(columns)) - sequence_feature_layer = SequenceFeatureLayer(columns) - input_layer, sequence_length = sequence_feature_layer(features) + sequence_feature_layer = SequenceFeatures(columns) + sequence_input, sequence_length = sequence_feature_layer(features) + sequence_length_mask = tf.sequence_mask(sequence_length) - rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size) - outputs, state = tf.nn.dynamic_rnn( - rnn_cell, inputs=input_layer, sequence_length=sequence_length) + rnn_cell = tf.keras.layers.SimpleRNNCell(hidden_size) + rnn_layer = tf.keras.layers.RNN(rnn_cell) + outputs, state = rnn_layer(sequence_input, mask=sequence_length_mask) ``` Args: @@ -415,12 +428,13 @@ def sequence_numeric_column( columns = [temperature] features = tf.parse_example(..., features=make_parse_example_spec(columns)) - sequence_feature_layer = SequenceFeatureLayer(columns) - input_layer, sequence_length = sequence_feature_layer(features) + sequence_feature_layer = SequenceFeatures(columns) + sequence_input, sequence_length = sequence_feature_layer(features) + sequence_length_mask = tf.sequence_mask(sequence_length) - rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size) - outputs, state = tf.nn.dynamic_rnn( - rnn_cell, inputs=input_layer, sequence_length=sequence_length) + rnn_cell = tf.keras.layers.SimpleRNNCell(hidden_size) + rnn_layer = tf.keras.layers.RNN(rnn_cell) + outputs, state = rnn_layer(sequence_input, mask=sequence_length_mask) ``` Args: @@ -445,7 +459,7 @@ def sequence_numeric_column( ValueError: if any dimension in shape is not a positive integer. ValueError: if `dtype` is not convertible to `tf.float32`. """ - shape = fc_old._check_shape(shape=shape, key=key) + shape = fc_v2._check_shape(shape=shape, key=key) if not (dtype.is_integer or dtype.is_floating): raise ValueError('dtype must be convertible to float. ' 'dtype: {}, key: {}'.format(dtype, key)) @@ -540,8 +554,10 @@ class SequenceNumericColumn( # For the 2D case, the raw values are grouped according to num_elements; # for the 3D case, the grouping happens in the third dimension, and # sequence length is not affected. - num_elements = (self.variable_shape.num_elements() - if sp_tensor.shape.ndims == 2 else 1) + if sp_tensor.shape.ndims == 2: + num_elements = self.variable_shape.num_elements() + else: + num_elements = 1 seq_length = fc_old._sequence_length_from_sparse_tensor( sp_tensor, num_elements=num_elements) diff --git a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2_integration_test.py b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2_integration_test.py new file mode 100644 index 0000000000000000000000000000000000000000..1b165a620ae67e855400eb297ec17db80eac7937 --- /dev/null +++ b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2_integration_test.py @@ -0,0 +1,283 @@ +# 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. +# ============================================================================== +"""Integration test for sequence feature columns with SequenceExamples.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import string +import tempfile + +from google.protobuf import text_format + +from tensorflow.contrib.feature_column.python.feature_column import sequence_feature_column_v2 as sfc +from tensorflow.core.example import example_pb2 +from tensorflow.core.example import feature_pb2 +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.feature_column import feature_column_v2 as fc +from tensorflow.python.keras.layers import recurrent +from tensorflow.python.ops import parsing_ops +from tensorflow.python.ops import variables +from tensorflow.python.platform import test +from tensorflow.python.util import compat + + +class SequenceFeatureColumnIntegrationTest(test.TestCase): + + def _make_sequence_example(self): + example = example_pb2.SequenceExample() + example.context.feature['int_ctx'].int64_list.value.extend([5]) + example.context.feature['float_ctx'].float_list.value.extend([123.6]) + for val in range(0, 10, 2): + feat = feature_pb2.Feature() + feat.int64_list.value.extend([val] * val) + example.feature_lists.feature_list['int_list'].feature.extend([feat]) + for val in range(1, 11, 2): + feat = feature_pb2.Feature() + feat.bytes_list.value.extend([compat.as_bytes(str(val))] * val) + example.feature_lists.feature_list['str_list'].feature.extend([feat]) + + return example + + def _build_feature_columns(self): + col = fc.categorical_column_with_identity('int_ctx', num_buckets=100) + ctx_cols = [ + fc.embedding_column(col, dimension=10), + fc.numeric_column('float_ctx') + ] + + identity_col = sfc.sequence_categorical_column_with_identity( + 'int_list', num_buckets=10) + bucket_col = sfc.sequence_categorical_column_with_hash_bucket( + 'bytes_list', hash_bucket_size=100) + seq_cols = [ + fc.embedding_column(identity_col, dimension=10), + fc.embedding_column(bucket_col, dimension=20) + ] + + return ctx_cols, seq_cols + + def test_sequence_example_into_input_layer(self): + examples = [_make_sequence_example().SerializeToString()] * 100 + ctx_cols, seq_cols = self._build_feature_columns() + + def _parse_example(example): + ctx, seq = parsing_ops.parse_single_sequence_example( + example, + context_features=fc.make_parse_example_spec_v2(ctx_cols), + sequence_features=fc.make_parse_example_spec_v2(seq_cols)) + ctx.update(seq) + return ctx + + ds = dataset_ops.Dataset.from_tensor_slices(examples) + ds = ds.map(_parse_example) + ds = ds.batch(20) + + # Test on a single batch + features = ds.make_one_shot_iterator().get_next() + + # Tile the context features across the sequence features + sequence_input_layer = sfc.SequenceFeatures(seq_cols) + seq_layer, _ = sequence_input_layer(features) + input_layer = fc.DenseFeatures(ctx_cols) + ctx_layer = input_layer(features) + input_layer = sfc.concatenate_context_input(ctx_layer, seq_layer) + + rnn_layer = recurrent.RNN(recurrent.SimpleRNNCell(10)) + output = rnn_layer(input_layer) + + with self.cached_session() as sess: + sess.run(variables.global_variables_initializer()) + features_r = sess.run(features) + self.assertAllEqual(features_r['int_list'].dense_shape, [20, 3, 6]) + + output_r = sess.run(output) + self.assertAllEqual(output_r.shape, [20, 10]) + + +class SequenceExampleParsingTest(test.TestCase): + + def test_seq_ex_in_sequence_categorical_column_with_identity(self): + self._test_parsed_sequence_example( + 'int_list', sfc.sequence_categorical_column_with_identity, + 10, [3, 6], [2, 4, 6]) + + def test_seq_ex_in_sequence_categorical_column_with_hash_bucket(self): + self._test_parsed_sequence_example( + 'bytes_list', sfc.sequence_categorical_column_with_hash_bucket, + 10, [3, 4], [compat.as_bytes(x) for x in 'acg']) + + def test_seq_ex_in_sequence_categorical_column_with_vocabulary_list(self): + self._test_parsed_sequence_example( + 'bytes_list', sfc.sequence_categorical_column_with_vocabulary_list, + list(string.ascii_lowercase), [3, 4], + [compat.as_bytes(x) for x in 'acg']) + + def test_seq_ex_in_sequence_categorical_column_with_vocabulary_file(self): + _, fname = tempfile.mkstemp() + with open(fname, 'w') as f: + f.write(string.ascii_lowercase) + self._test_parsed_sequence_example( + 'bytes_list', sfc.sequence_categorical_column_with_vocabulary_file, + fname, [3, 4], [compat.as_bytes(x) for x in 'acg']) + + def _test_parsed_sequence_example( + self, col_name, col_fn, col_arg, shape, values): + """Helper function to check that each FeatureColumn parses correctly. + + Args: + col_name: string, name to give to the feature column. Should match + the name that the column will parse out of the features dict. + col_fn: function used to create the feature column. For example, + sequence_numeric_column. + col_arg: second arg that the target feature column is expecting. + shape: the expected dense_shape of the feature after parsing into + a SparseTensor. + values: the expected values at index [0, 2, 6] of the feature + after parsing into a SparseTensor. + """ + example = _make_sequence_example() + columns = [ + fc.categorical_column_with_identity('int_ctx', num_buckets=100), + fc.numeric_column('float_ctx'), + col_fn(col_name, col_arg) + ] + context, seq_features = parsing_ops.parse_single_sequence_example( + example.SerializeToString(), + context_features=fc.make_parse_example_spec_v2(columns[:2]), + sequence_features=fc.make_parse_example_spec_v2(columns[2:])) + + with self.cached_session() as sess: + ctx_result, seq_result = sess.run([context, seq_features]) + self.assertEqual(list(seq_result[col_name].dense_shape), shape) + self.assertEqual( + list(seq_result[col_name].values[[0, 2, 6]]), values) + self.assertEqual(list(ctx_result['int_ctx'].dense_shape), [1]) + self.assertEqual(ctx_result['int_ctx'].values[0], 5) + self.assertEqual(list(ctx_result['float_ctx'].shape), [1]) + self.assertAlmostEqual(ctx_result['float_ctx'][0], 123.6, places=1) + + +_SEQ_EX_PROTO = """ +context { + feature { + key: "float_ctx" + value { + float_list { + value: 123.6 + } + } + } + feature { + key: "int_ctx" + value { + int64_list { + value: 5 + } + } + } +} +feature_lists { + feature_list { + key: "bytes_list" + value { + feature { + bytes_list { + value: "a" + } + } + feature { + bytes_list { + value: "b" + value: "c" + } + } + feature { + bytes_list { + value: "d" + value: "e" + value: "f" + value: "g" + } + } + } + } + feature_list { + key: "float_list" + value { + feature { + float_list { + value: 1.0 + } + } + feature { + float_list { + value: 3.0 + value: 3.0 + value: 3.0 + } + } + feature { + float_list { + value: 5.0 + value: 5.0 + value: 5.0 + value: 5.0 + value: 5.0 + } + } + } + } + feature_list { + key: "int_list" + value { + feature { + int64_list { + value: 2 + value: 2 + } + } + feature { + int64_list { + value: 4 + value: 4 + value: 4 + value: 4 + } + } + feature { + int64_list { + value: 6 + value: 6 + value: 6 + value: 6 + value: 6 + value: 6 + } + } + } + } +} +""" + + +def _make_sequence_example(): + example = example_pb2.SequenceExample() + return text_format.Parse(_SEQ_EX_PROTO, example) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2_test.py b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2_test.py index be012a87690c24c6d9b7808790393e1aa6d01211..a1feaddcc00d5fac86dca3138dfa1c6314bb6a8b 100644 --- a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2_test.py +++ b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2_test.py @@ -22,9 +22,7 @@ import os from absl.testing import parameterized import numpy as np -from tensorflow.contrib.feature_column.python.feature_column import sequence_feature_column as sfc_old from tensorflow.contrib.feature_column.python.feature_column import sequence_feature_column_v2 as sfc -from tensorflow.python.feature_column import feature_column as fc_old from tensorflow.python.feature_column import feature_column_lib as fc from tensorflow.python.feature_column.feature_column_v2_test import _TestStateManager from tensorflow.python.framework import dtypes @@ -32,13 +30,15 @@ from tensorflow.python.framework import errors 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 lookup_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import sparse_ops +from tensorflow.python.ops import variables as variables_lib from tensorflow.python.platform import test from tensorflow.python.training import monitored_session -class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): +class SequenceFeaturesTest(test.TestCase, parameterized.TestCase): @parameterized.named_parameters( {'testcase_name': '2D', @@ -111,29 +111,27 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): categorical_column_a = sfc.sequence_categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - embedding_column_a = fc_old._embedding_column( + embedding_column_a = fc.embedding_column( categorical_column_a, dimension=embedding_dimension_a, initializer=_get_initializer(embedding_dimension_a, embedding_values_a)) categorical_column_b = sfc.sequence_categorical_column_with_identity( key='bbb', num_buckets=vocabulary_size) - embedding_column_b = fc_old._embedding_column( + embedding_column_b = fc.embedding_column( categorical_column_b, dimension=embedding_dimension_b, initializer=_get_initializer(embedding_dimension_b, embedding_values_b)) - input_layer, sequence_length = sfc.sequence_input_layer( - features={ - 'aaa': sparse_input_a, - 'bbb': sparse_input_b, - }, - # Test that columns are reordered alphabetically. - feature_columns=[embedding_column_b, embedding_column_a]) + # Test that columns are reordered alphabetically. + sequence_input_layer = sfc.SequenceFeatures( + [embedding_column_b, embedding_column_a]) + input_layer, sequence_length = sequence_input_layer({ + 'aaa': sparse_input_a, 'bbb': sparse_input_b,}) global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) self.assertCountEqual( - ('sequence_input_layer/aaa_embedding/embedding_weights:0', - 'sequence_input_layer/bbb_embedding/embedding_weights:0'), + ('sequence_features/aaa_embedding/embedding_weights:0', + 'sequence_features/bbb_embedding/embedding_weights:0'), tuple([v.name for v in global_vars])) with monitored_session.MonitoredSession() as sess: self.assertAllEqual(embedding_values_a, global_vars[0].eval(session=sess)) @@ -152,18 +150,17 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): values=(2, 0, 1), dense_shape=(2, 2)) - categorical_column_a = fc_old._categorical_column_with_identity( + categorical_column_a = fc.categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - embedding_column_a = fc_old._embedding_column( + embedding_column_a = fc.embedding_column( categorical_column_a, dimension=2) with self.assertRaisesRegexp( ValueError, r'In embedding_column: aaa_embedding\. categorical_column must be of ' - r'type _SequenceCategoricalColumn to use sequence_input_layer\.'): - _, _ = sfc.sequence_input_layer( - features={'aaa': sparse_input}, - feature_columns=[embedding_column_a]) + r'type SequenceCategoricalColumn to use SequenceFeatures\.'): + sequence_input_layer = sfc.SequenceFeatures([embedding_column_a]) + _, _ = sequence_input_layer({'aaa': sparse_input}) def test_shared_embedding_column(self): vocabulary_size = 3 @@ -210,21 +207,18 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): categorical_column_b = sfc.sequence_categorical_column_with_identity( key='bbb', num_buckets=vocabulary_size) # Test that columns are reordered alphabetically. - shared_embedding_columns = fc.shared_embedding_columns( + shared_embedding_columns = fc.shared_embedding_columns_v2( [categorical_column_b, categorical_column_a], dimension=embedding_dimension, initializer=_get_initializer(embedding_dimension, embedding_values)) - input_layer, sequence_length = sfc.sequence_input_layer( - features={ - 'aaa': sparse_input_a, - 'bbb': sparse_input_b, - }, - feature_columns=shared_embedding_columns) + sequence_input_layer = sfc.SequenceFeatures(shared_embedding_columns) + input_layer, sequence_length = sequence_input_layer({ + 'aaa': sparse_input_a, 'bbb': sparse_input_b}) global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) self.assertCountEqual( - ('sequence_input_layer/aaa_bbb_shared_embedding/embedding_weights:0',), + ('aaa_bbb_shared_embedding:0',), tuple([v.name for v in global_vars])) with monitored_session.MonitoredSession() as sess: self.assertAllEqual(embedding_values, global_vars[0].eval(session=sess)) @@ -248,23 +242,20 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): values=(2, 0, 1), dense_shape=(2, 2)) - categorical_column_a = fc_old._categorical_column_with_identity( + categorical_column_a = fc.categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - categorical_column_b = fc_old._categorical_column_with_identity( + categorical_column_b = fc.categorical_column_with_identity( key='bbb', num_buckets=vocabulary_size) - shared_embedding_columns = fc.shared_embedding_columns( + shared_embedding_columns = fc.shared_embedding_columns_v2( [categorical_column_a, categorical_column_b], dimension=2) with self.assertRaisesRegexp( ValueError, r'In embedding_column: aaa_shared_embedding\. categorical_column must ' - r'be of type _SequenceCategoricalColumn to use sequence_input_layer\.'): - _, _ = sfc.sequence_input_layer( - features={ - 'aaa': sparse_input_a, - 'bbb': sparse_input_b - }, - feature_columns=shared_embedding_columns) + r'be of type SequenceCategoricalColumn to use SequenceFeatures\.'): + sequence_input_layer = sfc.SequenceFeatures(shared_embedding_columns) + _, _ = sequence_input_layer({'aaa': sparse_input_a, + 'bbb': sparse_input_b}) @parameterized.named_parameters( {'testcase_name': '2D', @@ -319,17 +310,15 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): categorical_column_a = sfc.sequence_categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size_a) - indicator_column_a = fc_old._indicator_column(categorical_column_a) + indicator_column_a = fc.indicator_column(categorical_column_a) categorical_column_b = sfc.sequence_categorical_column_with_identity( key='bbb', num_buckets=vocabulary_size_b) - indicator_column_b = fc_old._indicator_column(categorical_column_b) - input_layer, sequence_length = sfc.sequence_input_layer( - features={ - 'aaa': sparse_input_a, - 'bbb': sparse_input_b, - }, - # Test that columns are reordered alphabetically. - feature_columns=[indicator_column_b, indicator_column_a]) + indicator_column_b = fc.indicator_column(categorical_column_b) + # Test that columns are reordered alphabetically. + sequence_input_layer = sfc.SequenceFeatures( + [indicator_column_b, indicator_column_a]) + input_layer, sequence_length = sequence_input_layer({ + 'aaa': sparse_input_a, 'bbb': sparse_input_b}) with monitored_session.MonitoredSession() as sess: self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess)) @@ -346,17 +335,16 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): values=(2, 0, 1), dense_shape=(2, 2)) - categorical_column_a = fc_old._categorical_column_with_identity( + categorical_column_a = fc.categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - indicator_column_a = fc_old._indicator_column(categorical_column_a) + indicator_column_a = fc.indicator_column(categorical_column_a) with self.assertRaisesRegexp( ValueError, r'In indicator_column: aaa_indicator\. categorical_column must be of ' - r'type _SequenceCategoricalColumn to use sequence_input_layer\.'): - _, _ = sfc.sequence_input_layer( - features={'aaa': sparse_input}, - feature_columns=[indicator_column_a]) + r'type SequenceCategoricalColumn to use SequenceFeatures\.'): + sequence_input_layer = sfc.SequenceFeatures([indicator_column_a]) + _, _ = sequence_input_layer({'aaa': sparse_input}) @parameterized.named_parameters( {'testcase_name': '2D', @@ -375,7 +363,7 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): # feature 0, ids [[20, 3], [5]] # feature 1, ids [[3], [8]] 'indices': ((0, 0, 0), (0, 0, 1), (0, 1, 0), (1, 0, 0), (1, 1, 0)), - 'values': (20, 3, 5., 3., 8.), + 'values': (20., 3., 5., 3., 8.), 'dense_shape': (2, 2, 2)}, 'expected_input_layer': [ [[20.], [3.], [5.], [0.]], @@ -386,11 +374,10 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): self, sparse_input_args, expected_input_layer, expected_sequence_length): sparse_input = sparse_tensor.SparseTensorValue(**sparse_input_args) - numeric_column = sfc_old.sequence_numeric_column('aaa') + numeric_column = sfc.sequence_numeric_column('aaa') - input_layer, sequence_length = sfc.sequence_input_layer( - features={'aaa': sparse_input}, - feature_columns=[numeric_column]) + sequence_input_layer = sfc.SequenceFeatures([numeric_column]) + input_layer, sequence_length = sequence_input_layer({'aaa': sparse_input}) with monitored_session.MonitoredSession() as sess: self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess)) @@ -428,14 +415,13 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): ) def test_numeric_column_multi_dim( self, sparse_input_args, expected_input_layer, expected_sequence_length): - """Tests sequence_input_layer for multi-dimensional numeric_column.""" + """Tests SequenceFeatures for multi-dimensional numeric_column.""" sparse_input = sparse_tensor.SparseTensorValue(**sparse_input_args) - numeric_column = sfc_old.sequence_numeric_column('aaa', shape=(2, 2)) + numeric_column = sfc.sequence_numeric_column('aaa', shape=(2, 2)) - input_layer, sequence_length = sfc.sequence_input_layer( - features={'aaa': sparse_input}, - feature_columns=[numeric_column]) + sequence_input_layer = sfc.SequenceFeatures([numeric_column]) + input_layer, sequence_length = sequence_input_layer({'aaa': sparse_input}) with monitored_session.MonitoredSession() as sess: self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess)) @@ -454,22 +440,20 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): indices=((0, 0), (1, 0)), values=(1., 10.), dense_shape=(2, 2)) - numeric_column_a = sfc_old.sequence_numeric_column('aaa') - numeric_column_b = sfc_old.sequence_numeric_column('bbb') + numeric_column_a = sfc.sequence_numeric_column('aaa') + numeric_column_b = sfc.sequence_numeric_column('bbb') - _, sequence_length = sfc.sequence_input_layer( - features={ - 'aaa': sparse_input_a, - 'bbb': sparse_input_b, - }, - feature_columns=[numeric_column_a, numeric_column_b]) + sequence_input_layer = sfc.SequenceFeatures( + [numeric_column_a, numeric_column_b]) + _, sequence_length = sequence_input_layer({ + 'aaa': sparse_input_a, 'bbb': sparse_input_b}) with monitored_session.MonitoredSession() as sess: with self.assertRaisesRegexp( errors.InvalidArgumentError, r'\[Condition x == y did not hold element-wise:\] ' - r'\[x \(sequence_input_layer/aaa/sequence_length:0\) = \] \[2 1\] ' - r'\[y \(sequence_input_layer/bbb/sequence_length:0\) = \] \[1 1\]'): + r'\[x \(sequence_features/aaa/sequence_length:0\) = \] \[2 1\] ' + r'\[y \(sequence_features/bbb/sequence_length:0\) = \] \[1 1\]'): sess.run(sequence_length) @parameterized.named_parameters( @@ -497,11 +481,10 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): self, sparse_input_args, expected_shape): """Tests that we return a known static shape when we have one.""" sparse_input = sparse_tensor.SparseTensorValue(**sparse_input_args) - numeric_column = sfc_old.sequence_numeric_column('aaa', shape=(2, 2)) + numeric_column = sfc.sequence_numeric_column('aaa', shape=(2, 2)) - input_layer, _ = sfc.sequence_input_layer( - features={'aaa': sparse_input}, - feature_columns=[numeric_column]) + sequence_input_layer = sfc.SequenceFeatures([numeric_column]) + input_layer, _ = sequence_input_layer({'aaa': sparse_input}) shape = input_layer.get_shape() self.assertEqual(shape, expected_shape) @@ -534,13 +517,49 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): sparse_input = sparse_tensor.SparseTensorValue(**sparse_input_args) categorical_column = sfc.sequence_categorical_column_with_identity( key='aaa', num_buckets=3) - indicator_column = fc_old._indicator_column(categorical_column) + indicator_column = fc.indicator_column(categorical_column) - input_layer, _ = sfc.sequence_input_layer( - features={'aaa': sparse_input}, feature_columns=[indicator_column]) + sequence_input_layer = sfc.SequenceFeatures([indicator_column]) + input_layer, _ = sequence_input_layer({'aaa': sparse_input}) shape = input_layer.get_shape() self.assertEqual(shape, expected_shape) + def test_compute_output_shape(self): + price1 = sfc.sequence_numeric_column('price1', shape=2) + price2 = sfc.sequence_numeric_column('price2') + with ops.Graph().as_default(): + features = { + 'price1': sparse_tensor.SparseTensor( + indices=[[0, 0, 0], [0, 0, 1], + [0, 1, 0], [0, 1, 1], + [1, 0, 0], [1, 0, 1], + [2, 0, 0], [2, 0, 1], + [3, 0, 0], [3, 0, 1]], + values=[0., 1., 10., 11., 100., 101., 200., 201., 300., 301.], + dense_shape=(4, 3, 2)), + 'price2': sparse_tensor.SparseTensor( + indices=[[0, 0], + [0, 1], + [1, 0], + [2, 0], + [3, 0]], + values=[10., 11., 20., 30., 40.], + dense_shape=(4, 3))} + sequence_features = sfc.SequenceFeatures([price1, price2]) + seq_input, seq_len = sequence_features(features) + self.assertEqual( + sequence_features.compute_output_shape((None, None)), + (None, None, 3)) + self.evaluate(variables_lib.global_variables_initializer()) + self.evaluate(lookup_ops.tables_initializer()) + + self.assertAllClose([[[0., 1., 10.], [10., 11., 11.], [0., 0., 0.]], + [[100., 101., 20.], [0., 0., 0.], [0., 0., 0.]], + [[200., 201., 30.], [0., 0., 0.], [0., 0., 0.]], + [[300., 301., 40.], [0., 0., 0.], [0., 0., 0.]]], + self.evaluate(seq_input)) + self.assertAllClose([2, 1, 1, 1], self.evaluate(seq_len)) + class ConcatenateContextInputTest(test.TestCase, parameterized.TestCase): """Tests the utility fn concatenate_context_input.""" @@ -605,8 +624,8 @@ class ConcatenateContextInputTest(test.TestCase, parameterized.TestCase): sfc.concatenate_context_input(context_input, seq_input) -class InputLayerTest(test.TestCase): - """Tests input_layer with sequence feature columns.""" +class DenseFeaturesTest(test.TestCase): + """Tests DenseFeatures with sequence feature columns.""" def test_embedding_column(self): """Tests that error is raised for sequence embedding column.""" @@ -620,16 +639,15 @@ class InputLayerTest(test.TestCase): categorical_column_a = sfc.sequence_categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - embedding_column_a = fc_old._embedding_column( + embedding_column_a = fc.embedding_column( categorical_column_a, dimension=2) with self.assertRaisesRegexp( ValueError, r'In embedding_column: aaa_embedding\. categorical_column must not be ' - r'of type _SequenceCategoricalColumn\.'): - _ = fc_old.input_layer( - features={'aaa': sparse_input}, - feature_columns=[embedding_column_a]) + r'of type SequenceCategoricalColumn\.'): + input_layer = fc.DenseFeatures([embedding_column_a]) + _ = input_layer({'aaa': sparse_input}) def test_indicator_column(self): """Tests that error is raised for sequence indicator column.""" @@ -643,15 +661,14 @@ class InputLayerTest(test.TestCase): categorical_column_a = sfc.sequence_categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - indicator_column_a = fc_old._indicator_column(categorical_column_a) + indicator_column_a = fc.indicator_column(categorical_column_a) with self.assertRaisesRegexp( ValueError, r'In indicator_column: aaa_indicator\. categorical_column must not be ' - r'of type _SequenceCategoricalColumn\.'): - _ = fc_old.input_layer( - features={'aaa': sparse_input}, - feature_columns=[indicator_column_a]) + r'of type SequenceCategoricalColumn\.'): + input_layer = fc.DenseFeatures([indicator_column_a]) + _ = input_layer({'aaa': sparse_input}) def _assert_sparse_tensor_value(test_case, expected, actual): @@ -946,7 +963,7 @@ class SequenceEmbeddingColumnTest( embedding_column, {'aaa': inputs}) global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) - self.assertItemsEqual( + self.assertCountEqual( ('embedding_weights:0',), tuple([v.name for v in global_vars])) with monitored_session.MonitoredSession() as sess: self.assertAllEqual(embedding_values, global_vars[0].eval(session=sess)) diff --git a/tensorflow/contrib/framework/BUILD b/tensorflow/contrib/framework/BUILD index 88a14a2a94cc683f021d032ea11358e0cfb63faa..8fd2b5f39bc88b76fe5583f8d18389e232ea9f40 100644 --- a/tensorflow/contrib/framework/BUILD +++ b/tensorflow/contrib/framework/BUILD @@ -32,7 +32,6 @@ tf_custom_op_py_library( "python/ops/arg_scope.py", "python/ops/audio_ops.py", "python/ops/checkpoint_ops.py", - "python/ops/critical_section_ops.py", "python/ops/ops.py", "python/ops/prettyprint_ops.py", "python/ops/script_ops.py", @@ -51,6 +50,7 @@ tf_custom_op_py_library( "//learning/brain:__subpackages__", "//tensorflow:__subpackages__", "//tensorflow_estimator:__subpackages__", + "//tensorflow_model_optimization:__subpackages__", "//video/youtube/personalization:__subpackages__", ], deps = [ @@ -171,26 +171,6 @@ py_test( ], ) -cuda_py_test( - name = "critical_section_test", - size = "medium", - srcs = ["python/ops/critical_section_test.py"], - additional_deps = [ - "//tensorflow/python:client_testlib", - ":framework_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:control_flow_ops", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:gradients", - "//tensorflow/python:platform_test", - "//tensorflow/python:resource_variable_ops", - "//tensorflow/python:tensor_array_ops", - "//tensorflow/python/data/ops:dataset_ops", - "//tensorflow/python/eager:context", - ], -) - py_test( name = "ops_test", size = "small", diff --git a/tensorflow/contrib/framework/__init__.py b/tensorflow/contrib/framework/__init__.py index e72e50585a3861d4527b66f89e1659d76c85960a..94fb35b3346ecd64cec5a89e495c7a2d1af3584b 100644 --- a/tensorflow/contrib/framework/__init__.py +++ b/tensorflow/contrib/framework/__init__.py @@ -94,8 +94,6 @@ @@smart_constant_value @@smart_case -@@CriticalSection - @@BoundedTensorSpec @@TensorSpec @@ -130,17 +128,22 @@ _allowed_symbols = ['nest'] _nest_allowed_symbols = [ 'assert_same_structure', 'is_sequence', + 'is_sequence_or_composite', 'flatten', 'flatten_dict_items', 'pack_sequence_as', 'map_structure', 'map_structure_with_paths', + 'map_structure_with_tuple_paths', 'assert_shallow_structure', 'flatten_up_to', + 'flatten_with_tuple_paths_up_to', 'map_structure_up_to', + 'map_structure_with_tuple_paths_up_to', 'get_traverse_shallow_structure', 'yield_flat_paths', 'flatten_with_joined_string_paths', + 'flatten_with_tuple_paths', ] remove_undocumented(nest.__name__, allowed_exception_list=_nest_allowed_symbols) diff --git a/tensorflow/contrib/framework/python/ops/__init__.py b/tensorflow/contrib/framework/python/ops/__init__.py index c4976497f5fa95d82e492153b117681f693eaa13..8113bf7c095bd0817e40cfd08bdf1ef7275ba55b 100644 --- a/tensorflow/contrib/framework/python/ops/__init__.py +++ b/tensorflow/contrib/framework/python/ops/__init__.py @@ -22,7 +22,6 @@ from __future__ import print_function # pylint: disable=wildcard-import from tensorflow.contrib.framework.python.ops.arg_scope import * from tensorflow.contrib.framework.python.ops.checkpoint_ops import * -from tensorflow.contrib.framework.python.ops.critical_section_ops import * from tensorflow.contrib.framework.python.ops.ops import * from tensorflow.contrib.framework.python.ops.prettyprint_ops import * from tensorflow.contrib.framework.python.ops.script_ops import * diff --git a/tensorflow/contrib/fused_conv/BUILD b/tensorflow/contrib/fused_conv/BUILD index 57a5bfbf43c915775c6b0ef05baac19581213a09..f65f450eba49163c319af54ec2bd7f6b61e34c1e 100644 --- a/tensorflow/contrib/fused_conv/BUILD +++ b/tensorflow/contrib/fused_conv/BUILD @@ -171,6 +171,7 @@ cuda_py_test( main = "python/ops/fused_conv2d_bias_activation_benchmark.py", tags = [ "manual", # TODO(b/117128481): re-enable after fixing OSS build + "nogpu", "requires-gpu-sm70", ], ) diff --git a/tensorflow/contrib/fused_conv/kernels/fused_conv2d_bias_activation_op.cc b/tensorflow/contrib/fused_conv/kernels/fused_conv2d_bias_activation_op.cc index c541c71f996c7a1b36cf28ae9a1783f8dca0a72c..f13a66717f67a1a627f66af9468c6f2897aaf7a4 100644 --- a/tensorflow/contrib/fused_conv/kernels/fused_conv2d_bias_activation_op.cc +++ b/tensorflow/contrib/fused_conv/kernels/fused_conv2d_bias_activation_op.cc @@ -19,13 +19,13 @@ limitations under the License. #include "tensorflow/contrib/fused_conv/kernels/fused_conv2d_bias_activation_op.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/numeric_op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_slice.h" -#include "tensorflow/core/kernels/bounds_check.h" #include "tensorflow/core/kernels/conv_2d.h" #include "tensorflow/core/kernels/ops_util.h" #include "tensorflow/core/lib/core/errors.h" @@ -565,6 +565,20 @@ void LaunchFusedConv2DBiasActivationOp:: fused_conv_parameters.ShouldIncludeWinogradNonfusedAlgo( stream->parent()), &algorithms)); + if (activation_mode == ActivationMode::NONE) { + // Only CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM is supported for + // identity activation, other algs seem to quietly do Relu. + // See + // https://docs.nvidia.com/deeplearning/sdk/cudnn-developer-guide/index.html#cudnnConvolutionBiasActivationForward + algorithms.erase( + std::remove_if( + algorithms.begin(), algorithms.end(), + [](dnn::AlgorithmDesc alg) { + return alg.algo_id() != + CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM; + }), + algorithms.end()); + } dnn::ProfileResult best_result; dnn::ProfileResult best_result_no_scratch; for (auto profile_algorithm : algorithms) { diff --git a/tensorflow/contrib/gan/BUILD b/tensorflow/contrib/gan/BUILD index 97184dabb051b2ff4403d8a8ed73f76234f9c6cd..db0868fb2c43464a811b3d6dfcd96480ba2463ee 100644 --- a/tensorflow/contrib/gan/BUILD +++ b/tensorflow/contrib/gan/BUILD @@ -2,7 +2,6 @@ load("//tensorflow:tensorflow.bzl", "py_test") package(default_visibility = [ - "//learning/brain/contrib/tfgan:__subpackages__", "//tensorflow:__subpackages__", ]) @@ -107,6 +106,7 @@ py_library( deps = [ ":gan_estimator", ":head", + ":latent_gan_estimator", ":stargan_estimator", ":tpu_gan_estimator", "//tensorflow/python:util", @@ -132,6 +132,7 @@ py_library( ":clip_weights", ":conditioning_utils", ":random_tensor_pool", + ":spectral_normalization", ":virtual_batchnorm", "//tensorflow/python:util", ], @@ -642,6 +643,41 @@ py_test( ], ) +py_library( + name = "latent_gan_estimator", + srcs = [ + "python/estimator/python/latent_gan_estimator.py", + "python/estimator/python/latent_gan_estimator_impl.py", + ], + srcs_version = "PY2AND3", + deps = [ + ":train", + "//tensorflow/python:clip_ops", + "//tensorflow/python:gradients", + "//tensorflow/python:random_ops", + "//tensorflow/python:summary", + "//tensorflow/python:training_util", + "//tensorflow/python:variable_scope", + "//tensorflow/python/estimator:estimator_py", + ], +) + +py_test( + name = "latent_gan_estimator_test", + srcs = [ + "python/estimator/python/latent_gan_estimator_test.py", + ], + srcs_version = "PY2AND3", + deps = [ + ":latent_gan_estimator", + "//tensorflow/python:array_ops", + "//tensorflow/python:training", + "//tensorflow/python:variable_scope", + "//tensorflow/python/estimator:run_config", + "//tensorflow/python/ops/losses", + ], +) + py_library( name = "sliced_wasserstein", srcs = [ @@ -676,3 +712,45 @@ py_test( "//third_party/py/numpy", ], ) + +py_library( + name = "spectral_normalization", + srcs = [ + "python/features/python/spectral_normalization.py", + "python/features/python/spectral_normalization_impl.py", + ], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:array_ops", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:init_ops", + "//tensorflow/python:math_ops", + "//tensorflow/python:standard_ops", + "//tensorflow/python:variable_scope", + "//tensorflow/python/keras:engine", + ], +) + +py_test( + name = "spectral_normalization_test", + srcs = ["python/features/python/spectral_normalization_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":spectral_normalization", + "//tensorflow/contrib/layers:layers_py", + "//tensorflow/contrib/slim", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:init_ops", + "//tensorflow/python:layers", + "//tensorflow/python:linalg_ops", + "//tensorflow/python:math_ops", + "//tensorflow/python:variable_scope", + "//tensorflow/python:variables", + "//tensorflow/python/keras:layers", + "//third_party/py/numpy", + ], +) diff --git a/tensorflow/contrib/gan/README.md b/tensorflow/contrib/gan/README.md index db7dc51daa78ecee12ecb7f6d33df4511e068243..4eac4e80cdacd779fdbedef19e4a654196f0caf1 100644 --- a/tensorflow/contrib/gan/README.md +++ b/tensorflow/contrib/gan/README.md @@ -57,11 +57,10 @@ These include the following main pieces (explained in detail below). generative models. * [examples](https://github.com/tensorflow/models/tree/master/research/gan/) - and [tutorial](http://https://github.com/tensorflow/models/tree/master/research/gan/tutorial.ipynb): See examples of how to use TF-GAN - to make GAN training easier, or use the more complicated examples to - jumpstart your own project. These include unconditional and conditional - GANs, InfoGANs, adversarial losses on existing networks, and image-to-image - translation. + and [tutorial](https://github.com/tensorflow/models/tree/master/research/gan/tutorial.ipynb): See examples of how to use TF-GAN to make + GAN training easier, or use the more complicated examples to jump-start your + own project. These include unconditional and conditional GANs, InfoGANs, + adversarial losses on existing networks, and image-to-image translation. ## Training a GAN model diff --git a/tensorflow/contrib/gan/python/estimator/__init__.py b/tensorflow/contrib/gan/python/estimator/__init__.py index 03a639165badfb80fc9f0f77082770fd29076ec5..430266555b723e6ca39dccffc1442dbef5d4a385 100644 --- a/tensorflow/contrib/gan/python/estimator/__init__.py +++ b/tensorflow/contrib/gan/python/estimator/__init__.py @@ -26,11 +26,13 @@ from __future__ import print_function # pylint: disable=unused-import,wildcard-import from tensorflow.contrib.gan.python.estimator.python import gan_estimator from tensorflow.contrib.gan.python.estimator.python import head +from tensorflow.contrib.gan.python.estimator.python import latent_gan_estimator from tensorflow.contrib.gan.python.estimator.python import stargan_estimator from tensorflow.contrib.gan.python.estimator.python import tpu_gan_estimator from tensorflow.contrib.gan.python.estimator.python.gan_estimator import * from tensorflow.contrib.gan.python.estimator.python.head import * +from tensorflow.contrib.gan.python.estimator.python.latent_gan_estimator import * from tensorflow.contrib.gan.python.estimator.python.stargan_estimator import * from tensorflow.contrib.gan.python.estimator.python.tpu_gan_estimator import * # pylint: enable=unused-import,wildcard-import @@ -41,7 +43,8 @@ _allowed_symbols = ([ 'gan_estimator', 'stargan_estimator', 'tpu_gan_estimator', + 'latent_gan_estimator', 'head', ] + gan_estimator.__all__ + stargan_estimator.__all__ + head.__all__ + - tpu_gan_estimator.__all__) + tpu_gan_estimator.__all__ + latent_gan_estimator.__all__) remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/gan/python/estimator/python/latent_gan_estimator.py b/tensorflow/contrib/gan/python/estimator/python/latent_gan_estimator.py new file mode 100644 index 0000000000000000000000000000000000000000..4e164e24168bb0cc5e9a7cc772081781ea088bb1 --- /dev/null +++ b/tensorflow/contrib/gan/python/estimator/python/latent_gan_estimator.py @@ -0,0 +1,28 @@ +# 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. +# ============================================================================== +"""`tf.Learn` components for `Train Input Estimator`.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.gan.python.estimator.python import latent_gan_estimator_impl +# pylint: disable=wildcard-import +from tensorflow.contrib.gan.python.estimator.python.latent_gan_estimator_impl import * +# pylint: enable=wildcard-import +from tensorflow.python.util.all_util import remove_undocumented + +__all__ = latent_gan_estimator_impl.__all__ +remove_undocumented(__name__, __all__) diff --git a/tensorflow/contrib/gan/python/estimator/python/latent_gan_estimator_impl.py b/tensorflow/contrib/gan/python/estimator/python/latent_gan_estimator_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..f5afc7731937ed1a82c8ebb5969b2687ffdd583b --- /dev/null +++ b/tensorflow/contrib/gan/python/estimator/python/latent_gan_estimator_impl.py @@ -0,0 +1,205 @@ +# 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. +# ============================================================================== +"""Implements an estimator wrapper that allows training the input latent space. + +This file implements a latent gan estimator that wraps around a previously +trained GAN. The latent gan estimator trains a single variable z, representing +the hidden latent distribution that is the 'noise' input to the GAN. By training +z, the inpainting estimator can move around the latent z space towards +minimizing a specific loss function. + +The latent gan estimator has a few key differences from a normal estimator. + +First: the variables in the estimator should not be saved, as we are not +updating the original GAN and are only adding a new z variable that is meant +to be different for each run. In order to do distributed training using +train_and_evaluate, the Tensorflow RunConfig is expected to save checkpoints +by having either save_checkpoints_steps or save_checkpoints_secs saved. +To avoid this conflict, we purposely set the save_checkpoints_steps value in +the RunConfig to be one step more than the total number of steps that the +inpainter estimator will run. + +Second: we need to specify warm start settings, as we are reloading the +GAN model into a different graph (specifically, one with a new z variable). +The warm start settings defined below reload all GAN variables and ignore the +new z variable (and the optimizer). + +Usage: + + def _generator(net, mode): + ... + + def _discriminator(net, condition, mode): + ... + + def _loss(gan_model, features, labels, add_summaries): + ... + + def optimizer(): + ... + + params = {} + config = tf.estimator.RunConfig() + tmp_dir = path/to/output/storage + + estimator = latent_gan_estimator.get_latent_gan_estimator( + _generator, _discriminator, _loss, optimizer, params, config, tmp_dir) + + def input_fn(): + ... + + estimator.train(input_fn=input_fn) + +See latent_gan_estimator_test.py or tensorflow_models/gan/face_inpainting for +further examples. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import functools +from tensorflow.contrib.gan.python import train as tfgan_train +from tensorflow.python.estimator import estimator +from tensorflow.python.estimator import model_fn as model_fn_lib +from tensorflow.python.ops import clip_ops +from tensorflow.python.ops import gradients_impl +from tensorflow.python.ops import random_ops +from tensorflow.python.ops import variable_scope +from tensorflow.python.summary import summary +from tensorflow.python.training import training_util + + +INPUT_NAME = 'new_var_z_input' # The name for the new z space input variable. +OPTIMIZER_NAME = 'latent_gan_optimizer' # The name for the new optimizer vars. + +__all__ = [ + 'get_latent_gan_estimator', +] + + +def _get_latent_gan_model_fn(generator_fn, discriminator_fn, loss_fn, + optimizer): + """Sets up a model function that wraps around a given GAN.""" + def model_fn(features, labels, mode, params): + """Model function defining an inpainting estimator.""" + batch_size = params['batch_size'] + z_shape = [batch_size] + params['z_shape'] + add_summaries = params['add_summaries'] + input_clip = params['input_clip'] + + z = variable_scope.get_variable( + name=INPUT_NAME, initializer=random_ops.truncated_normal(z_shape), + constraint=lambda x: clip_ops.clip_by_value(x, -input_clip, input_clip)) + + generator = functools.partial(generator_fn, mode=mode) + discriminator = functools.partial(discriminator_fn, mode=mode) + gan_model = tfgan_train.gan_model(generator_fn=generator, + discriminator_fn=discriminator, + real_data=labels, + generator_inputs=z, + check_shapes=False) + + loss = loss_fn(gan_model, features, labels, add_summaries) + + # Use a variable scope to make sure that estimator variables dont cause + # save/load problems when restoring from ckpts. + with variable_scope.variable_scope(OPTIMIZER_NAME): + opt = optimizer(learning_rate=params['learning_rate'], + **params['opt_kwargs']) + train_op = opt.minimize( + loss=loss, global_step=training_util.get_or_create_global_step(), + var_list=[z]) + + if add_summaries: + z_grads = gradients_impl.gradients(loss, z) + summary.scalar('z_loss/z_grads', clip_ops.global_norm(z_grads)) + summary.scalar('z_loss/loss', loss) + + return model_fn_lib.EstimatorSpec(mode=mode, + predictions=gan_model.generated_data, + loss=loss, + train_op=train_op) + return model_fn + + +def get_latent_gan_estimator(generator_fn, discriminator_fn, loss_fn, + optimizer, params, config, ckpt_dir, + warmstart_options=True): + """Gets an estimator that passes gradients to the input. + + This function takes in a generator and adds a trainable z variable that is + used as input to this generator_fn. The generator itself is treated as a black + box through which gradients can pass through without updating any weights. The + result is a trainable way to traverse the GAN latent space. The loss_fn is + used to actually train the z variable. The generator_fn and discriminator_fn + should be previously trained by the tfgan library (on reload, the variables + are expected to follow the tfgan format. It may be possible to use the + latent gan estimator with entirely custom GANs that do not use the tfgan + library as long as the appropriate variables are wired properly). + + Args: + generator_fn: a function defining a Tensorflow graph for a GAN generator. + The weights defined in this graph should already be defined in the given + checkpoint location. Should have 'mode' as an argument. + discriminator_fn: a function defining a Tensorflow graph for a GAN + discriminator. Should have 'mode' as an argument. + loss_fn: a function defining a Tensorflow graph for a GAN loss. Takes in a + GANModel tuple, features, labels, and add_summaries as inputs. + optimizer: a tf.Optimizer or a function that returns a tf.Optimizer with no + inputs. + params: An object containing the following parameters: + - batch_size: an int indicating the size of the training batch. + - z_shape: the desired shape of the input z values (not counting batch). + - learning_rate: a scalar or function defining a learning rate applied to + optimizer. + - input_clip: the amount to clip the x training variable by. + - add_summaries: whether or not to add summaries. + - opt_kwargs: optimizer kwargs. + config: tf.RunConfig. Should point model to output dir and should indicate + whether to save checkpoints (to avoid saving checkpoints, set + save_checkpoints_steps to a number larger than the number of train steps). + The model_dir field in the RunConfig should point to a directory WITHOUT + any saved checkpoints. + ckpt_dir: the directory where the model checkpoints live. The checkpoint is + used to warm start the underlying GAN. This should NOT be the same as + config.model_dir. + warmstart_options: boolean, None, or a WarmStartSettings object. If set to + True, uses a default WarmStartSettings object. If set to False or None, + does not use warm start. If using a custom WarmStartSettings object, make + sure that new variables are properly accounted for when reloading the + underlying GAN. Defaults to True. + Returns: + An estimator spec defining a GAN input training estimator. + """ + model_fn = _get_latent_gan_model_fn(generator_fn, discriminator_fn, + loss_fn, optimizer) + + if isinstance(warmstart_options, estimator.WarmStartSettings): + ws = warmstart_options + elif warmstart_options: + # Default WarmStart loads all variable names except INPUT_NAME and + # OPTIMIZER_NAME. + var_regex = '^(?!.*(%s|%s).*)' % (INPUT_NAME, OPTIMIZER_NAME) + ws = estimator.WarmStartSettings(ckpt_to_initialize_from=ckpt_dir, + vars_to_warm_start=var_regex) + else: + ws = None + + if 'opt_kwargs' not in params: + params['opt_kwargs'] = {} + + return estimator.Estimator(model_fn=model_fn, config=config, params=params, + warm_start_from=ws) diff --git a/tensorflow/contrib/gan/python/estimator/python/latent_gan_estimator_test.py b/tensorflow/contrib/gan/python/estimator/python/latent_gan_estimator_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ac139e532e35f7aae6da0655103a7249fe3382d4 --- /dev/null +++ b/tensorflow/contrib/gan/python/estimator/python/latent_gan_estimator_test.py @@ -0,0 +1,119 @@ +# 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 latent_gan_estimator. + +See g3.tp.tensorflow.contrib.gan.python.estimator.python.latent_gan_estimator. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import tempfile +import numpy as np +from tensorflow.contrib.gan.python.estimator.python import latent_gan_estimator +from tensorflow.python.estimator import run_config as run_config +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import variable_scope +from tensorflow.python.ops.losses import losses +from tensorflow.python.platform import test +from tensorflow.python.training import training + + +class TrainInputEstimatorTest(test.TestCase): + + def test_get_input_training_estimator(self): + """Integration test to make sure the input_training_estimator works.""" + + # Create dummy test input tensors. + true_features = np.reshape(np.random.uniform(size=100), (10, 10)) + true_labels = np.reshape(np.random.uniform(size=100), (5, 20)) + expected_z_output = [[1, -1], [-1, 1]] + + # Fill out required parameters randomly, includes optimizer kwargs. + params = { + 'batch_size': 2, + 'z_shape': [2], + 'learning_rate': 1.0, + 'input_clip': 1.0, + 'add_summaries': False, + 'opt_kwargs': { + 'beta1': 0.1 + } + } + + input_z_shape = [params['batch_size']] + params['z_shape'] + + # Create dummy model functions that represent an underlying GANEstimator and + # the input training wrapper. Make sure that everything is wired up + # correctly in the internals of each dummy function. + def _generator(net, mode): + """The generator function will get the newly created z variable.""" + del mode + self.assertSequenceEqual(net.shape, input_z_shape) + gen_dummy_var = variable_scope.get_variable( + name='generator_dummy_variable', + initializer=array_ops.ones(input_z_shape)) + return net * gen_dummy_var + + def _discriminator(net, condition, mode): + """The discriminator function will get either the z variable or labels.""" + del condition, mode + try: + self.assertSequenceEqual(net.shape, true_labels.shape) + except AssertionError: + self.assertSequenceEqual(net.shape, input_z_shape) + return net + + def _loss(gan_model, features, labels, _): + """Make sure that features and labels are passed in from input.""" + self.assertTrue(np.array_equal(features, true_features)) + self.assertTrue(np.array_equal(labels, true_labels)) + return losses.absolute_difference(expected_z_output, + gan_model.generated_data) + + optimizer = training.AdamOptimizer + + # We are not loading checkpoints, so set the corresponding directory to a + # dummy directories. + tmp_dir = tempfile.mkdtemp() + config = run_config.RunConfig(model_dir=tmp_dir, + save_summary_steps=None, + save_checkpoints_steps=1, + save_checkpoints_secs=None) + + # Get the estimator. Disable warm start so that there is no attempted + # checkpoint reloading. + estimator = latent_gan_estimator.get_latent_gan_estimator( + _generator, _discriminator, _loss, optimizer, params, config, tmp_dir, + warmstart_options=None) + + # Train for a few steps. + def dummy_input(): + return true_features, true_labels + estimator.train(input_fn=dummy_input, steps=10) + + # Make sure the generator variables did not change, but the z variables did + # change. + self.assertTrue(np.array_equal( + estimator.get_variable_value('Generator/generator_dummy_variable'), + np.ones(input_z_shape))) + self.assertTrue(np.array_equal( + estimator.get_variable_value('new_var_z_input'), + expected_z_output)) + + +if __name__ == '__main__': + test.main() 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 9fdcc08334d50b4ddf3a0bc9bc755e55d51b0bd8..0d9e6489bdd1d89cc49bfedc2eed784999c31d2b 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 @@ -184,8 +184,7 @@ class TPUGANEstimatorIntegrationTest(test.TestCase, parameterized.TestCase): # Evaluate. num_steps_eval = 2 scores = est.evaluate(eval_input_fn, steps=num_steps_eval) - self.assertEqual(num_steps_train + num_steps_eval, - scores[ops.GraphKeys.GLOBAL_STEP]) + self.assertIn(ops.GraphKeys.GLOBAL_STEP, six.iterkeys(scores)) self.assertIn('loss', six.iterkeys(scores)) self.assertEqual(scores['discriminator_loss'] + scores['generator_loss'], scores['loss']) diff --git a/tensorflow/contrib/gan/python/eval/python/classifier_metrics_impl.py b/tensorflow/contrib/gan/python/eval/python/classifier_metrics_impl.py index 31f0d34ed68a6adc25cca102236079d0f66615cb..ff19ce2f78e9c86400089e454c88450f01c41764 100644 --- a/tensorflow/contrib/gan/python/eval/python/classifier_metrics_impl.py +++ b/tensorflow/contrib/gan/python/eval/python/classifier_metrics_impl.py @@ -41,9 +41,9 @@ from tensorflow.python.framework import importer from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops -from tensorflow.python.ops import functional_ops from tensorflow.python.ops import image_ops from tensorflow.python.ops import linalg_ops +from tensorflow.python.ops import map_fn from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_impl from tensorflow.python.ops import nn_ops @@ -346,7 +346,7 @@ def classifier_score(images, classifier_fn, num_batches=1): images, num_or_size_splits=num_batches) # Compute the classifier splits using the memory-efficient `map_fn`. - logits = functional_ops.map_fn( + logits = map_fn.map_fn( fn=classifier_fn, elems=array_ops.stack(generated_images_list), parallel_iterations=1, @@ -505,12 +505,12 @@ def frechet_classifier_distance(real_images, # Compute the activations using the memory-efficient `map_fn`. def compute_activations(elems): - return functional_ops.map_fn(fn=classifier_fn, - elems=elems, - parallel_iterations=1, - back_prop=False, - swap_memory=True, - name='RunClassifier') + return map_fn.map_fn(fn=classifier_fn, + elems=elems, + parallel_iterations=1, + back_prop=False, + swap_memory=True, + name='RunClassifier') real_a = compute_activations(real_imgs) gen_a = compute_activations(generated_imgs) @@ -895,7 +895,7 @@ def kernel_classifier_distance_and_std(real_images, # Compute the activations using the memory-efficient `map_fn`. def compute_activations(elems): - return functional_ops.map_fn( + return map_fn.map_fn( fn=classifier_fn, elems=elems, parallel_iterations=1, @@ -1099,7 +1099,7 @@ def kernel_classifier_distance_and_std_from_activations(real_activations, (math_ops.reduce_sum(k_rr) - math_ops.trace(k_rr)) / (m * (m - 1)) + (math_ops.reduce_sum(k_gg) - math_ops.trace(k_gg)) / (n * (n - 1))) - ests = functional_ops.map_fn( + ests = map_fn.map_fn( compute_kid_block, math_ops.range(n_blocks), dtype=dtype, back_prop=False) mn = math_ops.reduce_mean(ests) diff --git a/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py b/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py index bd17571a0535a3c8e9dfee24a8da16eb2e72f165..bc7c1057b478fe2656898e68c1a14013b5a71d12 100644 --- a/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py +++ b/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py @@ -365,7 +365,7 @@ class ClassifierMetricsTest(test.TestCase, parameterized.TestCase): unused_image = array_ops.zeros([2, 299, 299, 3]) incscore = _run_with_mock(classifier_metrics.inception_score, unused_image) - with self.test_session(use_gpu=True) as sess: + with self.cached_session(use_gpu=True) as sess: incscore_np = sess.run(incscore, {'concat:0': logits}) self.assertAllClose(_expected_inception_score(logits), incscore_np) @@ -473,7 +473,7 @@ class ClassifierMetricsTest(test.TestCase, parameterized.TestCase): classifier_fn=lambda x: x, max_block_size=600) - with self.test_session() as sess: + with self.cached_session() as sess: actual_kid, actual_std = sess.run(kid_op) expected_kid, expected_std = _expected_kid_and_std(test_pool_real_a, @@ -500,7 +500,7 @@ class ClassifierMetricsTest(test.TestCase, parameterized.TestCase): max_block_size=max_block_size) for block_size in [50, 512, 1000]: - with self.test_session() as sess: + with self.cached_session() as sess: actual_kid, actual_std = sess.run(kid_op, {max_block_size: block_size}) expected_kid, expected_std = _expected_kid_and_std( diff --git a/tensorflow/contrib/gan/python/eval/python/summaries_impl.py b/tensorflow/contrib/gan/python/eval/python/summaries_impl.py index 9f448d3a1602c503093214201bdc96fc9bee85b5..c7bbd65bbff41c25327733ae1f17a090fb69cb52 100644 --- a/tensorflow/contrib/gan/python/eval/python/summaries_impl.py +++ b/tensorflow/contrib/gan/python/eval/python/summaries_impl.py @@ -22,7 +22,7 @@ from tensorflow.contrib.gan.python import namedtuples from tensorflow.contrib.gan.python.eval.python import eval_utils from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops -from tensorflow.python.ops import functional_ops +from tensorflow.python.ops import map_fn from tensorflow.python.ops import math_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops.losses import util as loss_util @@ -261,7 +261,7 @@ def add_stargan_image_summaries(stargan_model, summary.image( 'stargan_image_generation', - functional_ops.map_fn( + map_fn.map_fn( _build_image, stargan_model.input_data[:num_images], parallel_iterations=num_images, diff --git a/tensorflow/contrib/gan/python/features/__init__.py b/tensorflow/contrib/gan/python/features/__init__.py index 4816daf760143af9f1502873b123ffad8e5ec8ce..410c3a02052cd3a07a36a0ba332a80b3c2705d89 100644 --- a/tensorflow/contrib/gan/python/features/__init__.py +++ b/tensorflow/contrib/gan/python/features/__init__.py @@ -27,11 +27,13 @@ from __future__ import print_function from tensorflow.contrib.gan.python.features.python import clip_weights from tensorflow.contrib.gan.python.features.python import conditioning_utils from tensorflow.contrib.gan.python.features.python import random_tensor_pool +from tensorflow.contrib.gan.python.features.python import spectral_normalization from tensorflow.contrib.gan.python.features.python import virtual_batchnorm from tensorflow.contrib.gan.python.features.python.clip_weights import * from tensorflow.contrib.gan.python.features.python.conditioning_utils import * from tensorflow.contrib.gan.python.features.python.random_tensor_pool import * +from tensorflow.contrib.gan.python.features.python.spectral_normalization import * from tensorflow.contrib.gan.python.features.python.virtual_batchnorm import * # pylint: enable=unused-import,wildcard-import @@ -40,5 +42,6 @@ from tensorflow.python.util.all_util import remove_undocumented _allowed_symbols = clip_weights.__all__ _allowed_symbols += conditioning_utils.__all__ _allowed_symbols += random_tensor_pool.__all__ +_allowed_symbols += spectral_normalization.__all__ _allowed_symbols += virtual_batchnorm.__all__ remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/gan/python/features/python/spectral_normalization.py b/tensorflow/contrib/gan/python/features/python/spectral_normalization.py new file mode 100644 index 0000000000000000000000000000000000000000..54d3d0a218dec3588844333cd47e1f92489d8df9 --- /dev/null +++ b/tensorflow/contrib/gan/python/features/python/spectral_normalization.py @@ -0,0 +1,32 @@ +# 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. +# ============================================================================== +"""Keras-like layers and utilities that implement Spectral Normalization. + +Based on "Spectral Normalization for Generative Adversarial Networks" by Miyato, +et al in ICLR 2018. https://openreview.net/pdf?id=B1QRgziT- +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.gan.python.features.python import spectral_normalization_impl +# pylint: disable=wildcard-import +from tensorflow.contrib.gan.python.features.python.spectral_normalization_impl import * +# pylint: enable=wildcard-import +from tensorflow.python.util.all_util import remove_undocumented + +__all__ = spectral_normalization_impl.__all__ +remove_undocumented(__name__, __all__) diff --git a/tensorflow/contrib/gan/python/features/python/spectral_normalization_impl.py b/tensorflow/contrib/gan/python/features/python/spectral_normalization_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..0cc653f0a7907f407e66add5537d1e0a5adb6d8b --- /dev/null +++ b/tensorflow/contrib/gan/python/features/python/spectral_normalization_impl.py @@ -0,0 +1,315 @@ +# 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. +# ============================================================================== +"""Keras-like layers and utilities that implement Spectral Normalization. + +Based on "Spectral Normalization for Generative Adversarial Networks" by Miyato, +et al in ICLR 2018. https://openreview.net/pdf?id=B1QRgziT- +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import contextlib +import numbers +import re + +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.keras.engine import base_layer_utils as keras_base_layer_utils +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import init_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import nn +from tensorflow.python.ops import variable_scope +from tensorflow.python.platform import tf_logging as logging + +__all__ = [ + 'compute_spectral_norm', 'spectral_normalize', 'spectral_norm_regularizer', + 'spectral_normalization_custom_getter', 'keras_spectral_normalization' +] + +# tf.bfloat16 should work, but tf.matmul converts those to tf.float32 which then +# can't directly be assigned back to the tf.bfloat16 variable. +_OK_DTYPES_FOR_SPECTRAL_NORM = (dtypes.float16, dtypes.float32, dtypes.float64) +_PERSISTED_U_VARIABLE_SUFFIX = 'spectral_norm_u' + + +def compute_spectral_norm(w_tensor, power_iteration_rounds=1, name=None): + """Estimates the largest singular value in the weight tensor. + + 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. + name: An optional scope name. + + Returns: + The largest singular value (the spectral norm) of w. + """ + with variable_scope.variable_scope(name, 'spectral_norm'): + # The paper says to flatten convnet kernel weights from + # (C_out, C_in, KH, KW) to (C_out, C_in * KH * KW). But TensorFlow's Conv2D + # kernel weight shape is (KH, KW, C_in, C_out), so it should be reshaped to + # (KH * KW * C_in, C_out), and similarly for other layers that put output + # channels as last dimension. + # n.b. this means that w here is equivalent to w.T in the paper. + w = array_ops.reshape(w_tensor, (-1, w_tensor.get_shape()[-1])) + + # Persisted approximation of first left singular vector of matrix `w`. + u_var = variable_scope.get_variable( + _PERSISTED_U_VARIABLE_SUFFIX, + shape=(w.shape[0], 1), + dtype=w.dtype, + initializer=init_ops.random_normal_initializer(), + trainable=False) + u = u_var + + # Use power iteration method to approximate spectral norm. + for _ in range(power_iteration_rounds): + # `v` approximates the first right singular vector of matrix `w`. + v = nn.l2_normalize(math_ops.matmul(array_ops.transpose(w), u)) + u = nn.l2_normalize(math_ops.matmul(w, v)) + + # Update persisted approximation. + with ops.control_dependencies([u_var.assign(u, name='update_u')]): + u = array_ops.identity(u) + + u = array_ops.stop_gradient(u) + v = array_ops.stop_gradient(v) + + # Largest singular value of `w`. + spectral_norm = math_ops.matmul( + math_ops.matmul(array_ops.transpose(u), w), v) + spectral_norm.shape.assert_is_fully_defined() + spectral_norm.shape.assert_is_compatible_with([1, 1]) + + return spectral_norm[0][0] + + +def spectral_normalize(w, power_iteration_rounds=1, name=None): + """Normalizes a weight matrix by its spectral norm. + + 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. + name: An optional scope name. + + Returns: + A normalized weight matrix tensor. + """ + with variable_scope.variable_scope(name, 'spectral_normalize'): + w_normalized = w / compute_spectral_norm( + w, power_iteration_rounds=power_iteration_rounds) + return array_ops.reshape(w_normalized, w.get_shape()) + + +def spectral_norm_regularizer(scale, power_iteration_rounds=1, scope=None): + """Returns a functions that can be used to apply spectral norm regularization. + + Small spectral norms enforce a small Lipschitz constant, which is necessary + for Wasserstein GANs. + + 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. + scope: An optional scope name. + + Returns: + A function with the signature `sn(weights)` that applies spectral norm + regularization. + + Raises: + ValueError: If scale is negative or if scale is not a float. + """ + if isinstance(scale, numbers.Integral): + raise ValueError('scale cannot be an integer: %s' % scale) + if isinstance(scale, numbers.Real): + if scale < 0.0: + raise ValueError( + 'Setting a scale less than 0 on a regularizer: %g' % scale) + if scale == 0.0: + logging.info('Scale of 0 disables regularizer.') + return lambda _: None + + def sn(weights, name=None): + """Applies spectral norm regularization to weights.""" + with ops.name_scope(scope, 'SpectralNormRegularizer', [weights]) as name: + scale_t = ops.convert_to_tensor( + scale, dtype=weights.dtype.base_dtype, name='scale') + return math_ops.multiply( + scale_t, + compute_spectral_norm( + weights, power_iteration_rounds=power_iteration_rounds), + name=name) + + return sn + + +def _default_name_filter(name): + """A filter function to identify common names of weight variables. + + Args: + name: The variable name. + + Returns: + Whether `name` is a standard name for a weight/kernel variables used in the + Keras, tf.layers, tf.contrib.layers or tf.contrib.slim libraries. + """ + match = re.match(r'(.*\/)?(depthwise_|pointwise_)?(weights|kernel)$', name) + return match is not None + + +def spectral_normalization_custom_getter(name_filter=_default_name_filter, + power_iteration_rounds=1): + """Custom getter that performs Spectral Normalization on a weight tensor. + + Specifically it divides the weight tensor by its largest singular value. This + is intended to stabilize GAN training, by making the discriminator satisfy a + local 1-Lipschitz constraint. + + Based on [Spectral Normalization for Generative Adversarial Networks][sn-gan]. + + [sn-gan]: https://openreview.net/forum?id=B1QRgziT- + + To reproduce an SN-GAN, apply this custom_getter to every weight tensor of + your discriminator. The last dimension of the weight tensor must be the number + of output channels. + + Apply this to layers by supplying this as the `custom_getter` of a + `tf.variable_scope`. For example: + + with tf.variable_scope('discriminator', + custom_getter=spectral_norm_getter()): + net = discriminator_fn(net) + + IMPORTANT: Keras does not respect the custom_getter supplied by the + VariableScope, so Keras users should use `keras_spectral_normalization` + instead of (or in addition to) this approach. + + It is important to carefully select to which weights you want to apply + Spectral Normalization. In general you want to normalize the kernels of + convolution and dense layers, but you do not want to normalize biases. You + also want to avoid normalizing batch normalization (and similar) variables, + but in general such layers play poorly with Spectral Normalization, since the + gamma can cancel out the normalization in other layers. By default we supply a + filter that matches the kernel variable names of the dense and convolution + layers of the tf.layers, tf.contrib.layers, tf.keras and tf.contrib.slim + libraries. If you are using anything else you'll need a custom `name_filter`. + + This custom getter internally creates a variable used to compute the spectral + norm by power iteration. It will update every time the variable is accessed, + which means the normalized discriminator weights may change slightly whilst + training the generator. Whilst unusual, this matches how the paper's authors + implement it, and in general additional rounds of power iteration can't hurt. + + Args: + 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 + true spectral norm. + + Returns: + A custom getter function that applies Spectral Normalization to all + Variables whose names match `name_filter`. + + Raises: + ValueError: If name_filter is not callable. + """ + if not callable(name_filter): + raise ValueError('name_filter must be callable') + + def _internal_getter(getter, name, *args, **kwargs): + """A custom getter function that applies Spectral Normalization. + + Args: + getter: The true getter to call. + name: Name of new/existing variable, in the same format as + tf.get_variable. + *args: Other positional arguments, in the same format as tf.get_variable. + **kwargs: Keyword arguments, in the same format as tf.get_variable. + + Returns: + The return value of `getter(name, *args, **kwargs)`, spectrally + normalized. + + Raises: + ValueError: If used incorrectly, or if `dtype` is not supported. + """ + if not name_filter(name): + return getter(name, *args, **kwargs) + + if name.endswith(_PERSISTED_U_VARIABLE_SUFFIX): + raise ValueError( + 'Cannot apply Spectral Normalization to internal variables created ' + 'for Spectral Normalization. Tried to normalized variable [%s]' % + name) + + if kwargs['dtype'] not in _OK_DTYPES_FOR_SPECTRAL_NORM: + raise ValueError('Disallowed data type {}'.format(kwargs['dtype'])) + + # This layer's weight Variable/PartitionedVariable. + w_tensor = getter(name, *args, **kwargs) + + if len(w_tensor.get_shape()) < 2: + raise ValueError( + 'Spectral norm can only be applied to multi-dimensional tensors') + + return spectral_normalize( + w_tensor, + power_iteration_rounds=power_iteration_rounds, + name=(name + '/spectral_normalize')) + + return _internal_getter + + +@contextlib.contextmanager +def keras_spectral_normalization(name_filter=_default_name_filter, + power_iteration_rounds=1): + """A context manager that enables Spectral Normalization for Keras. + + Keras doesn't respect the `custom_getter` in the VariableScope, so this is a + bit of a hack to make things work. + + Usage: + with keras_spectral_normalization(): + net = discriminator_fn(net) + + Args: + 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 + true spectral norm. + + Yields: + A context manager that wraps the standard Keras variable creation method + with the `spectral_normalization_custom_getter`. + """ + original_make_variable = keras_base_layer_utils.make_variable + sn_getter = spectral_normalization_custom_getter( + name_filter=name_filter, power_iteration_rounds=power_iteration_rounds) + + def make_variable_wrapper(name, *args, **kwargs): + return sn_getter(original_make_variable, name, *args, **kwargs) + + keras_base_layer_utils.make_variable = make_variable_wrapper + + yield + + keras_base_layer_utils.make_variable = original_make_variable diff --git a/tensorflow/contrib/gan/python/features/python/spectral_normalization_test.py b/tensorflow/contrib/gan/python/features/python/spectral_normalization_test.py new file mode 100644 index 0000000000000000000000000000000000000000..4ea21f70ec01950cfef5e4fa851c78b219d6062f --- /dev/null +++ b/tensorflow/contrib/gan/python/features/python/spectral_normalization_test.py @@ -0,0 +1,354 @@ +# 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. +# ============================================================================== +"""Tests for features.spectral_normalization.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib import slim +from tensorflow.contrib.gan.python.features.python import spectral_normalization_impl as spectral_normalization +from tensorflow.contrib.layers.python.layers import layers as contrib_layers +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.keras.layers import convolutional as keras_convolutional +from tensorflow.python.keras.layers import core as keras_core +from tensorflow.python.layers import convolutional as layers_convolutional +from tensorflow.python.layers import core as layers_core +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import init_ops +from tensorflow.python.ops import linalg_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import variable_scope +from tensorflow.python.ops import variables +from tensorflow.python.platform import test + + +class SpectralNormalizationTest(test.TestCase): + + def testComputeSpectralNorm(self): + weights = variable_scope.get_variable( + 'w', dtype=dtypes.float32, shape=[2, 3, 50, 100]) + weights = math_ops.multiply(weights, 10.0) + s = linalg_ops.svd( + array_ops.reshape(weights, [-1, weights.shape[-1]]), compute_uv=False) + true_sn = s[..., 0] + estimated_sn = spectral_normalization.compute_spectral_norm(weights) + + with self.cached_session() as sess: + sess.run(variables.global_variables_initializer()) + np_true_sn = sess.run(true_sn) + for i in range(50): + est = sess.run(estimated_sn) + if i < 1: + np_est_1 = est + if i < 4: + np_est_5 = est + if i < 9: + np_est_10 = est + np_est_50 = est + + # Check that the estimate improves with more iterations. + self.assertAlmostEqual(np_true_sn, np_est_50, 0) + self.assertGreater( + abs(np_true_sn - np_est_10), abs(np_true_sn - np_est_50)) + self.assertGreater( + abs(np_true_sn - np_est_5), abs(np_true_sn - np_est_10)) + self.assertGreater(abs(np_true_sn - np_est_1), abs(np_true_sn - np_est_5)) + + def testSpectralNormalize(self): + weights = variable_scope.get_variable( + 'w', dtype=dtypes.float32, shape=[2, 3, 50, 100]) + weights = math_ops.multiply(weights, 10.0) + normalized_weights = spectral_normalization.spectral_normalize( + weights, power_iteration_rounds=1) + + unnormalized_sigma = linalg_ops.svd( + array_ops.reshape(weights, [-1, weights.shape[-1]]), + compute_uv=False)[..., 0] + normalized_sigma = linalg_ops.svd( + array_ops.reshape(normalized_weights, [-1, weights.shape[-1]]), + compute_uv=False)[..., 0] + + with self.cached_session() as sess: + sess.run(variables.global_variables_initializer()) + s0 = sess.run(unnormalized_sigma) + + for i in range(50): + sigma = sess.run(normalized_sigma) + if i < 1: + s1 = sigma + if i < 5: + s5 = sigma + if i < 10: + s10 = sigma + s50 = sigma + + self.assertAlmostEqual(1., s50, 0) + self.assertGreater(abs(s10 - 1.), abs(s50 - 1.)) + self.assertGreater(abs(s5 - 1.), abs(s10 - 1.)) + self.assertGreater(abs(s1 - 1.), abs(s5 - 1.)) + self.assertGreater(abs(s0 - 1.), abs(s1 - 1.)) + + def _testLayerHelper(self, build_layer_fn, w_shape, b_shape, is_keras=False): + x = array_ops.placeholder(dtypes.float32, shape=[2, 10, 10, 3]) + + w_initial = np.random.randn(*w_shape) * 10 + w_initializer = init_ops.constant_initializer(w_initial) + b_initial = np.random.randn(*b_shape) + b_initializer = init_ops.constant_initializer(b_initial) + + if is_keras: + context_manager = spectral_normalization.keras_spectral_normalization() + else: + getter = spectral_normalization.spectral_normalization_custom_getter() + context_manager = variable_scope.variable_scope('', custom_getter=getter) + + with context_manager: + (net, + expected_normalized_vars, expected_not_normalized_vars) = build_layer_fn( + x, w_initializer, b_initializer) + + x_data = np.random.rand(*x.shape) + + with self.cached_session() as sess: + sess.run(variables.global_variables_initializer()) + + # Before running a forward pass we still expect the variables values to + # differ from the initial value because of the normalizer. + w_befores = [] + for name, var in expected_normalized_vars.items(): + w_before = sess.run(var) + w_befores.append(w_before) + self.assertFalse( + np.allclose(w_initial, w_before), + msg=('%s appears not to be normalized. Before: %s After: %s' % + (name, w_initial, w_before))) + + # Not true for the unnormalized variables. + for name, var in expected_not_normalized_vars.items(): + b_before = sess.run(var) + self.assertTrue( + np.allclose(b_initial, b_before), + msg=('%s appears to be unexpectedly normalized. ' + 'Before: %s After: %s' % (name, b_initial, b_before))) + + # Run a bunch of forward passes. + for _ in range(1000): + _ = sess.run(net, feed_dict={x: x_data}) + + # We expect this to have improved the estimate of the spectral norm, + # which should have changed the variable values and brought them close + # to the true Spectral Normalized values. + _, s, _ = np.linalg.svd(w_initial.reshape([-1, 3])) + exactly_normalized = w_initial / s[0] + for w_before, (name, var) in zip(w_befores, + expected_normalized_vars.items()): + w_after = sess.run(var) + self.assertFalse( + np.allclose(w_before, w_after, rtol=1e-8, atol=1e-8), + msg=('%s did not improve over many iterations. ' + 'Before: %s After: %s' % (name, w_before, w_after))) + self.assertAllClose( + exactly_normalized, + w_after, + rtol=1e-4, + atol=1e-4, + msg=('Estimate of spectral norm for %s was innacurate. ' + 'Normalized matrices do not match.' + 'Estimate: %s Actual: %s' % (name, w_after, + exactly_normalized))) + + def testConv2D_Layers(self): + + def build_layer_fn(x, w_initializer, b_initializer): + layer = layers_convolutional.Conv2D( + filters=3, + kernel_size=3, + padding='same', + kernel_initializer=w_initializer, + bias_initializer=b_initializer) + net = layer.apply(x) + expected_normalized_vars = {'tf.layers.Conv2d.kernel': layer.kernel} + expected_not_normalized_vars = {'tf.layers.Conv2d.bias': layer.bias} + + return net, expected_normalized_vars, expected_not_normalized_vars + + self._testLayerHelper(build_layer_fn, (3, 3, 3, 3), (3,)) + + def testConv2D_ContribLayers(self): + + def build_layer_fn(x, w_initializer, b_initializer): + var_collection = { + 'weights': ['CONTRIB_LAYERS_CONV2D_WEIGHTS'], + 'biases': ['CONTRIB_LAYERS_CONV2D_BIASES'] + } + net = contrib_layers.conv2d( + x, + 3, + 3, + weights_initializer=w_initializer, + biases_initializer=b_initializer, + variables_collections=var_collection) + weight_vars = ops.get_collection('CONTRIB_LAYERS_CONV2D_WEIGHTS') + self.assertEquals(1, len(weight_vars)) + bias_vars = ops.get_collection('CONTRIB_LAYERS_CONV2D_BIASES') + self.assertEquals(1, len(bias_vars)) + expected_normalized_vars = { + 'contrib.layers.conv2d.weights': weight_vars[0] + } + expected_not_normalized_vars = { + 'contrib.layers.conv2d.bias': bias_vars[0] + } + + return net, expected_normalized_vars, expected_not_normalized_vars + + self._testLayerHelper(build_layer_fn, (3, 3, 3, 3), (3,)) + + def testConv2D_Slim(self): + + def build_layer_fn(x, w_initializer, b_initializer): + var_collection = { + 'weights': ['SLIM_CONV2D_WEIGHTS'], + 'biases': ['SLIM_CONV2D_BIASES'] + } + net = slim.conv2d( + x, + 3, + 3, + weights_initializer=w_initializer, + biases_initializer=b_initializer, + variables_collections=var_collection) + weight_vars = ops.get_collection('SLIM_CONV2D_WEIGHTS') + self.assertEquals(1, len(weight_vars)) + bias_vars = ops.get_collection('SLIM_CONV2D_BIASES') + self.assertEquals(1, len(bias_vars)) + expected_normalized_vars = {'slim.conv2d.weights': weight_vars[0]} + expected_not_normalized_vars = {'slim.conv2d.bias': bias_vars[0]} + + return net, expected_normalized_vars, expected_not_normalized_vars + + self._testLayerHelper(build_layer_fn, (3, 3, 3, 3), (3,)) + + def testConv2D_Keras(self): + + def build_layer_fn(x, w_initializer, b_initializer): + layer = keras_convolutional.Conv2D( + filters=3, + kernel_size=3, + padding='same', + kernel_initializer=w_initializer, + bias_initializer=b_initializer) + net = layer.apply(x) + expected_normalized_vars = {'keras.layers.Conv2d.kernel': layer.kernel} + expected_not_normalized_vars = {'keras.layers.Conv2d.bias': layer.bias} + + return net, expected_normalized_vars, expected_not_normalized_vars + + self._testLayerHelper(build_layer_fn, (3, 3, 3, 3), (3,), is_keras=True) + + def testFC_Layers(self): + + def build_layer_fn(x, w_initializer, b_initializer): + x = layers_core.Flatten()(x) + layer = layers_core.Dense( + units=3, + kernel_initializer=w_initializer, + bias_initializer=b_initializer) + net = layer.apply(x) + expected_normalized_vars = {'tf.layers.Dense.kernel': layer.kernel} + expected_not_normalized_vars = {'tf.layers.Dense.bias': layer.bias} + + return net, expected_normalized_vars, expected_not_normalized_vars + + self._testLayerHelper(build_layer_fn, (300, 3), (3,)) + + def testFC_ContribLayers(self): + + def build_layer_fn(x, w_initializer, b_initializer): + var_collection = { + 'weights': ['CONTRIB_LAYERS_FC_WEIGHTS'], + 'biases': ['CONTRIB_LAYERS_FC_BIASES'] + } + x = contrib_layers.flatten(x) + net = contrib_layers.fully_connected( + x, + 3, + weights_initializer=w_initializer, + biases_initializer=b_initializer, + variables_collections=var_collection) + weight_vars = ops.get_collection('CONTRIB_LAYERS_FC_WEIGHTS') + self.assertEquals(1, len(weight_vars)) + bias_vars = ops.get_collection('CONTRIB_LAYERS_FC_BIASES') + self.assertEquals(1, len(bias_vars)) + expected_normalized_vars = { + 'contrib.layers.fully_connected.weights': weight_vars[0] + } + expected_not_normalized_vars = { + 'contrib.layers.fully_connected.bias': bias_vars[0] + } + + return net, expected_normalized_vars, expected_not_normalized_vars + + self._testLayerHelper(build_layer_fn, (300, 3), (3,)) + + def testFC_Slim(self): + + def build_layer_fn(x, w_initializer, b_initializer): + var_collection = { + 'weights': ['SLIM_FC_WEIGHTS'], + 'biases': ['SLIM_FC_BIASES'] + } + x = slim.flatten(x) + net = slim.fully_connected( + x, + 3, + weights_initializer=w_initializer, + biases_initializer=b_initializer, + variables_collections=var_collection) + weight_vars = ops.get_collection('SLIM_FC_WEIGHTS') + self.assertEquals(1, len(weight_vars)) + bias_vars = ops.get_collection('SLIM_FC_BIASES') + self.assertEquals(1, len(bias_vars)) + expected_normalized_vars = { + 'slim.fully_connected.weights': weight_vars[0] + } + expected_not_normalized_vars = {'slim.fully_connected.bias': bias_vars[0]} + + return net, expected_normalized_vars, expected_not_normalized_vars + + self._testLayerHelper(build_layer_fn, (300, 3), (3,)) + + def testFC_Keras(self): + + def build_layer_fn(x, w_initializer, b_initializer): + x = keras_core.Flatten()(x) + layer = keras_core.Dense( + units=3, + kernel_initializer=w_initializer, + bias_initializer=b_initializer) + net = layer.apply(x) + expected_normalized_vars = {'keras.layers.Dense.kernel': layer.kernel} + expected_not_normalized_vars = {'keras.layers.Dense.bias': layer.bias} + + return net, expected_normalized_vars, expected_not_normalized_vars + + self._testLayerHelper(build_layer_fn, (300, 3), (3,), is_keras=True) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/gan/python/losses/python/losses_impl_test.py b/tensorflow/contrib/gan/python/losses/python/losses_impl_test.py index e3c780ac1a0f0ef15ff993bd3a9bf9730dcb45b8..44ee0f52696dc1cdcd91286a80b2d4b42be93a4d 100644 --- a/tensorflow/contrib/gan/python/losses/python/losses_impl_test.py +++ b/tensorflow/contrib/gan/python/losses/python/losses_impl_test.py @@ -403,7 +403,9 @@ class _PenaltyTest(object): def test_all_correct(self): loss = self._penalty_fn(**self._kwargs) self.assertEqual(self._expected_dtype, loss.dtype) - self.assertEqual(self._expected_op_name, loss.op.name) + # NOTE: Op names will change, it is inappropriate to include them in tests. + # See go/tf-breaking-change. + # self.assertEqual(self._expected_op_name, loss.op.name) with self.cached_session(): variables.global_variables_initializer().run() self.assertAlmostEqual(self._expected_loss, loss.eval(), 6) diff --git a/tensorflow/contrib/gdr/BUILD b/tensorflow/contrib/gdr/BUILD index 704be917b3680a1b5712f4f1dc5059b354db8610..bf8b66dcfa5e44a03107cdf1ef8b04e1dbff4a9c 100644 --- a/tensorflow/contrib/gdr/BUILD +++ b/tensorflow/contrib/gdr/BUILD @@ -17,11 +17,6 @@ filegroup( ]), ) -load( - "//tensorflow:tensorflow.bzl", - "tf_cuda_library", -) - # For platform specific build config load( "//tensorflow/core:platform/default/build_config.bzl", @@ -66,7 +61,6 @@ cc_library( ":gdr_memory_manager", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", - "//tensorflow/core:gpu_runtime", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core/distributed_runtime:graph_mgr", @@ -100,15 +94,37 @@ cc_library( ], ) +cc_library( + name = "gdr_collective_executor_mgr", + srcs = ["gdr_collective_executor_mgr.cc"], + hdrs = ["gdr_collective_executor_mgr.h"], + deps = [ + ":gdr_memory_manager", + "//tensorflow/core:core_cpu_internal", + "//tensorflow/core:framework", + "//tensorflow/core:lib_internal", + "//tensorflow/core/distributed_runtime:cancellable_call", + "//tensorflow/core/distributed_runtime:collective_param_resolver_distributed", + "//tensorflow/core/distributed_runtime:device_resolver_distributed", + "//tensorflow/core/distributed_runtime:request_id", + "//tensorflow/core/distributed_runtime:rpc_collective_executor_mgr", + "//tensorflow/core/distributed_runtime:worker_cache", + ], +) + cc_library( name = "gdr_server_lib", srcs = ["gdr_server_lib.cc"], hdrs = ["gdr_server_lib.h"], linkstatic = 1, # Seems to be needed since alwayslink is broken in bazel deps = [ + ":gdr_collective_executor_mgr", ":gdr_memory_manager", ":gdr_rendezvous_mgr", ":gdr_worker", + "//tensorflow/core:core_cpu_internal", + "//tensorflow/core/distributed_runtime:collective_param_resolver_distributed", + "//tensorflow/core/distributed_runtime:device_resolver_distributed", "//tensorflow/core/distributed_runtime/rpc:grpc_server_lib", ], alwayslink = 1, diff --git a/tensorflow/contrib/gdr/README.md b/tensorflow/contrib/gdr/README.md index 8242d93f129904828a11b61d48f2df8fb0f88bc3..711adc865f37fc84550e4b45d9f0c7fff421a0dc 100644 --- a/tensorflow/contrib/gdr/README.md +++ b/tensorflow/contrib/gdr/README.md @@ -114,7 +114,16 @@ Caveats In current implementation, only tensors that reside in host memory or in GPU memory such that the GPU is adjacent to an RDMA capable NIC will use direct RDMA as its transport. When RDMA is available but not GDR, a temporary tensor copy on host memory will be used as RDMA source/destination (and copied from/to the target device). When there is no RDMA device present, it can even fallback to the original gRPC runtime. While it is theoretically possible to mix GDR enabled TF with non-GDR deployments in the same job, make sure the environment is properly setup so the GDR mode is enabled whenever possible (i.e. do not fall back to gRPC when it is not absolutely necessary). -In the original design (as in the reference), tensor buffers are only registered to NIC when we could determine that the tensor will be either a source of Send or a sink of Recv across physical machine boundary. However, to implement the precise allocations, we need to change all the devices to possibly return a NIC compatible allocator. As GDR is currently in contrib, we would like to avoid the unnecessary code disruption to the TF core, so we allocate all tensors from NIC-registered buffers using a BFC allocator. This behaviour is similar to the effect of enabling the extra GPU option `force_gpu_compatible`, which allocate all host tensors in GPU-registered buffers no matter they will be transferred from/to GPUs or not. +In the original design (as in the reference), tensor buffers are only registered +to NIC when we could determine that the tensor will be either a source of Send +or a sink of Recv across physical machine boundary. However, to implement the +precise allocations, we need to change all the devices to possibly return a NIC +compatible allocator. As GDR is currently in contrib, we would like to avoid the +unnecessary code disruption to the TF core, so we allocate all tensors from +NIC-registered buffers using a BFC allocator. This behavior is similar to the +effect of enabling the extra GPU option `force_gpu_compatible`, which allocate +all host tensors in GPU-registered buffers no matter they will be transferred +from/to GPUs or not. Reference === diff --git a/tensorflow/contrib/gdr/gdr_collective_executor_mgr.cc b/tensorflow/contrib/gdr/gdr_collective_executor_mgr.cc new file mode 100644 index 0000000000000000000000000000000000000000..755cbdff31cd7ca31579e0d64399d681dc24ad81 --- /dev/null +++ b/tensorflow/contrib/gdr/gdr_collective_executor_mgr.cc @@ -0,0 +1,159 @@ +/* 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/contrib/gdr/gdr_collective_executor_mgr.h" + +#include "tensorflow/core/common_runtime/base_collective_executor.h" +#include "tensorflow/core/common_runtime/collective_executor_mgr.h" +#include "tensorflow/core/common_runtime/collective_rma_local.h" +#include "tensorflow/core/common_runtime/dma_helper.h" +#include "tensorflow/core/distributed_runtime/cancellable_call.h" +#include "tensorflow/core/distributed_runtime/collective_param_resolver_distributed.h" +#include "tensorflow/core/distributed_runtime/device_resolver_distributed.h" +#include "tensorflow/core/distributed_runtime/request_id.h" +#include "tensorflow/core/distributed_runtime/worker_cache.h" +#include "tensorflow/core/lib/random/random.h" + +namespace tensorflow { + +class WorkerCacheInterface; + +namespace { + +class RecvBufCall : public CancellableCall { + public: + RecvBufCall(int64 step_id, const string& peer_device, const string& peer_task, + const string& key, Device* to_device, + DeviceContext* to_device_ctx, + const AllocatorAttributes& to_alloc_attr, Tensor* to_tensor, + const DeviceLocality& client_locality, + const DeviceLocality& server_locality, + CancellationManager* cancel_mgr, WorkerCacheInterface* wc) + : CancellableCall(cancel_mgr, peer_task, wc) { + req_.set_step_id(step_id); + req_.set_buf_rendezvous_key(key); + *req_.mutable_client_locality() = client_locality; + *req_.mutable_server_locality() = server_locality; + req_.set_num_bytes(to_tensor->TotalBytes()); + req_.set_buf_ptr(reinterpret_cast(DMAHelper::base(to_tensor))); + req_.set_src_device(peer_device); + req_.set_dst_device(to_device->name()); + req_.set_request_id(GetUniqueRequestId()); + } + + ~RecvBufCall() override {} + + void IssueCall(const StatusCallback& done) override { + wi_->RecvBufAsync(&opts_, &req_, &resp_, done); + } + + RecvBufRequest req_; + RecvBufResponse resp_; +}; + +class CollectiveRemoteAccessDistributed : public CollectiveRemoteAccessLocal { + public: + CollectiveRemoteAccessDistributed(const DeviceMgr* dev_mgr, + DeviceResolverInterface* dev_resolver, + WorkerCacheInterface* worker_cache, + int64 step_id, + RemoteMemoryManager* remote_memory_manager) + : CollectiveRemoteAccessLocal(dev_mgr, dev_resolver, step_id), + worker_cache_(worker_cache), + remote_memory_manager_(remote_memory_manager) {} + + ~CollectiveRemoteAccessDistributed() override {} + + void RecvFromPeer(const string& peer_device, const string& peer_task, + bool peer_is_local, const string& key, Device* to_device, + DeviceContext* to_device_ctx, + const AllocatorAttributes& to_alloc_attr, Tensor* to_tensor, + const DeviceLocality& client_locality, + int dev_to_dev_stream_index, + const StatusCallback& done) override { + if (peer_is_local) { + CollectiveRemoteAccessLocal::RecvFromPeer( + peer_device, peer_task, peer_is_local, key, to_device, to_device_ctx, + to_alloc_attr, to_tensor, client_locality, dev_to_dev_stream_index, + done); + return; + } + + // State that needs to be threaded through a couple of async calls + // in order to make this function completely non-blocking. + struct State { + DeviceLocality server_locality; + std::unique_ptr call; + }; + State* state = new State; + + // Logic to be executed on the RecvBufAsync callback. + auto recv_buf_callback = [this, state, peer_task, to_device, to_alloc_attr, + to_device_ctx, to_tensor, done](const Status& s) { + if (s.ok()) { + remote_memory_manager_->TensorFromTransportOptions( + to_tensor, state->call->resp_.transport_options(), to_device, + to_device_ctx, to_alloc_attr.on_host(), done); + } + if (!s.ok() && errors::IsFailedPrecondition(s)) { + dev_resolver_->ClearTask(peer_task); + } + + delete state; + }; + + // Logic to execute once we have the device locality for the server-side + // device. + auto dev_locality_callback = [this, state, peer_device, peer_task, key, + to_device, to_device_ctx, to_alloc_attr, + to_tensor, client_locality, + recv_buf_callback](const Status& s) { + if (!s.ok()) { + recv_buf_callback(s); + } else { + state->call.reset(new RecvBufCall( + step_id_, peer_device, peer_task, key, to_device, to_device_ctx, + to_alloc_attr, to_tensor, client_locality, state->server_locality, + &cancel_mgr_, worker_cache_)); + state->call->Start(recv_buf_callback); + } + }; + + dev_resolver_->GetLocalityAsync( + peer_device, peer_task, &state->server_locality, dev_locality_callback); + } + + void StartAbort(const Status& s) override { + CollectiveRemoteAccessLocal::StartAbort(s); + cancel_mgr_.StartCancel(); + } + + protected: + WorkerCacheInterface* worker_cache_; // Not owned + CancellationManager cancel_mgr_; + RemoteMemoryManager* remote_memory_manager_; +}; + +} // namespace + +CollectiveExecutor* GdrCollectiveExecutorMgr::Create(int64 step_id) { + CollectiveRemoteAccessDistributed* rma = + new CollectiveRemoteAccessDistributed(dev_mgr_, dev_resolver_.get(), + worker_cache_, step_id, + remote_memory_manager_); + return new BaseCollectiveExecutor(this, rma, step_id, dev_mgr_, + &gpu_ring_order_); +} + +} // namespace tensorflow diff --git a/tensorflow/contrib/gdr/gdr_collective_executor_mgr.h b/tensorflow/contrib/gdr/gdr_collective_executor_mgr.h new file mode 100644 index 0000000000000000000000000000000000000000..1417e51e82c31035f058e8e9b546e04fb0ad97b8 --- /dev/null +++ b/tensorflow/contrib/gdr/gdr_collective_executor_mgr.h @@ -0,0 +1,56 @@ +/* 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_CONTRIB_GDR_GDR_COLLECTIVE_EXECUTOR_MGR_H_ +#define TENSORFLOW_CONTRIB_GDR_GDR_COLLECTIVE_EXECUTOR_MGR_H_ + +#include "tensorflow/contrib/gdr/gdr_memory_manager.h" +#include "tensorflow/core/distributed_runtime/collective_param_resolver_distributed.h" +#include "tensorflow/core/distributed_runtime/device_resolver_distributed.h" +#include "tensorflow/core/distributed_runtime/rpc_collective_executor_mgr.h" +#include "tensorflow/core/framework/collective.h" + +namespace tensorflow { +class ConfigProto; +class DeviceMgr; +class WorkerCacheInterface; +class StepSequenceRequest; +class StepSequenceResponse; + +// An implementation of CollectiveExecutorMgr for a distributed environment +// that uses WorkerInterface::RecvBufAsync to route data transfers over RDMA. +class GdrCollectiveExecutorMgr : public RpcCollectiveExecutorMgr { + public: + GdrCollectiveExecutorMgr( + const ConfigProto& config, const DeviceMgr* dev_mgr, + std::unique_ptr dev_resolver, + std::unique_ptr param_resolver, + WorkerCacheInterface* worker_cache, const string& task_name, + RemoteMemoryManager* remote_memory_manager) + : RpcCollectiveExecutorMgr(config, dev_mgr, std::move(dev_resolver), + std::move(param_resolver), worker_cache, + task_name), + remote_memory_manager_(remote_memory_manager) {} + + ~GdrCollectiveExecutorMgr() override {} + + protected: + virtual CollectiveExecutor* Create(int64 step_id) override; + + private: + RemoteMemoryManager* remote_memory_manager_; // Not owned. +}; + +} // namespace tensorflow +#endif // TENSORFLOW_CONTRIB_GDR_GDR_COLLECTIVE_EXECUTOR_MGR_H_ diff --git a/tensorflow/contrib/gdr/gdr_memory_manager.cc b/tensorflow/contrib/gdr/gdr_memory_manager.cc index ce1875151597f926aeb6392e7fc8307312da123f..7321e973191c4cc45f88735c6be7f2f67fe71c39 100644 --- a/tensorflow/contrib/gdr/gdr_memory_manager.cc +++ b/tensorflow/contrib/gdr/gdr_memory_manager.cc @@ -73,7 +73,10 @@ int TryToReadNumaNode(ibv_device* device) { std::ifstream ifs(filename.c_str()); string content; - CHECK(std::getline(ifs, content)); + const auto& ret = std::getline(ifs, content); + if (!ret) { + return port::kNUMANoAffinity; + } int32 value; if (strings::safe_strto32(content, &value)) { diff --git a/tensorflow/contrib/gdr/gdr_server_lib.cc b/tensorflow/contrib/gdr/gdr_server_lib.cc index dc0d5d548b80d36409778ef34e63171441f10142..c39cc0f9bcecc26aedfaf9707113210acf670244 100644 --- a/tensorflow/contrib/gdr/gdr_server_lib.cc +++ b/tensorflow/contrib/gdr/gdr_server_lib.cc @@ -16,11 +16,13 @@ limitations under the License. #include "tensorflow/contrib/gdr/gdr_server_lib.h" #include "grpc/support/alloc.h" +#include "tensorflow/contrib/gdr/gdr_collective_executor_mgr.h" #include "tensorflow/contrib/gdr/gdr_memory_manager.h" #include "tensorflow/contrib/gdr/gdr_rendezvous_mgr.h" #include "tensorflow/contrib/gdr/gdr_worker.h" - -#include "grpc/support/alloc.h" +#include "tensorflow/core/common_runtime/collective_rma_local.h" +#include "tensorflow/core/distributed_runtime/collective_param_resolver_distributed.h" +#include "tensorflow/core/distributed_runtime/device_resolver_distributed.h" namespace tensorflow { @@ -57,10 +59,34 @@ Status GdrServer::Init() { return std::unique_ptr( new GdrWorker(env, config, remote_memory_manager_.get())); }; - + CollectiveMgrCreationFunction collective_mgr_func = + [this](const ConfigProto& config, const WorkerEnv* env, + WorkerCacheInterface* worker_cache) { + string unused; + string default_worker_name; + DeviceNameUtils::SplitDeviceName( + env->device_mgr->ListDevices()[0]->name(), &default_worker_name, + &unused); + + std::unique_ptr dev_resolver( + new DeviceResolverDistributed(env->device_mgr, worker_cache, + default_worker_name)); + std::unique_ptr param_resolver( + new CollectiveParamResolverDistributed( + config, env->device_mgr, dev_resolver.get(), worker_cache, + default_worker_name)); + return new GdrCollectiveExecutorMgr( + config, env->device_mgr, std::move(dev_resolver), + std::move(param_resolver), worker_cache, default_worker_name, + remote_memory_manager_.get()); + }; TF_RETURN_IF_ERROR(remote_memory_manager_->Init()); - return GrpcServer::Init(nullptr, rendezvous_mgr_func, nullptr, worker_func); + GrpcServerOptions opts; + opts.rendezvous_mgr_func = rendezvous_mgr_func; + opts.collective_mgr_func = collective_mgr_func; + opts.worker_func = worker_func; + return GrpcServer::Init(opts); } Status GdrServer::Start() { diff --git a/tensorflow/contrib/gdr/gdr_worker.cc b/tensorflow/contrib/gdr/gdr_worker.cc index 016e5ea27b397830c69b6e1761b5994ebcfa9c3d..1204b8ca501a8f99ea6abd6c047ab2d91350bae1 100644 --- a/tensorflow/contrib/gdr/gdr_worker.cc +++ b/tensorflow/contrib/gdr/gdr_worker.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/contrib/gdr/gdr_worker.h" +#include "tensorflow/core/common_runtime/buf_rendezvous.h" #include "tensorflow/core/common_runtime/device.h" #include "tensorflow/core/common_runtime/device_mgr.h" #include "tensorflow/core/common_runtime/dma_helper.h" @@ -29,6 +30,7 @@ limitations under the License. #include "tensorflow/core/distributed_runtime/worker_cache.h" #include "tensorflow/core/distributed_runtime/worker_session.h" #include "tensorflow/core/framework/cancellation.h" +#include "tensorflow/core/framework/collective.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/platform/logging.h" @@ -40,13 +42,13 @@ GdrWorker::GdrWorker(WorkerEnv* worker_env, const ConfigProto& config, RemoteMemoryManager* remote_memory_manager) : GrpcWorker(worker_env, config), remote_memory_manager_(remote_memory_manager), - recv_tensor_recent_request_ids_(100000) {} + recent_request_ids_(100000) {} void GdrWorker::GrpcRecvTensorAsync(CallOptions* opts, const RecvTensorRequest* request, ::grpc::ByteBuffer* response, StatusCallback done) { - Status s = recv_tensor_recent_request_ids_.TrackUnique( + Status s = recent_request_ids_.TrackUnique( request->request_id(), "RecvTensor (GdrWorker)", *request); if (!s.ok()) { done(s); @@ -145,4 +147,41 @@ void GdrWorker::GrpcRecvTensorAsync(CallOptions* opts, }); } +void GdrWorker::RecvBufAsync(CallOptions* opts, const RecvBufRequest* request, + RecvBufResponse* response, StatusCallback done) { + // This is an RDMA enabled implementation augmenting grpc. + Status s = recent_request_ids_.TrackUnique(request->request_id(), + "RecvBuf (GdrWorker)", *request); + if (!s.ok()) { + done(s); + return; + } + CollectiveExecutor::Handle ce_handle( + env_->collective_executor_mgr->FindOrCreate(request->step_id()), true); + CollectiveRemoteAccess* rma = ce_handle.get()->remote_access(); + rma->buf_rendezvous()->ConsumeBuf( + request->buf_rendezvous_key(), + [this, request, response, done](const Status& status, + BufRendezvous::Hook* hook) { + Status s = status; + if (s.ok()) { + if (!DMAHelper::CanUseDMA(hook->prod_value)) { + s = errors::Internal("Tensor value for key ", + request->buf_rendezvous_key(), + " is not of a type supported by RecvBuf"); + } + } + if (s.ok()) { + remote_memory_manager_->TransportOptionsFromTensor( + response->mutable_transport_options(), *hook->prod_value, + hook->prod_dev, hook->prod_ctx, hook->prod_attr.on_host(), + [this, response, done, hook](const Status& s) { + response->set_send_start_micros(env_->env->NowMicros()); + done(s); + BufRendezvous::DoneWithHook(hook); + }); + } + }); +} + } // namespace tensorflow diff --git a/tensorflow/contrib/gdr/gdr_worker.h b/tensorflow/contrib/gdr/gdr_worker.h index 39f11e6bde5a1ca7ae91ead02279d22d70af027b..9a85cfd4263ad86f6579eedce95969c2829ff62c 100644 --- a/tensorflow/contrib/gdr/gdr_worker.h +++ b/tensorflow/contrib/gdr/gdr_worker.h @@ -38,9 +38,13 @@ class GdrWorker : public GrpcWorker { ::grpc::ByteBuffer* response, StatusCallback done) override; + virtual void RecvBufAsync(CallOptions* opts, const RecvBufRequest* request, + RecvBufResponse* response, + StatusCallback done) override; + private: RemoteMemoryManager* remote_memory_manager_; // Not owned - RecentRequestIds recv_tensor_recent_request_ids_; + RecentRequestIds recent_request_ids_; }; } // namespace tensorflow diff --git a/tensorflow/contrib/hvx/hvx_ops_support_checker/BUILD b/tensorflow/contrib/hvx/hvx_ops_support_checker/BUILD index 0081fb61770075a2c36e92f65e01126f657edeb4..d319aa7986d81cf9ac2d1dc2e15b053a0aa0c31b 100644 --- a/tensorflow/contrib/hvx/hvx_ops_support_checker/BUILD +++ b/tensorflow/contrib/hvx/hvx_ops_support_checker/BUILD @@ -16,9 +16,22 @@ tf_cc_binary( srcs = ["hvx_ops_support_checker_main.cc"], visibility = ["//visibility:public"], deps = [ + "//tensorflow/core:array_ops_op_lib", + "//tensorflow/core:candidate_sampling_ops_op_lib", + "//tensorflow/core:control_flow_ops_op_lib", "//tensorflow/core:framework_internal", + "//tensorflow/core:functional_ops_op_lib", "//tensorflow/core:lib", + "//tensorflow/core:list_ops_op_lib", + "//tensorflow/core:manip_ops_op_lib", + "//tensorflow/core:math_ops_op_lib", + "//tensorflow/core:nn_ops_op_lib", "//tensorflow/core:protos_all_cc", + "//tensorflow/core:random_ops_op_lib", + "//tensorflow/core:remote_fused_graph_ops_op_lib", + "//tensorflow/core:string_ops_op_lib", + "//tensorflow/core:training_ops_op_lib", + "//tensorflow/core:user_ops_op_lib", "//tensorflow/core/kernels:remote_fused_graph_execute_utils", "//tensorflow/core/kernels/hexagon:graph_transferer", "//tensorflow/tools/graph_transforms:file_utils", diff --git a/tensorflow/contrib/ignite/README.md b/tensorflow/contrib/ignite/README.md index 5a8c650fb927be0c835aaceffc516c048195c7bf..c1f6cac4942436d32f9867d4b5557c6b9e376c69 100644 --- a/tensorflow/contrib/ignite/README.md +++ b/tensorflow/contrib/ignite/README.md @@ -30,7 +30,8 @@ system based on Apache Ignite. ## Features -Ignite Dataset provides features that that you can use in a wide range of cases. The most important and interesting features are described below. +Ignite Dataset provides features that you can use in a wide range of cases. The +most important and interesting features are described below. ### Distributed In-Memory Datasource [Apache Ignite](https://ignite.apache.org/) is a distributed in-memory database, caching, and processing platform that provides fast data access. It allows you to avoid limitations of hard drive and store and operate with as much data as you need in distributed cluster. You can utilize @@ -97,6 +98,7 @@ jdbc:ignite:thin://localhost/> INSERT INTO KITTEN_CACHE VALUES (3, 'LITTLE BALL ```python >>> import tensorflow as tf >>> from tensorflow.contrib.ignite import IgniteDataset +>>> tf.enable_eager_execution() >>> >>> dataset = IgniteDataset(cache_name="IMAGES").map(lambda obj: obj['val']['pixels']) >>> @@ -116,7 +118,15 @@ Using this ability we can calculate gradients on the nodes the data is stored on Apache Ignite uses horizontal partitioning to store data in distributed cluster. When we create Apache Ignite cache (or table in terms of SQL), we can specify the number of partitions the data will be partitioned on. For example, if an Apache Ignite cluster consists of 10 machines and we create cache with 10 partitions, then every machine will maintain approximately one data partition. -Ignite Dataset allows using these two aspects of distributed neural network training (using TensorFlow) and Apache Ignite partitioning. Ignite Dataset is a computation graph operation that can be performed on a remote worker. The remote worker can override Ignite Dataset parameters (such as `host`, `port` or `part`) by setting correstondent environment variables for worker process (such as `IGNITE_DATASET_HOST`, `IGNITE_DATASET_PORT` or `IGNITE_DATASET_PART`). Using this overriding approach, we can assign a specific partition to every worker so that one worker handles one partition and, at the same time, transparently work with single dataset. +Ignite Dataset allows using these two aspects of distributed neural network +training (using TensorFlow) and Apache Ignite partitioning. Ignite Dataset is a +computation graph operation that can be performed on a remote worker. The remote +worker can override Ignite Dataset parameters (such as `host`, `port` or `part`) +by setting correspondent environment variables for worker process (such as +`IGNITE_DATASET_HOST`, `IGNITE_DATASET_PORT` or `IGNITE_DATASET_PART`). Using +this overriding approach, we can assign a specific partition to every worker so +that one worker handles one partition and, at the same time, transparently work +with single dataset. ```python >>> import tensorflow as tf @@ -149,23 +159,31 @@ system called [IGFS](https://ignite.apache.org/features/igfs.html). IGFS delivers a similar functionality to Hadoop HDFS, but only in-memory. In fact, in addition to its own APIs, IGFS implements Hadoop FileSystem API and can be transparently plugged into Hadoop or Spark deployments. This contrib package -contains an integration between IGFS and TensorFlow. The integration is based -on [custom filesystem plugin](https://www.tensorflow.org/extend/add_filesys) -from TensorFlow side and +contains an integration between IGFS and TensorFlow. The integration is based on +[custom filesystem plugin](https://www.tensorflow.org/extend/add_filesys) from +TensorFlow side and [IGFS Native API](https://ignite.apache.org/features/igfs.html) from Apache -Ignite side. It has numerous uses, for example: * Checkpoints of state can be -saved to IGFS for reliability and fault-tolerance. * Training processes -communicate with TensorBoard by writing event files to a directory, which -TensorBoard watches. IGFS allows this communication to work even when -TensorBoard runs in a different process or machine. +Ignite side. It has numerous uses, for example: + +* Checkpoints of state can be saved to IGFS for reliability and + fault-tolerance. +* Training processes communicate with TensorBoard by writing event files to a + directory, which TensorBoard watches. IGFS allows this communication to work + even when TensorBoard runs in a different process or machine. ### SSL Connection -Apache Ignite allows to protect data transfer channels by [SSL](https://en.wikipedia.org/wiki/Transport_Layer_Security) and authentification. Ignite Dataset supports both SSL connection with and without authntication. For more information, please refer to the [Apache Ignite SSL/TLS](https://apacheignite.readme.io/docs/ssltls) documentation. +Apache Ignite allows to protect data transfer channels by +[SSL](https://en.wikipedia.org/wiki/Transport_Layer_Security) and +authentication. Ignite Dataset supports both SSL connection with and without +authentication. For more information, please refer to the +[Apache Ignite SSL/TLS](https://apacheignite.readme.io/docs/ssltls) +documentation. ```python >>> import tensorflow as tf >>> from tensorflow.contrib.ignite import IgniteDataset +>>> tf.enable_eager_execution() >>> >>> dataset = IgniteDataset(cache_name="IMAGES", certfile="client.pem", @@ -186,7 +204,7 @@ Following examples will help you to easily start working with this module. The simplest way to try Ignite Dataset is to run a [Docker](https://www.docker.com/) container with Apache Ignite and loaded -[MNIST](http://yann.lecun.com/exdb/mnist/) data and after start interruct with +[MNIST](http://yann.lecun.com/exdb/mnist/) data and after start interrupt with it using Ignite Dataset. Such container is available on Docker Hub: [dmitrievanthony/ignite-with-mnist](https://hub.docker.com/r/dmitrievanthony/ignite-with-mnist/). You need to start this container on your machine: @@ -197,13 +215,13 @@ docker run -it -p 10800:10800 dmitrievanthony/ignite-with-mnist After that you will be able to work with it following way: -![ignite-dataset-mnist](https://s3.amazonaws.com/helloworld23423423ew23/ignite-dataset-mnist.png "Ignite Dataset Mnist") +![ignite-dataset-mnist](https://s3.amazonaws.com/helloworld23423423ew23/ignite-dataset-mnist-2.png "Ignite Dataset Mnist") ### IGFS The simplest way to try IGFS with TensorFlow is to run [Docker](https://www.docker.com/) container with Apache Ignite and enabled IGFS -and then interruct with it using TensorFlow +and then interrupt with it using TensorFlow [tf.gfile](https://www.tensorflow.org/api_docs/python/tf/gfile). Such container is available on Docker Hub: [dmitrievanthony/ignite-with-igfs](https://hub.docker.com/r/dmitrievanthony/ignite-with-igfs/). diff --git a/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py b/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py index 66e654ca636a5a051c6f9cd35bf9001dfbcbf7f4..3ffceef8070e0fc3b3cebae2522f89fe98ce4413 100644 --- a/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py +++ b/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py @@ -735,8 +735,6 @@ class IgniteDataset(dataset_ops.DatasetSource): cert_password: Password to be used if the private key is encrypted and a password is necessary. """ - super(IgniteDataset, self).__init__() - with IgniteClient(host, port, username, password, certfile, keyfile, cert_password) as client: client.handshake() @@ -760,6 +758,8 @@ class IgniteDataset(dataset_ops.DatasetSource): self.cache_type.to_output_types(), self.cache_type.to_output_shapes(), self.cache_type.to_output_classes()) + super(IgniteDataset, self).__init__(self._as_variant_tensor()) + def _as_variant_tensor(self): return gen_dataset_ops.ignite_dataset(self.cache_name, self.host, self.port, self.local, self.part, self.page_size, diff --git a/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py b/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py index ff5d4c458c859fd8e5e3ae65ee41a454d55d6538..89b74fbfdc38c9f42795d5c778889210baf6387f 100644 --- a/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py +++ b/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py @@ -19,9 +19,9 @@ from __future__ import print_function import os +from tensorflow import compat from tensorflow.contrib.ignite import IgniteDataset from tensorflow.python.client import session -from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.platform import test @@ -66,7 +66,7 @@ class IgniteDatasetTest(test.TestCase): self.assertEqual(dtypes.string, dataset.output_types["val"]["NAME"]) self.assertEqual(dtypes.int64, dataset.output_types["val"]["VAL"]) - it = dataset_ops.make_one_shot_iterator(dataset) + it = compat.v1.data.make_one_shot_iterator(dataset) ne = it.get_next() with session.Session() as sess: diff --git a/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.sh b/tensorflow/contrib/kafka/python/kernel_tests/kafka_test.sh old mode 100644 new mode 100755 diff --git a/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py b/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py index b399e1b6c2ac47db205b5d8bbc81875ef5c08a31..5591c3b0cc8c8bf196bb4821c018cbf155cba4ce 100644 --- a/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py +++ b/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py @@ -52,7 +52,6 @@ class KafkaDataset(dataset_ops.DatasetSource): timeout: The timeout value for the Kafka Consumer to wait (in millisecond). """ - super(KafkaDataset, self).__init__() self._topics = ops.convert_to_tensor( topics, dtype=dtypes.string, name="topics") self._servers = ops.convert_to_tensor( @@ -63,6 +62,8 @@ class KafkaDataset(dataset_ops.DatasetSource): self._timeout = ops.convert_to_tensor( timeout, dtype=dtypes.int64, name="timeout") + super(KafkaDataset, self).__init__(self._as_variant_tensor()) + def _as_variant_tensor(self): return gen_dataset_ops.kafka_dataset(self._topics, self._servers, self._group, self._eof, self._timeout) diff --git a/tensorflow/contrib/kernel_methods/python/losses.py b/tensorflow/contrib/kernel_methods/python/losses.py index 4ef0a66a52429233c6e6f70667a451466493629c..294a7d69a704b3c06ab9e30489af116929ab6c2a 100644 --- a/tensorflow/contrib/kernel_methods/python/losses.py +++ b/tensorflow/contrib/kernel_methods/python/losses.py @@ -34,7 +34,7 @@ def sparse_multiclass_hinge_loss( scope=None, loss_collection=ops.GraphKeys.LOSSES, reduction=losses.Reduction.SUM_BY_NONZERO_WEIGHTS): - """Adds Ops for computing the multiclass hinge loss. + r"""Adds Ops for computing the multiclass hinge loss. The implementation is based on the following paper: On the Algorithmic Implementation of Multiclass Kernel-based Vector Machines diff --git a/tensorflow/contrib/kinesis/python/ops/kinesis_dataset_ops.py b/tensorflow/contrib/kinesis/python/ops/kinesis_dataset_ops.py index 2b1d478a9b0fd12ca25c72da6872acccfd7285fc..9479afb180df7bb4a08d6aafa4fc3bf63489d9f3 100644 --- a/tensorflow/contrib/kinesis/python/ops/kinesis_dataset_ops.py +++ b/tensorflow/contrib/kinesis/python/ops/kinesis_dataset_ops.py @@ -71,7 +71,6 @@ class KinesisDataset(dataset_ops.DatasetSource): interval: The interval for the Kinesis Client to wait before it tries to get records again (in millisecond). """ - super(KinesisDataset, self).__init__() self._stream = ops.convert_to_tensor( stream, dtype=dtypes.string, name="stream") self._shard = ops.convert_to_tensor( @@ -80,6 +79,7 @@ class KinesisDataset(dataset_ops.DatasetSource): read_indefinitely, dtype=dtypes.bool, name="read_indefinitely") self._interval = ops.convert_to_tensor( interval, dtype=dtypes.int64, name="interval") + super(KinesisDataset, self).__init__(self._as_variant_tensor()) def _as_variant_tensor(self): return gen_dataset_ops.kinesis_dataset( diff --git a/tensorflow/contrib/labeled_tensor/BUILD b/tensorflow/contrib/labeled_tensor/BUILD index 588f15b867c1fedbadd5a5d945d870a356549468..7e19ae7c13df421ec5bb9cb0e07dff0d00fb9548 100644 --- a/tensorflow/contrib/labeled_tensor/BUILD +++ b/tensorflow/contrib/labeled_tensor/BUILD @@ -155,7 +155,7 @@ py_library( ":core", "//tensorflow/python:array_ops", "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:functional_ops", + "//tensorflow/python:map_fn", "//tensorflow/python:math_ops", "//tensorflow/python:numerics", "//tensorflow/python:random_ops", diff --git a/tensorflow/contrib/labeled_tensor/python/ops/ops.py b/tensorflow/contrib/labeled_tensor/python/ops/ops.py index 2ede5daee74223e812cc29e9708b1989b698fb4e..a65f045cc886f4d4f351423858d92412baa3a622 100644 --- a/tensorflow/contrib/labeled_tensor/python/ops/ops.py +++ b/tensorflow/contrib/labeled_tensor/python/ops/ops.py @@ -29,6 +29,7 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import functional_ops +from tensorflow.python.ops import map_fn as map_fn_lib from tensorflow.python.ops import math_ops from tensorflow.python.ops import numerics from tensorflow.python.ops import random_ops @@ -629,7 +630,7 @@ def map_fn(fn, labeled_tensor, name=None): # TODO(ericmc): Fix this upstream. if labeled_tensor.dtype == dtypes.string: - # We must construct the full graph here, because functional_ops.map_fn + # We must construct the full graph here, because map_fn_lib.map_fn # doesn't work for string-valued tensors. # Constructing the full graph may be slow. map_lts = [fn(t) for t in unpack_lts] @@ -652,7 +653,7 @@ def map_fn(fn, labeled_tensor, name=None): tensor_lt = core.LabeledTensor(tensor, original_axes) return fn(tensor_lt).tensor - map_op = functional_ops.map_fn( + map_op = map_fn_lib.map_fn( tf_fn, labeled_tensor.tensor, dtype=first_map_lt.dtype) map_lt = core.LabeledTensor(map_op, final_axes) diff --git a/tensorflow/contrib/layers/BUILD b/tensorflow/contrib/layers/BUILD index 9ca6f8df5dbe3c236c4cd85095176ce69ad9deaa..69d5496f8aebb9b89c5d79f80a1a439f556093d7 100644 --- a/tensorflow/contrib/layers/BUILD +++ b/tensorflow/contrib/layers/BUILD @@ -81,6 +81,7 @@ tf_custom_op_py_library( visibility = [ "//learning/brain:__subpackages__", "//tensorflow:__subpackages__", + "//tensorflow_model_optimization:__subpackages__", "//video/youtube/personalization:__subpackages__", ], deps = [ diff --git a/tensorflow/contrib/layers/python/layers/feature_column_ops_test.py b/tensorflow/contrib/layers/python/layers/feature_column_ops_test.py index 7e6eafaa0d6f60cfc28a4c422abac0b6d5a991fb..00e41026d0038409ace178e6affd2c1cdc812122 100644 --- a/tensorflow/contrib/layers/python/layers/feature_column_ops_test.py +++ b/tensorflow/contrib/layers/python/layers/feature_column_ops_test.py @@ -1757,7 +1757,7 @@ class WeightedSumTest(test.TestCase): logits_core = fc_core.linear_model(features, [movies]) with self.cached_session() as sess: - variables_lib.initialize_all_variables().run() + variables_lib.global_variables_initializer().run() lookup_ops.tables_initializer().run() weights = column_to_variable[movies][0] diff --git a/tensorflow/contrib/layers/python/layers/layers_test.py b/tensorflow/contrib/layers/python/layers/layers_test.py index d791418c9d0f887058ceb535092fa8122da1aa75..1c0088186c030437454c0f764decab9e5a276adc 100644 --- a/tensorflow/contrib/layers/python/layers/layers_test.py +++ b/tensorflow/contrib/layers/python/layers/layers_test.py @@ -1356,7 +1356,7 @@ class DropoutTest(test.TestCase): with self.cached_session(): images = np.random.uniform(size=(5, height, width, 3)) output = _layers.dropout(images) - self.assertEqual(output.op.name, 'Dropout/dropout_1/mul') + self.assertEqual(output.op.name, 'Dropout/dropout_1/mul_1') output.get_shape().assert_is_compatible_with( ops.convert_to_tensor(images).get_shape()) diff --git a/tensorflow/contrib/layers/python/layers/normalization.py b/tensorflow/contrib/layers/python/layers/normalization.py index 11033a2e9cb646c2e7cd2f45de1f751d88c6921a..76b03ff514821d3459f84c5f46a64d1134e0d4de 100644 --- a/tensorflow/contrib/layers/python/layers/normalization.py +++ b/tensorflow/contrib/layers/python/layers/normalization.py @@ -186,7 +186,7 @@ def group_norm(inputs, Args: inputs: A Tensor with at least 2 dimensions one which is channels. All - shape dimensions must be fully defined. + shape dimensions except for batch must be fully defined. groups: Integer. Divide the channels into this number of groups over which normalization statistics are computed. This number must be commensurate with the number of channels in `inputs`. @@ -249,13 +249,21 @@ def group_norm(inputs, """ # TODO(shlens): Support partially defined shapes for the inputs. inputs = ops.convert_to_tensor(inputs) - original_shape = inputs.shape if inputs.shape.ndims is None: raise ValueError('Inputs %s has undefined rank.' % inputs.name) if channels_axis > (inputs.shape.ndims - 1): raise ValueError('Axis is out of bounds.') + # Use dynamic shape for not fully defined dimensions in the inputs. + dyanmic_shape = array_ops.shape(inputs) + input_shape_list = [] + for i, dim in enumerate(inputs.shape): + if dim.value is None: + input_shape_list.append(dyanmic_shape[i]) + else: + input_shape_list.append(dim) + # Standardize the channels_axis to be positive and identify # of channels. if channels_axis < 0: channels_axis = inputs.shape.ndims + channels_axis @@ -289,8 +297,8 @@ def group_norm(inputs, # Determine axes before channels. Some examples of common image formats: # 'NCHW': before = [N], after = [HW] # 'NHWC': before = [NHW], after = [] - axes_before_channels = inputs.shape.as_list()[:channels_axis] - axes_after_channels = inputs.shape.as_list()[channels_axis+1:] + axes_before_channels = input_shape_list[:channels_axis] + axes_after_channels = input_shape_list[channels_axis+1:] # Manually broadcast the parameters to conform to the number of groups. params_shape_broadcast = ([1] * len(axes_before_channels) + @@ -369,7 +377,7 @@ def group_norm(inputs, outputs = inputs * gain + offset # Collapse the groups into the channel dimension. - outputs = array_ops.reshape(outputs, original_shape) + outputs = array_ops.reshape(outputs, input_shape_list) if activation_fn is not None: outputs = activation_fn(outputs) diff --git a/tensorflow/contrib/layers/python/layers/normalization_test.py b/tensorflow/contrib/layers/python/layers/normalization_test.py index c8d3c91b10dbe3b959e91182f9924b78352d370d..9a85084b239837ade87d8c778393ef8e885f5bdd 100644 --- a/tensorflow/contrib/layers/python/layers/normalization_test.py +++ b/tensorflow/contrib/layers/python/layers/normalization_test.py @@ -221,6 +221,15 @@ class GroupNormTest(test.TestCase): normalization.group_norm(inputs, channels_axis=-1, reduction_axes=[-3, -2]) + def testParamsShapeNotFullyDefinedBatchAxis(self): + height, width, groups = 3, 3, 4 + inputs = array_ops.placeholder(dtypes.float32, + shape=(None, height, width, 2*groups)) + output = normalization.group_norm(inputs, channels_axis=-1, + reduction_axes=[-3, -2], groups=groups) + self.assertListEqual([None, height, width, 2 * groups], + output.shape.as_list()) + def testCreateOp(self): height, width, groups = 3, 3, 4 images = random_ops.random_uniform((5, height, width, 2*groups), seed=1) diff --git a/tensorflow/contrib/layers/python/layers/target_column.py b/tensorflow/contrib/layers/python/layers/target_column.py index 8a6b4f68a8b33d497ddb16614a7e3cdf32f2c422..5234869718b427d7e275b76ae12021a096241a56 100644 --- a/tensorflow/contrib/layers/python/layers/target_column.py +++ b/tensorflow/contrib/layers/python/layers/target_column.py @@ -399,7 +399,7 @@ def _mean_squared_loss(logits, target): target = array_ops.expand_dims(target, axis=1) logits.get_shape().assert_is_compatible_with(target.get_shape()) - return math_ops.square(logits - math_ops.to_float(target)) + return math_ops.squared_difference(logits, math_ops.to_float(target)) def _log_loss_with_two_classes(logits, target): diff --git a/tensorflow/contrib/learn/BUILD b/tensorflow/contrib/learn/BUILD index 14065fcee51c014a1af227504eaaca1fa39941e1..4749371248ee89a033912132986d7f76c85dbaa6 100644 --- a/tensorflow/contrib/learn/BUILD +++ b/tensorflow/contrib/learn/BUILD @@ -357,9 +357,9 @@ py_test( py_test( name = "dnn_linear_combined_test", - size = "large", + size = "medium", srcs = ["python/learn/estimators/dnn_linear_combined_test.py"], - shard_count = 4, + shard_count = 8, srcs_version = "PY2AND3", tags = ["no_oss"], # flaky b/70524820 deps = [ diff --git a/tensorflow/contrib/learn/README.md b/tensorflow/contrib/learn/README.md index b0bff915a993c9a01e2e6d9ef9f71c14d2f29a73..b2d3a6273abba7e3a893f30bbdd4f8b2662bd54a 100644 --- a/tensorflow/contrib/learn/README.md +++ b/tensorflow/contrib/learn/README.md @@ -111,18 +111,17 @@ Some arguments are renamed, please refer to documentation. In addition: Switch to `tf.estimator.train_and_evaluate`. Some differences: -* Most of the constructor arguments, like `train_input_fn`, `eval_input_fn`, - should be wrapped into `tf.estimator.TrainSpec` and `tf.estimator.EvalSpec`. -* Remove the `experiment_fn`. Instead, create the `Estimator`, - `train_spec` and `eval_spec`, then call `tf.estimator.train_and_evaluate` - directly. -* Inside `tf.estimator.EvalSpec`, the `exporter` field is the replacement - for `export_strategy`. To be precise, `tf.estimator.LatestExporter` is the - replacement for `tf.contrib.learn.make_export_strategy`. If you want to export - only at the end of training use `tf.estimator.FinalExporter`. -* If the `TF_CONFIG` environment variable is constructed manually, please read - the `train_and_evaluate` documentation for the new requirementds (in - particular, the chief node and evaluator node). +* Most of the constructor arguments, like `train_input_fn`, `eval_input_fn`, + should be wrapped into `tf.estimator.TrainSpec` and `tf.estimator.EvalSpec`. +* Remove the `experiment_fn`. Instead, create the `Estimator`, `train_spec` + and `eval_spec`, then call `tf.estimator.train_and_evaluate` directly. +* Inside `tf.estimator.EvalSpec`, the `exporter` field is the replacement for + `export_strategy`. To be precise, `tf.estimator.LatestExporter` is the + replacement for `tf.contrib.learn.make_export_strategy`. If you want to + export only at the end of training use `tf.estimator.FinalExporter`. +* If the `TF_CONFIG` environment variable is constructed manually, please read + the `train_and_evaluate` documentation for the new requirements (in + particular, the chief node and evaluator node). ## Others Classes and Functions diff --git a/tensorflow/contrib/learn/python/learn/estimators/dynamic_rnn_estimator_test.py b/tensorflow/contrib/learn/python/learn/estimators/dynamic_rnn_estimator_test.py index 28c4964527bb034c8c6b1642366c6c82c1a72201..c3e9e3af9427037a4e7be6b86417cd081c42ef67 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/dynamic_rnn_estimator_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/dynamic_rnn_estimator_test.py @@ -37,8 +37,8 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import random_seed from tensorflow.python.framework import sparse_tensor from tensorflow.python.ops import array_ops -from tensorflow.python.ops import functional_ops from tensorflow.python.ops import lookup_ops +from tensorflow.python.ops import map_fn from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import rnn_cell @@ -524,7 +524,7 @@ class DynamicRNNEstimatorLearningTest(test.TestCase): def input_fn(): starts = random_ops.random_uniform( [batch_size], maxval=(2 * np.pi), seed=seed) - sin_curves = functional_ops.map_fn( + sin_curves = map_fn.map_fn( _sin_fn, (starts,), dtype=dtypes.float32) inputs = array_ops.expand_dims( array_ops.slice(sin_curves, [0, 0], [batch_size, sequence_length]), diff --git a/tensorflow/contrib/learn/python/learn/estimators/head.py b/tensorflow/contrib/learn/python/learn/estimators/head.py index c1b97d8b49613ea49d9813954da3b7a63d3ba04c..4bb14a6e63b159fa4d09c9ef20947d4b125de657 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/head.py +++ b/tensorflow/contrib/learn/python/learn/estimators/head.py @@ -567,7 +567,8 @@ def _mean_squared_loss(labels, logits, weights=None): if len(logits.get_shape()) == 1: logits = array_ops.expand_dims(logits, axis=1) logits.get_shape().assert_is_compatible_with(labels.get_shape()) - loss = math_ops.square(logits - math_ops.to_float(labels), name=name) + loss = math_ops.squared_difference( + logits, math_ops.to_float(labels), name=name) return _compute_weighted_loss(loss, weights) diff --git a/tensorflow/contrib/learn/python/learn/learn_io/generator_io_test.py b/tensorflow/contrib/learn/python/learn/learn_io/generator_io_test.py index 5e90d1fa20535de3b5e25bc7ff8c3862cea5514c..318046733bf75a6d661d26f478118c8e944afe15 100644 --- a/tensorflow/contrib/learn/python/learn/learn_io/generator_io_test.py +++ b/tensorflow/contrib/learn/python/learn/learn_io/generator_io_test.py @@ -174,7 +174,7 @@ class GeneratorIoTest(test.TestCase): return np.arange(32, 36) with self.cached_session(): - with self.assertRaisesRegexp(TypeError, 'x\(\) must be generator'): + with self.assertRaisesRegexp(TypeError, r'x\(\) must be generator'): failing_input_fn = generator_io.generator_input_fn( generator, batch_size=2, shuffle=False, num_epochs=1) failing_input_fn() @@ -185,7 +185,7 @@ class GeneratorIoTest(test.TestCase): yield np.arange(32, 36) with self.cached_session(): - with self.assertRaisesRegexp(TypeError, 'x\(\) must yield dict'): + with self.assertRaisesRegexp(TypeError, r'x\(\) must yield dict'): failing_input_fn = generator_io.generator_input_fn( generator, batch_size=2, shuffle=False, num_epochs=1) failing_input_fn() diff --git a/tensorflow/contrib/learn/python/learn/utils/gc_test.py b/tensorflow/contrib/learn/python/learn/utils/gc_test.py index e7d091e18a8f186f89f5217442c24fb106c5cdab..af93e517f51ed33a8968982945ac1f65ec915ab1 100644 --- a/tensorflow/contrib/learn/python/learn/utils/gc_test.py +++ b/tensorflow/contrib/learn/python/learn/utils/gc_test.py @@ -36,10 +36,10 @@ def _create_parser(base_dir): # Modify the path object for RegEx match for Windows Paths if os.name == "nt": match = re.match( - "^" + compat.as_str_any(base_dir).replace("\\", "/") + "/(\\d+)$", + r"^" + compat.as_str_any(base_dir).replace("\\", "/") + r"/(\d+)$", compat.as_str_any(path.path).replace("\\", "/")) else: - match = re.match("^" + compat.as_str_any(base_dir) + "/(\\d+)$", + match = re.match(r"^" + compat.as_str_any(base_dir) + r"/(\d+)$", compat.as_str_any(path.path)) if not match: return None diff --git a/tensorflow/contrib/lookup/lookup_ops.py b/tensorflow/contrib/lookup/lookup_ops.py index 229a72a780d5ccce8263444ffeae7700f6ac8613..3d21fb68a1452c97f7eb85491fc850d9e846266a 100644 --- a/tensorflow/contrib/lookup/lookup_ops.py +++ b/tensorflow/contrib/lookup/lookup_ops.py @@ -18,8 +18,6 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import functools - from tensorflow.python.eager import context from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -28,7 +26,6 @@ from tensorflow.python.ops import lookup_ops # pylint: disable=unused-import from tensorflow.python.ops.lookup_ops import FastHashSpec from tensorflow.python.ops.lookup_ops import HasherSpec -from tensorflow.python.ops.lookup_ops import HashTable 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 @@ -42,7 +39,6 @@ from tensorflow.python.ops.lookup_ops import TextFileIndex from tensorflow.python.ops.lookup_ops import TextFileInitializer from tensorflow.python.ops.lookup_ops import TextFileStringTableInitializer # pylint: enable=unused-import -from tensorflow.python.training.saver import BaseSaverBuilder from tensorflow.python.util.deprecation import deprecated @@ -288,353 +284,52 @@ def index_to_string(tensor, mapping, default_value="UNK", name=None): return table.lookup(tensor) -class MutableHashTable(LookupInterface): - """A generic mutable hash table implementation. - - Data can be inserted by calling the insert method and removed by calling the - remove method. It does not support initialization via the init method. +class HashTable(InitializableLookupTableBase): + """A generic hash table implementation. Example usage: ```python - table = tf.contrib.lookup.MutableHashTable(key_dtype=tf.string, - value_dtype=tf.int64, - default_value=-1) - sess.run(table.insert(keys, values)) - out = table.lookup(query_keys) + table = tf.HashTable( + tf.KeyValueTensorInitializer(keys, values), -1) + out = table.lookup(input_tensor) + table.init.run() print(out.eval()) ``` """ - def __init__(self, - key_dtype, - value_dtype, - default_value, - shared_name=None, - name="MutableHashTable", - checkpoint=True): - """Creates an empty `MutableHashTable` object. + def __init__(self, initializer, default_value, shared_name=None, name=None): + """Creates a non-initialized `HashTable` object. - Creates a table, the type of its keys and values are specified by key_dtype - and value_dtype, respectively. + Creates a table, the type of its keys and values are specified by the + initializer. + Before using the table you will have to initialize it. After initialization + the table will be immutable. Args: - key_dtype: the type of the key tensors. - value_dtype: the type of the value tensors. + initializer: The table initializer to use. See `HashTable` kernel for + supported key and value types. default_value: The value to use if a key is missing in the table. - shared_name: If non-empty, this table will be shared under - the given name across multiple sessions. + shared_name: If non-empty, this table will be shared under the given name + across multiple sessions. name: A name for the operation (optional). - checkpoint: if True, the contents of the table are saved to and restored - from checkpoints. If `shared_name` is empty for a checkpointed table, it - is shared using the table node name. Returns: - A `MutableHashTable` object. - - Raises: - ValueError: If checkpoint is True and no name was specified. + A `HashTable` object. """ - self._default_value = ops.convert_to_tensor(default_value, - dtype=value_dtype) - self._value_shape = self._default_value.get_shape() - self._checkpoint = checkpoint - self._key_dtype = key_dtype - self._value_dtype = value_dtype - self._name = name - - if context.executing_eagerly() and shared_name is None: - # TODO(allenl): This will leak memory due to kernel caching by the - # shared_name attribute value (but is better than the alternative of - # sharing everything by default when executing eagerly; hopefully creating - # tables in a loop is uncommon). - shared_name = "table_%d" % (ops.uid(),) + self._initializer = initializer + self._default_value = default_value self._shared_name = shared_name - super(MutableHashTable, self).__init__(key_dtype, value_dtype) - - self._resource_handle = self.create_resource() - if checkpoint: - saveable = MutableHashTable._Saveable(self, name) - if not context.executing_eagerly(): - ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable) - - def create_resource(self): - # The table must be shared if checkpointing is requested for multi-worker - # training to work correctly. Use the node name if no shared_name has been - # explicitly specified. - use_node_name_sharing = self._checkpoint and self._shared_name is None - if self._default_value.get_shape().ndims == 0: - table_ref = gen_lookup_ops.mutable_hash_table_v2( - shared_name=self._shared_name, - use_node_name_sharing=use_node_name_sharing, - key_dtype=self._key_dtype, - value_dtype=self._value_dtype, - name=self._name) - else: - table_ref = gen_lookup_ops.mutable_hash_table_of_tensors_v2( - shared_name=self._shared_name, - use_node_name_sharing=use_node_name_sharing, - key_dtype=self._key_dtype, - value_dtype=self._value_dtype, - value_shape=self._default_value.get_shape(), - name=self._name) - - if context.executing_eagerly(): - self._table_name = None - else: - self._table_name = table_ref.op.name.split("/")[-1] - return table_ref - - @property - def name(self): - return self._table_name - - def size(self, name=None): - """Compute the number of elements in this table. - - Args: - name: A name for the operation (optional). - - Returns: - A scalar tensor containing the number of elements in this table. - """ - with ops.name_scope(name, "%s_Size" % self.name, - [self.resource_handle]) as name: - with ops.colocate_with(self.resource_handle): - return gen_lookup_ops.lookup_table_size_v2( - self.resource_handle, name=name) - - def remove(self, keys, name=None): - """Removes `keys` and its associated values from the table. - - If a key is not present in the table, it is silently ignored. - - Args: - keys: Keys to remove. Can be a tensor of any shape. Must match the table's - key type. - name: A name for the operation (optional). - - Returns: - The created Operation. - - Raises: - TypeError: when `keys` do not match the table data types. - """ - if keys.dtype != self._key_dtype: - raise TypeError("Signature mismatch. Keys must be dtype %s, got %s." % - (self._key_dtype, keys.dtype)) - - with ops.name_scope( - name, "%s_lookup_table_remove" % self.name, - (self.resource_handle, keys, self._default_value)) as name: - # pylint: disable=protected-access - op = gen_lookup_ops.lookup_table_remove_v2( - self.resource_handle, keys, name=name) - - return op - - def lookup(self, keys, name=None): - """Looks up `keys` in a table, outputs the corresponding values. - - The `default_value` is used for keys not present in the table. - - Args: - keys: Keys to look up. Can be a tensor of any shape. Must match the - table's key_dtype. - name: A name for the operation (optional). - - Returns: - A tensor containing the values in the same shape as `keys` using the - table's value type. - - Raises: - TypeError: when `keys` do not match the table data types. - """ - with ops.name_scope( - name, "%s_lookup_table_find" % self.name, - (self.resource_handle, keys, self._default_value)) as name: - keys = ops.convert_to_tensor(keys, dtype=self._key_dtype, name="keys") - with ops.colocate_with(self.resource_handle): - values = gen_lookup_ops.lookup_table_find_v2( - self.resource_handle, keys, self._default_value, name=name) - return values - - def insert(self, keys, values, name=None): - """Associates `keys` with `values`. - - Args: - keys: Keys to insert. Can be a tensor of any shape. Must match the - table's key type. - values: Values to be associated with keys. Must be a tensor of the same - shape as `keys` and match the table's value type. - name: A name for the operation (optional). - - Returns: - The created Operation. - - Raises: - TypeError: when `keys` or `values` doesn't match the table data - types. - """ - with ops.name_scope(name, "%s_lookup_table_insert" % self.name, - [self.resource_handle, keys, values]) as name: - keys = ops.convert_to_tensor(keys, self._key_dtype, name="keys") - values = ops.convert_to_tensor(values, self._value_dtype, name="values") - with ops.colocate_with(self.resource_handle): - # pylint: disable=protected-access - op = gen_lookup_ops.lookup_table_insert_v2( - self.resource_handle, keys, values, name=name) - return op - - def export(self, name=None): - """Returns tensors of all keys and values in the table. - - Args: - name: A name for the operation (optional). - - Returns: - A pair of tensors with the first tensor containing all keys and the - second tensors containing all values in the table. - """ - with ops.name_scope(name, "%s_lookup_table_export_values" % self.name, - [self.resource_handle]) as name: - with ops.colocate_with(self.resource_handle): - exported_keys, exported_values = gen_lookup_ops.lookup_table_export_v2( - self.resource_handle, self._key_dtype, self._value_dtype, name=name) - return exported_keys, exported_values - - def _gather_saveables_for_checkpoint(self): - """For object-based checkpointing.""" - return {"table": functools.partial(MutableHashTable._Saveable, table=self)} - - class _Saveable(BaseSaverBuilder.SaveableObject): - """SaveableObject implementation for MutableHashTable.""" - - def __init__(self, table, name): - tensors = table.export() - specs = [ - BaseSaverBuilder.SaveSpec(tensors[0], "", name + "-keys"), - BaseSaverBuilder.SaveSpec(tensors[1], "", name + "-values") - ] - # pylint: disable=protected-access - super(MutableHashTable._Saveable, self).__init__(table, specs, name) - - def restore(self, restored_tensors, restored_shapes): - del restored_shapes # unused - # pylint: disable=protected-access - with ops.colocate_with(self.op.resource_handle): - return gen_lookup_ops.lookup_table_import_v2( - self.op.resource_handle, restored_tensors[0], restored_tensors[1]) - - -class MutableDenseHashTable(LookupInterface): - """A generic mutable hash table implementation using tensors as backing store. - - Data can be inserted by calling the insert method and removed by calling the - remove method. It does not support initialization via the init method. - - It uses "open addressing" with quadratic reprobing to resolve collisions. - Compared to `MutableHashTable` the insert, remove and lookup operations in a - `MutableDenseHashTable` are typically faster, but memory usage can be higher. - However, `MutableDenseHashTable` does not require additional memory for - temporary tensors created during checkpointing and restore operations. - - Example usage: - - ```python - table = tf.contrib.lookup.MutableDenseHashTable(key_dtype=tf.int64, - value_dtype=tf.int64, - default_value=-1, - empty_key=0, - deleted_key=-1) - - sess.run(table.insert(keys, values)) - out = table.lookup(query_keys) - print(out.eval()) - ``` - """ - - # TODO(andreasst): consider extracting common code with MutableHashTable into - # a common superclass. - def __init__(self, - key_dtype, - value_dtype, - default_value, - empty_key, - deleted_key, - initial_num_buckets=None, - shared_name=None, - name="MutableDenseHashTable", - checkpoint=True): - """Creates an empty `MutableDenseHashTable` object. - - Creates a table, the type of its keys and values are specified by key_dtype - and value_dtype, respectively. - - Args: - key_dtype: the type of the key tensors. - value_dtype: the type of the value tensors. - default_value: The value to use if a key is missing in the table. - empty_key: the key to use to represent empty buckets internally. Must not - be used in insert, remove or lookup operations. - initial_num_buckets: the initial number of buckets. - shared_name: If non-empty, this table will be shared under - the given name across multiple sessions. - name: A name for the operation (optional). - checkpoint: if True, the contents of the table are saved to and restored - from checkpoints. If `shared_name` is empty for a checkpointed table, it - is shared using the table node name. - deleted_key: the key to use to represent deleted buckets internally. Must - not be used in insert, remove or lookup operations and be different from - the empty_key. - - Returns: - A `MutableDenseHashTable` object. - - Raises: - ValueError: If checkpoint is True and no name was specified. - """ - self._default_value = ops.convert_to_tensor( - default_value, dtype=value_dtype, name="default_value") - self._key_dtype = key_dtype - self._value_dtype = value_dtype - self._initial_num_buckets = initial_num_buckets + self._name = name or "hash_table" + self._table_name = None + super(HashTable, self).__init__(default_value, initializer) self._value_shape = self._default_value.get_shape() - self._checkpoint = checkpoint - self._name = name - - self._empty_key = ops.convert_to_tensor( - empty_key, dtype=key_dtype, name="empty_key") - self._deleted_key = ops.convert_to_tensor( - deleted_key, dtype=key_dtype, name="deleted_key") - if context.executing_eagerly() and shared_name is None: - # TODO(allenl): This will leak memory due to kernel caching by the - # shared_name attribute value (but is better than the alternative of - # sharing everything by default when executing eagerly; hopefully creating - # tables in a loop is uncommon). - shared_name = "table_%d" % (ops.uid(),) - self._shared_name = shared_name - super(MutableDenseHashTable, self).__init__(key_dtype, value_dtype) - - self._resource_handle = self.create_resource() - if checkpoint: - saveable = MutableDenseHashTable._Saveable(self, name) - if not context.executing_eagerly(): - ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable) def create_resource(self): - # The table must be shared if checkpointing is requested for multi-worker - # training to work correctly. Use the node name if no shared_name has been - # explicitly specified. - use_node_name_sharing = self._checkpoint and self._shared_name is None - table_ref = gen_lookup_ops.mutable_dense_hash_table_v2( - empty_key=self._empty_key, - deleted_key=self._deleted_key, + table_ref = gen_lookup_ops.hash_table_v2( shared_name=self._shared_name, - use_node_name_sharing=use_node_name_sharing, - value_dtype=self._value_dtype, - value_shape=self._value_shape, - initial_num_buckets=self._initial_num_buckets, + key_dtype=self._initializer.key_dtype, + value_dtype=self._initializer.value_dtype, name=self._name) if context.executing_eagerly(): self._table_name = None @@ -646,103 +341,6 @@ class MutableDenseHashTable(LookupInterface): def name(self): return self._table_name - def size(self, name=None): - """Compute the number of elements in this table. - - Args: - name: A name for the operation (optional). - - Returns: - A scalar tensor containing the number of elements in this table. - """ - with ops.name_scope(name, "%s_Size" % self.name, - [self.resource_handle]) as name: - with ops.colocate_with(self.resource_handle): - return gen_lookup_ops.lookup_table_size_v2( - self.resource_handle, name=name) - - def lookup(self, keys, name=None): - """Looks up `keys` in a table, outputs the corresponding values. - - The `default_value` is used for keys not present in the table. - - Args: - keys: Keys to look up. Can be a tensor of any shape. Must match the - table's key_dtype. - name: A name for the operation (optional). - - Returns: - A tensor containing the values in the same shape as `keys` using the - table's value type. - - Raises: - TypeError: when `keys` do not match the table data types. - """ - with ops.name_scope(name, "%s_lookup_table_find" % self.name, - [self.resource_handle, keys]) as name: - keys = ops.convert_to_tensor(keys, dtype=self._key_dtype, name="keys") - with ops.colocate_with(self.resource_handle): - values = gen_lookup_ops.lookup_table_find_v2( - self.resource_handle, keys, self._default_value, name=name) - - return values - - def insert(self, keys, values, name=None): - """Associates `keys` with `values`. - - Args: - keys: Keys to insert. Can be a tensor of any shape. Must match the - table's key type. - values: Values to be associated with keys. Must be a tensor of the same - shape as `keys` and match the table's value type. - name: A name for the operation (optional). - - Returns: - The created Operation. - - Raises: - TypeError: when `keys` or `values` doesn't match the table data - types. - """ - with ops.name_scope(name, "%s_lookup_table_insert" % self.name, - [self.resource_handle, keys, values]) as name: - keys = ops.convert_to_tensor(keys, dtype=self._key_dtype, name="keys") - values = ops.convert_to_tensor( - values, dtype=self._value_dtype, name="values") - with ops.colocate_with(self.resource_handle): - op = gen_lookup_ops.lookup_table_insert_v2( - self.resource_handle, keys, values, name=name) - return op - - def remove(self, keys, name=None): - """Removes `keys` and its associated values from the table. - - If a key is not present in the table, it is silently ignored. - - Args: - keys: Keys to remove. Can be a tensor of any shape. Must match the table's - key type. - name: A name for the operation (optional). - - Returns: - The created Operation. - - Raises: - TypeError: when `keys` do not match the table data types. - """ - if keys.dtype != self._key_dtype: - raise TypeError("Signature mismatch. Keys must be dtype %s, got %s." % - (self._key_dtype, keys.dtype)) - - with ops.name_scope( - name, "%s_lookup_table_remove" % self.name, - (self.resource_handle, keys, self._default_value)) as name: - # pylint: disable=protected-access - op = gen_lookup_ops.lookup_table_remove_v2( - self.resource_handle, keys, name=name) - - return op - def export(self, name=None): """Returns tensors of all keys and values in the table. @@ -753,34 +351,15 @@ class MutableDenseHashTable(LookupInterface): A pair of tensors with the first tensor containing all keys and the second tensors containing all values in the table. """ - with ops.name_scope(name, "%s_lookup_table_export_values" % self.name, + with ops.name_scope(name, "%s_Export" % self.name, [self.resource_handle]) as name: - with ops.colocate_with(self.resource_handle): - exported_keys, exported_values = gen_lookup_ops.lookup_table_export_v2( - self.resource_handle, self._key_dtype, self._value_dtype, name=name) + exported_keys, exported_values = gen_lookup_ops.lookup_table_export_v2( + self.resource_handle, self._key_dtype, self._value_dtype, name=name) + exported_values.set_shape(exported_keys.get_shape().concatenate( + self._value_shape)) return exported_keys, exported_values - def _gather_saveables_for_checkpoint(self): - """For object-based checkpointing.""" - return {"table": functools.partial( - MutableDenseHashTable._Saveable, table=self)} - - class _Saveable(BaseSaverBuilder.SaveableObject): - """SaveableObject implementation for MutableDenseHashTable.""" - - def __init__(self, table, name): - tensors = table.export() - specs = [ - BaseSaverBuilder.SaveSpec(tensors[0], "", name + "-keys"), - BaseSaverBuilder.SaveSpec(tensors[1], "", name + "-values") - ] - # pylint: disable=protected-access - super(MutableDenseHashTable._Saveable, self).__init__(table, specs, name) - - def restore(self, restored_tensors, restored_shapes): - del restored_shapes # unused - # pylint: disable=protected-access - with ops.colocate_with(self.op.resource_handle): - return gen_lookup_ops.lookup_table_import_v2( - self.op.resource_handle, restored_tensors[0], restored_tensors[1]) + +MutableHashTable = lookup_ops.MutableHashTable +MutableDenseHashTable = lookup_ops.MutableDenseHashTable diff --git a/tensorflow/contrib/lookup/lookup_ops_test.py b/tensorflow/contrib/lookup/lookup_ops_test.py index 9b2c2dd87cc8a92fbb6b45504939be3788b60839..591eabc66c49f301cf73cd912ebbef70cc9e1e3f 100644 --- a/tensorflow/contrib/lookup/lookup_ops_test.py +++ b/tensorflow/contrib/lookup/lookup_ops_test.py @@ -18,14 +18,10 @@ from __future__ import division from __future__ import print_function import os -import tempfile import numpy as np -import six from tensorflow.contrib import lookup from tensorflow.python.client import session -from tensorflow.python.data.experimental.ops import counter -from tensorflow.python.data.ops import dataset_ops from tensorflow.python.eager import context from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes @@ -37,9 +33,7 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import lookup_ops from tensorflow.python.ops import variables from tensorflow.python.platform import test -from tensorflow.python.training import saver from tensorflow.python.training import server_lib -from tensorflow.python.training.checkpointable import util as checkpointable class HashTableOpTest(test.TestCase): @@ -299,1240 +293,6 @@ class HashTableOpTest(test.TestCase): self.assertAllEqual([b"brain", b"salad", b"n/a"], result) -class MutableHashTableOpTest(test.TestCase): - - def testMutableHashTable(self): - with self.cached_session(): - default_val = -1 - keys = constant_op.constant(["brain", "salad", "surgery", "tarkus"]) - values = constant_op.constant([0, 1, 2, 3], dtypes.int64) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(4, table.size().eval()) - - remove_string = constant_op.constant(["tarkus", "tank"]) - table.remove(remove_string).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant(["brain", "salad", "tank"]) - output = table.lookup(input_string) - self.assertAllEqual([3], output.get_shape()) - - result = output.eval() - self.assertAllEqual([0, 1, -1], result) - - exported_keys, exported_values = table.export() - self.assertAllEqual([None], exported_keys.get_shape().as_list()) - self.assertAllEqual([None], exported_values.get_shape().as_list()) - - # exported data is in the order of the internal map, i.e. undefined - sorted_keys = np.sort(exported_keys.eval()) - sorted_values = np.sort(exported_values.eval()) - self.assertAllEqual([b"brain", b"salad", b"surgery"], sorted_keys) - self.assertAllEqual([0, 1, 2], sorted_values) - - def testSaveRestore(self): - save_dir = os.path.join(self.get_temp_dir(), "save_restore") - save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash") - - with self.session(graph=ops.Graph()) as sess: - v0 = variables.Variable(10.0, name="v0") - v1 = variables.Variable(20.0, name="v1") - - default_val = -1 - keys = constant_op.constant(["b", "c", "d"], dtypes.string) - values = constant_op.constant([0, 1, 2], dtypes.int64) - table = lookup.MutableHashTable( - dtypes.string, dtypes.int64, default_val, name="t1", checkpoint=True) - - save = saver.Saver() - variables.global_variables_initializer().run() - - # Check that the parameter nodes have been initialized. - self.assertEqual(10.0, v0.eval()) - self.assertEqual(20.0, v1.eval()) - - self.assertAllEqual(0, table.size().eval()) - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - val = save.save(sess, save_path) - self.assertTrue(isinstance(val, six.string_types)) - self.assertEqual(save_path, val) - - with self.session(graph=ops.Graph()) as sess: - v0 = variables.Variable(-1.0, name="v0") - v1 = variables.Variable(-1.0, name="v1") - default_val = -1 - table = lookup.MutableHashTable( - dtypes.string, dtypes.int64, default_val, name="t1", checkpoint=True) - table.insert( - constant_op.constant(["a", "c"], dtypes.string), - constant_op.constant([12, 24], dtypes.int64)).run() - self.assertAllEqual(2, table.size().eval()) - - save = saver.Saver() - - # Restore the saved values in the parameter nodes. - save.restore(sess, save_path) - # Check that the parameter nodes have been restored. - self.assertEqual(10.0, v0.eval()) - self.assertEqual(20.0, v1.eval()) - - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant(["a", "b", "c", "d", "e"], - dtypes.string) - output = table.lookup(input_string) - self.assertAllEqual([-1, 0, 1, 2, -1], output.eval()) - - @test_util.run_in_graph_and_eager_modes - def testObjectSaveRestore(self): - save_dir = os.path.join(self.get_temp_dir(), "save_restore") - save_prefix = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash") - - v0 = variables.Variable(10.0, name="v0") - v1 = variables.Variable(20.0, name="v1") - - default_val = -1 - keys = constant_op.constant(["b", "c", "d"], dtypes.string) - values = constant_op.constant([0, 1, 2], dtypes.int64) - table = lookup.MutableHashTable( - dtypes.string, dtypes.int64, default_val, name="t1", checkpoint=True) - - checkpoint = checkpointable.Checkpoint(table=table, v0=v0, v1=v1) - self.evaluate([v0.initializer, v1.initializer]) - - # Check that the parameter nodes have been initialized. - self.assertEqual(10.0, self.evaluate(v0)) - self.assertEqual(20.0, self.evaluate(v1)) - - self.assertAllEqual(0, self.evaluate(table.size())) - self.evaluate(table.insert(keys, values)) - self.assertAllEqual(3, self.evaluate(table.size())) - - save_path = checkpoint.save(save_prefix) - del table, checkpoint, v0, v1 - - v0 = variables.Variable(-1.0, name="v0") - v1 = variables.Variable(-1.0, name="v1") - default_val = -1 - table = lookup.MutableHashTable( - dtypes.string, dtypes.int64, default_val, name="t1", checkpoint=True) - self.evaluate(table.insert( - constant_op.constant(["a", "c"], dtypes.string), - constant_op.constant([12, 24], dtypes.int64))) - self.assertAllEqual(2, self.evaluate(table.size())) - - checkpoint = checkpointable.Checkpoint(table=table, v0=v0, v1=v1) - - # Restore the saved values in the parameter nodes. - checkpoint.restore(save_path).run_restore_ops() - # Check that the parameter nodes have been restored. - self.assertEqual(10.0, self.evaluate(v0)) - self.assertEqual(20.0, self.evaluate(v1)) - - self.assertAllEqual(3, self.evaluate(table.size())) - - input_string = constant_op.constant(["a", "b", "c", "d", "e"], - dtypes.string) - output = table.lookup(input_string) - self.assertAllEqual([-1, 0, 1, 2, -1], self.evaluate(output)) - - def testSharing(self): - # Start a server to store the table state - server = server_lib.Server( - { - "local0": ["localhost:0"] - }, protocol="grpc", start=True) - # Create two sessions sharing the same state - session1 = session.Session(server.target) - session2 = session.Session(server.target) - - table = lookup.MutableHashTable( - dtypes.int64, dtypes.string, "-", name="t1") - - # Populate the table in the first session - with session1: - self.assertAllEqual(0, table.size().eval()) - - keys = constant_op.constant([11, 12], dtypes.int64) - values = constant_op.constant(["a", "b"]) - table.insert(keys, values).run() - self.assertAllEqual(2, table.size().eval()) - - output = table.lookup(constant_op.constant([11, 12, 13], dtypes.int64)) - self.assertAllEqual([b"a", b"b", b"-"], output.eval()) - - # Verify that we can access the shared data from the second session - with session2: - self.assertAllEqual(2, table.size().eval()) - - output = table.lookup(constant_op.constant([10, 11, 12], dtypes.int64)) - self.assertAllEqual([b"-", b"a", b"b"], output.eval()) - - def testMutableHashTableOfTensors(self): - with self.cached_session(): - default_val = constant_op.constant([-1, -1], dtypes.int64) - keys = constant_op.constant(["brain", "salad", "surgery", "tarkus"]) - values = constant_op.constant([[0, 1], [2, 3], [4, 5], [6, 7]], - dtypes.int64) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(4, table.size().eval()) - - remove_string = constant_op.constant(["tarkus", "tank"]) - table.remove(remove_string).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant(["brain", "salad", "tank"]) - output = table.lookup(input_string) - self.assertAllEqual([3, 2], output.get_shape()) - - result = output.eval() - self.assertAllEqual([[0, 1], [2, 3], [-1, -1]], result) - - exported_keys, exported_values = table.export() - self.assertAllEqual([None], exported_keys.get_shape().as_list(), - msg="Saw shape %s" % exported_keys.shape) - self.assertAllEqual([None, 2], exported_values.get_shape().as_list(), - msg="Saw shape %s" % exported_values.shape) - # exported data is in the order of the internal map, i.e. undefined - sorted_keys = np.sort(exported_keys.eval()) - sorted_values = np.sort(exported_values.eval()) - self.assertAllEqual([b"brain", b"salad", b"surgery"], sorted_keys) - self.assertAllEqual([[4, 5], [2, 3], [0, 1]], sorted_values) - - def testMutableHashTableExportInsert(self): - with self.cached_session(): - default_val = constant_op.constant([-1, -1], dtypes.int64) - keys = constant_op.constant(["brain", "salad", "surgery"]) - values = constant_op.constant([[0, 1], [2, 3], [4, 5]], dtypes.int64) - table1 = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - self.assertAllEqual(0, table1.size().eval()) - table1.insert(keys, values).run() - self.assertAllEqual(3, table1.size().eval()) - - input_string = constant_op.constant(["brain", "salad", "tank"]) - expected_output = [[0, 1], [2, 3], [-1, -1]] - output1 = table1.lookup(input_string) - self.assertAllEqual(expected_output, output1.eval()) - - exported_keys, exported_values = table1.export() - self.assertAllEqual(3, exported_keys.eval().size) - self.assertAllEqual(6, exported_values.eval().size) - - # Populate a second table from the exported data - table2 = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - self.assertAllEqual(0, table2.size().eval()) - table2.insert(exported_keys, exported_values).run() - self.assertAllEqual(3, table2.size().eval()) - - # Verify lookup result is still the same - output2 = table2.lookup(input_string) - self.assertAllEqual(expected_output, output2.eval()) - - def testMutableHashTableOfTensorsInvalidShape(self): - with self.cached_session(): - default_val = constant_op.constant([-1, -1], dtypes.int64) - keys = constant_op.constant(["brain", "salad", "surgery"]) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - - # Shape [6] instead of [3, 2] - values = constant_op.constant([0, 1, 2, 3, 4, 5], dtypes.int64) - with self.assertRaisesOpError("Expected shape"): - table.insert(keys, values).run() - - # Shape [2,3] instead of [3, 2] - values = constant_op.constant([[0, 1, 2], [3, 4, 5]], dtypes.int64) - with self.assertRaisesOpError("Expected shape"): - table.insert(keys, values).run() - - # Shape [2, 2] instead of [3, 2] - values = constant_op.constant([[0, 1], [2, 3]], dtypes.int64) - with self.assertRaisesOpError("Expected shape"): - table.insert(keys, values).run() - - # Shape [3, 1] instead of [3, 2] - values = constant_op.constant([[0], [2], [4]], dtypes.int64) - with self.assertRaisesOpError("Expected shape"): - table.insert(keys, values).run() - - # Valid Insert - values = constant_op.constant([[0, 1], [2, 3], [4, 5]], dtypes.int64) - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - def testMutableHashTableInvalidDefaultValue(self): - with self.cached_session(): - default_val = constant_op.constant([[-1, -1]], dtypes.int64) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - with self.assertRaisesOpError("Default value must be a vector"): - self.assertAllEqual(0, table.size().eval()) - - def testMutableHashTableDuplicateInsert(self): - with self.cached_session(): - default_val = -1 - keys = constant_op.constant(["brain", "salad", "surgery", "brain"]) - values = constant_op.constant([0, 1, 2, 3], dtypes.int64) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant(["brain", "salad", "tank"]) - output = table.lookup(input_string) - - result = output.eval() - self.assertAllEqual([3, 1, -1], result) - - def testMutableHashTableFindHighRank(self): - with self.cached_session(): - default_val = -1 - keys = constant_op.constant(["brain", "salad", "surgery"]) - values = constant_op.constant([0, 1, 2], dtypes.int64) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant( - [["brain", "salad"], ["tank", "tarkus"]]) - output = table.lookup(input_string) - self.assertAllEqual([2, 2], output.get_shape()) - - result = output.eval() - self.assertAllEqual([[0, 1], [-1, -1]], result) - - def testMutableHashTableInsertHighRank(self): - with self.cached_session(): - default_val = -1 - keys = constant_op.constant([["brain", "salad"], ["surgery", "tank"]]) - values = constant_op.constant([[0, 1], [2, 3]], dtypes.int64) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - - table.insert(keys, values).run() - self.assertAllEqual(4, table.size().eval()) - - input_string = constant_op.constant(["brain", "salad", "tank", "tarkus"]) - output = table.lookup(input_string) - - result = output.eval() - self.assertAllEqual([0, 1, 3, -1], result) - - def testMutableHashTableRemoveHighRank(self): - with self.test_session(): - default_val = -1 - keys = constant_op.constant([["brain", "salad"], ["surgery", "tank"]]) - values = constant_op.constant([[0, 1], [2, 3]], dtypes.int64) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, default_val) - - table.insert(keys, values).run() - self.assertAllEqual(4, table.size().eval()) - - remove_string = constant_op.constant(["salad", "tarkus"]) - table.remove(remove_string).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant(["brain", "salad", "tank", "tarkus"]) - output = table.lookup(input_string) - - result = output.eval() - self.assertAllEqual([0, -1, 3, -1], result) - - def testMutableHashTableOfTensorsFindHighRank(self): - with self.cached_session(): - default_val = constant_op.constant([-1, -1, -1], dtypes.int64) - keys = constant_op.constant(["brain", "salad", "surgery"]) - values = constant_op.constant([[0, 1, 2], [2, 3, 4], [4, 5, 6]], - dtypes.int64) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant( - [["brain", "salad"], ["tank", "tarkus"]]) - output = table.lookup(input_string) - self.assertAllEqual([2, 2, 3], output.get_shape()) - - result = output.eval() - self.assertAllEqual( - [[[0, 1, 2], [2, 3, 4]], [[-1, -1, -1], [-1, -1, -1]]], result) - - def testMutableHashTableOfTensorsRemoveHighRank(self): - with self.test_session(): - default_val = constant_op.constant([-1, -1, -1], dtypes.int64) - keys = constant_op.constant(["brain", "salad", "surgery"]) - values = constant_op.constant([[0, 1, 2], [2, 3, 4], [4, 5, 6]], - dtypes.int64) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, default_val) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - remove_string = constant_op.constant([["brain", "tank"]]) - table.remove(remove_string).run() - self.assertAllEqual(2, table.size().eval()) - - input_string = constant_op.constant([["brain", "salad"], - ["surgery", "tank"]]) - output = table.lookup(input_string) - self.assertAllEqual([2, 2, 3], output.get_shape()) - - result = output.eval() - self.assertAllEqual( - [[[-1, -1, -1], [2, 3, 4]], [[4, 5, 6], [-1, -1, -1]]], result) - - def testMultipleMutableHashTables(self): - with self.cached_session() as sess: - default_val = -1 - keys = constant_op.constant(["brain", "salad", "surgery"]) - values = constant_op.constant([0, 1, 2], dtypes.int64) - - table1 = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - table2 = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - table3 = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - table1.insert(keys, values).run() - table2.insert(keys, values).run() - table3.insert(keys, values).run() - - self.assertAllEqual(3, table1.size().eval()) - self.assertAllEqual(3, table2.size().eval()) - self.assertAllEqual(3, table3.size().eval()) - - input_string = constant_op.constant(["brain", "salad", "tank"]) - output1 = table1.lookup(input_string) - output2 = table2.lookup(input_string) - output3 = table3.lookup(input_string) - - out1, out2, out3 = sess.run([output1, output2, output3]) - self.assertAllEqual([0, 1, -1], out1) - self.assertAllEqual([0, 1, -1], out2) - self.assertAllEqual([0, 1, -1], out3) - - def testMutableHashTableWithTensorDefault(self): - with self.cached_session(): - default_val = constant_op.constant(-1, dtypes.int64) - keys = constant_op.constant(["brain", "salad", "surgery"]) - values = constant_op.constant([0, 1, 2], dtypes.int64) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant(["brain", "salad", "tank"]) - output = table.lookup(input_string) - - result = output.eval() - self.assertAllEqual([0, 1, -1], result) - - def testSignatureMismatch(self): - with self.cached_session(): - default_val = -1 - keys = constant_op.constant(["brain", "salad", "surgery"]) - values = constant_op.constant([0, 1, 2], dtypes.int64) - table = lookup.MutableHashTable(dtypes.string, dtypes.int64, - default_val) - - # insert with keys of the wrong type - with self.assertRaises(ValueError): - table.insert(constant_op.constant([4, 5, 6]), values).run() - - # insert with values of the wrong type - with self.assertRaises(ValueError): - table.insert(keys, constant_op.constant(["a", "b", "c"])).run() - - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - input_string_ref = variables.Variable("brain") - input_int64_ref = variables.Variable(-1, dtype=dtypes.int64) - variables.global_variables_initializer().run() - - # Ref types do not produce an insert signature mismatch. - table.insert(input_string_ref, input_int64_ref).run() - self.assertAllEqual(3, table.size().eval()) - - # Ref types do not produce a lookup signature mismatch. - self.assertEqual(-1, table.lookup(input_string_ref).eval()) - - # lookup with keys of the wrong type - input_string = constant_op.constant([1, 2, 3], dtypes.int64) - with self.assertRaises(ValueError): - table.lookup(input_string).eval() - - # default value of the wrong type - with self.assertRaises(TypeError): - lookup.MutableHashTable(dtypes.string, dtypes.int64, "UNK") - - def testMutableHashTableStringFloat(self): - with self.cached_session(): - default_val = -1.5 - keys = constant_op.constant(["brain", "salad", "surgery"]) - values = constant_op.constant([0, 1.1, 2.2], dtypes.float32) - table = lookup.MutableHashTable(dtypes.string, dtypes.float32, - default_val) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant(["brain", "salad", "tank"]) - output = table.lookup(input_string) - - result = output.eval() - self.assertAllClose([0, 1.1, default_val], result) - - def testMutableHashTableIntFloat(self): - with self.cached_session(): - default_val = -1.0 - keys = constant_op.constant([3, 7, 0], dtypes.int64) - values = constant_op.constant([7.5, -1.2, 9.9], dtypes.float32) - table = lookup.MutableHashTable(dtypes.int64, dtypes.float32, - default_val) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant([7, 0, 11], dtypes.int64) - output = table.lookup(input_string) - - result = output.eval() - self.assertAllClose([-1.2, 9.9, default_val], result) - - def testMutableHashTableInt64String(self): - with self.cached_session(): - default_val = "n/a" - keys = constant_op.constant([0, 1, 2], dtypes.int64) - values = constant_op.constant(["brain", "salad", "surgery"]) - table = lookup.MutableHashTable(dtypes.int64, dtypes.string, - default_val) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant([0, 1, 3], dtypes.int64) - output = table.lookup(input_string) - - result = output.eval() - self.assertAllEqual((b"brain", b"salad", b"n/a"), result) - - -class MutableDenseHashTableOpTest(test.TestCase): - - def testBasic(self): - with self.cached_session(): - - keys = constant_op.constant([11, 12, 13, 14], dtypes.int64) - values = constant_op.constant([0, 1, 2, 3], dtypes.int64) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=-1, - empty_key=0, - deleted_key=-1) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(4, table.size().eval()) - - remove_string = constant_op.constant([12, 15], dtypes.int64) - table.remove(remove_string).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant([11, 12, 15], dtypes.int64) - output = table.lookup(input_string) - self.assertAllEqual([3], output.get_shape()) - - result = output.eval() - self.assertAllEqual([0, -1, -1], result) - - def testBasicBool(self): - with self.cached_session(): - - keys = constant_op.constant([11, 12, 13, 14], dtypes.int64) - values = constant_op.constant([True, True, True, True], dtypes.bool) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.bool, - default_value=False, - empty_key=0, - deleted_key=-1) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(4, table.size().eval()) - - remove_string = constant_op.constant([11, 15], dtypes.int64) - table.remove(remove_string).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant([11, 12, 15], dtypes.int64) - output = table.lookup(input_string) - self.assertAllEqual([3], output.get_shape()) - - result = output.eval() - self.assertAllEqual([False, True, False], result) - - def testSameEmptyAndDeletedKey(self): - with self.cached_session(): - with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, - "deleted_key"): - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=-1, - empty_key=42, - deleted_key=42) - self.assertAllEqual(0, table.size().eval()) - - def testLookupUnknownShape(self): - with self.cached_session(): - keys = constant_op.constant([11, 12, 13], dtypes.int64) - values = constant_op.constant([0, 1, 2], dtypes.int64) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=-1, - empty_key=0, - deleted_key=-1) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - placeholder_keys = array_ops.placeholder(dtypes.int64) - output = table.lookup(placeholder_keys) - self.assertAllEqual(None, output.get_shape()) - result = output.eval({placeholder_keys: [11, 12, 15]}) - self.assertAllEqual([0, 1, -1], result) - - def testMapStringToFloat(self): - with self.cached_session(): - - keys = constant_op.constant(["a", "b", "c", "d"], dtypes.string) - values = constant_op.constant([0.0, 1.1, 2.2, 3.3], dtypes.float32) - default_value = constant_op.constant(-1.5, dtypes.float32) - table = lookup.MutableDenseHashTable( - dtypes.string, - dtypes.float32, - default_value=default_value, - empty_key="", - deleted_key="$") - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(4, table.size().eval()) - - remove_string = constant_op.constant(["b", "e"]) - table.remove(remove_string).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant(["a", "b", "d", "e"], dtypes.string) - output = table.lookup(input_string) - self.assertAllEqual([4], output.get_shape()) - - result = output.eval() - self.assertAllClose([0, -1.5, 3.3, -1.5], result) - - def testMapInt64ToFloat(self): - for float_dtype in [dtypes.float32, dtypes.float64]: - with self.cached_session(): - - keys = constant_op.constant([11, 12, 13, 14], dtypes.int64) - values = constant_op.constant([0.0, 1.1, 2.2, 3.3], float_dtype) - default_value = constant_op.constant(-1.5, float_dtype) - table = lookup.MutableDenseHashTable( - dtypes.int64, - float_dtype, - default_value=default_value, - empty_key=0, - deleted_key=-1) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(4, table.size().eval()) - - remove_string = constant_op.constant([12, 15], dtypes.int64) - table.remove(remove_string).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant([11, 12, 14, 15], dtypes.int64) - output = table.lookup(input_string) - self.assertAllEqual([4], output.get_shape()) - - result = output.eval() - self.assertAllClose([0, -1.5, 3.3, -1.5], result) - - def testVectorValues(self): - with self.cached_session(): - keys = constant_op.constant([11, 12, 13], dtypes.int64) - values = constant_op.constant([[0, 1, 2, 3], [3, 4, 5, 6], [6, 7, 8, 9]], - dtypes.int64) - default_value = constant_op.constant([-1, -2, -3, -4], dtypes.int64) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=default_value, - empty_key=0, - deleted_key=-1, - initial_num_buckets=4) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - self.assertAllEqual(4, len(table.export()[0].eval())) - - table.insert( - constant_op.constant([14], dtypes.int64), - constant_op.constant([[2, 3, 4, 5]], dtypes.int64)).run() - self.assertAllEqual(4, table.size().eval()) - self.assertAllEqual(8, len(table.export()[0].eval())) - - remove_string = constant_op.constant([12, 16], dtypes.int64) - table.remove(remove_string).run() - self.assertAllEqual(3, table.size().eval()) - self.assertAllEqual(8, len(table.export()[0].eval())) - - input_string = constant_op.constant([11, 12, 14, 15], dtypes.int64) - output = table.lookup(input_string) - self.assertAllEqual([4, 4], - output.shape, - msg="Saw shape: %s" % output.shape) - - result = output.eval() - self.assertAllEqual( - [[0, 1, 2, 3], [-1, -2, -3, -4], [2, 3, 4, 5], [-1, -2, -3, -4]], - result) - - def testVectorKeys(self): - with self.cached_session(): - keys = constant_op.constant([[0, 1], [1, 2], [1, 3]], dtypes.int64) - values = constant_op.constant([10, 11, 12], dtypes.int64) - empty_key = constant_op.constant([0, 3], dtypes.int64) - deleted_key = constant_op.constant([-1, -1], dtypes.int64) - default_value = constant_op.constant(-1, dtypes.int64) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=default_value, - empty_key=empty_key, - deleted_key=deleted_key, - initial_num_buckets=8) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - table.insert( - constant_op.constant([[0, 0]], dtypes.int64), - constant_op.constant([13], dtypes.int64)).run() - self.assertAllEqual(4, table.size().eval()) - self.assertAllEqual(8, len(table.export()[0].eval())) - - remove_string = constant_op.constant([[1, 2], [7, 8]], dtypes.int64) - table.remove(remove_string).run() - self.assertAllEqual(3, table.size().eval()) - self.assertAllEqual(8, len(table.export()[0].eval())) - - input_string = constant_op.constant([[0, 1], [1, 2], [1, 3], [0, 2]], - dtypes.int64) - output = table.lookup(input_string) - self.assertAllEqual([4], output.get_shape()) - - result = output.eval() - self.assertAllEqual([10, -1, 12, -1], result) - - def testResize(self): - with self.cached_session(): - keys = constant_op.constant([11, 12, 13], dtypes.int64) - values = constant_op.constant([0, 1, 2], dtypes.int64) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=-1, - empty_key=0, - deleted_key=-1, - initial_num_buckets=4) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - self.assertAllEqual(4, len(table.export()[0].eval())) - - keys2 = constant_op.constant([12, 99], dtypes.int64) - table.remove(keys2).run() - self.assertAllEqual(2, table.size().eval()) - self.assertAllEqual(4, len(table.export()[0].eval())) - - keys3 = constant_op.constant([13, 14, 15, 16, 17], dtypes.int64) - values3 = constant_op.constant([3, 4, 5, 6, 7], dtypes.int64) - - table.insert(keys3, values3).run() - self.assertAllEqual(6, table.size().eval()) - self.assertAllEqual(16, len(table.export()[0].eval())) - - keys4 = constant_op.constant([10, 11, 12, 13, 14, 15, 16, 17, 18], - dtypes.int64) - output = table.lookup(keys4) - self.assertAllEqual([-1, 0, -1, 3, 4, 5, 6, 7, -1], output.eval()) - - def testExport(self): - with self.cached_session(): - - keys = constant_op.constant([11, 12, 13, 14], dtypes.int64) - values = constant_op.constant([1, 2, 3, 4], dtypes.int64) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=-1, - empty_key=100, - deleted_key=200, - initial_num_buckets=8) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(4, table.size().eval()) - - keys2 = constant_op.constant([12, 15], dtypes.int64) - table.remove(keys2).run() - self.assertAllEqual(3, table.size().eval()) - - exported_keys, exported_values = table.export() - self.assertAllEqual([None], exported_keys.get_shape().as_list()) - self.assertAllEqual([None], exported_values.get_shape().as_list()) - - np_keys = exported_keys.eval() - np_values = exported_values.eval() - - self.assertAllEqual(8, len(np_keys)) - self.assertAllEqual(8, len(np_values)) - - # pair up keys and values, drop extra added dimension - pairs = np.dstack((np_keys.flatten(), np_values.flatten()))[0] - # sort by key - pairs = pairs[pairs[:, 0].argsort()] - self.assertAllEqual([[11, 1], [13, 3], [14, 4], [100, 0], [100, 0], - [100, 0], [100, 0], [200, 2]], pairs) - - def testSaveRestore(self): - save_dir = os.path.join(self.get_temp_dir(), "save_restore") - save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash") - - with self.session(graph=ops.Graph()) as sess: - default_value = -1 - empty_key = 0 - deleted_key = -1 - keys = constant_op.constant([11, 12, 13, 14], dtypes.int64) - values = constant_op.constant([0, 1, 2, 3], dtypes.int64) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=default_value, - empty_key=empty_key, - deleted_key=deleted_key, - name="t1", - checkpoint=True, - initial_num_buckets=32) - - save = saver.Saver() - - self.assertAllEqual(0, table.size().eval()) - table.insert(keys, values).run() - self.assertAllEqual(4, table.size().eval()) - self.assertAllEqual(32, len(table.export()[0].eval())) - - keys2 = constant_op.constant([12, 15], dtypes.int64) - table.remove(keys2).run() - self.assertAllEqual(3, table.size().eval()) - self.assertAllEqual(32, len(table.export()[0].eval())) - - val = save.save(sess, save_path) - self.assertTrue(isinstance(val, six.string_types)) - self.assertEqual(save_path, val) - - with self.session(graph=ops.Graph()) as sess: - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=default_value, - empty_key=empty_key, - deleted_key=deleted_key, - name="t1", - checkpoint=True, - initial_num_buckets=64) - table.insert( - constant_op.constant([11, 14], dtypes.int64), - constant_op.constant([12, 24], dtypes.int64)).run() - self.assertAllEqual(2, table.size().eval()) - self.assertAllEqual(64, len(table.export()[0].eval())) - - save = saver.Saver() - - # Restore the saved values in the parameter nodes. - save.restore(sess, save_path) - - self.assertAllEqual(3, table.size().eval()) - self.assertAllEqual(32, len(table.export()[0].eval())) - - input_string = constant_op.constant([10, 11, 12, 13, 14], dtypes.int64) - output = table.lookup(input_string) - self.assertAllEqual([-1, 0, -1, 2, 3], output.eval()) - - @test_util.run_in_graph_and_eager_modes - def testObjectSaveRestore(self): - save_dir = os.path.join(self.get_temp_dir(), "save_restore") - save_prefix = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash") - - default_value = -1 - empty_key = 0 - deleted_key = -1 - keys = constant_op.constant([11, 12, 13], dtypes.int64) - values = constant_op.constant([0, 1, 2], dtypes.int64) - save_table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=default_value, - empty_key=empty_key, - deleted_key=deleted_key, - name="t1", - checkpoint=True, - initial_num_buckets=32) - - save_checkpoint = checkpointable.Checkpoint(table=save_table) - - self.assertAllEqual(0, self.evaluate(save_table.size())) - self.evaluate(save_table.insert(keys, values)) - self.assertAllEqual(3, self.evaluate(save_table.size())) - self.assertAllEqual(32, len(self.evaluate(save_table.export()[0]))) - - save_path = save_checkpoint.save(save_prefix) - del save_table, save_checkpoint - - load_table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=default_value, - empty_key=empty_key, - deleted_key=deleted_key, - name="t1", - checkpoint=True, - initial_num_buckets=64) - self.evaluate(load_table.insert( - constant_op.constant([11, 14], dtypes.int64), - constant_op.constant([12, 24], dtypes.int64))) - self.assertAllEqual(2, self.evaluate(load_table.size())) - self.assertAllEqual(64, len(self.evaluate(load_table.export()[0]))) - - restore_checkpoint = checkpointable.Checkpoint(table=load_table) - - # Restore the saved values in the parameter nodes. - restore_checkpoint.restore(save_path).run_restore_ops() - - self.assertAllEqual(3, self.evaluate(load_table.size())) - self.assertAllEqual(32, len(self.evaluate(load_table.export()[0]))) - - input_string = constant_op.constant([10, 11, 12, 13, 14], dtypes.int64) - output = load_table.lookup(input_string) - self.assertAllEqual([-1, 0, 1, 2, -1], self.evaluate(output)) - - def testVectorSaveRestore(self): - save_dir = os.path.join(self.get_temp_dir(), "vector_save_restore") - save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash") - - with self.session(graph=ops.Graph()) as sess: - empty_key = constant_op.constant([11, 13], dtypes.int64) - deleted_key = constant_op.constant([-2, -3], dtypes.int64) - default_value = constant_op.constant([-1, -2], dtypes.int64) - keys = constant_op.constant([[11, 12], [11, 14], [12, 13], [13, 14]], - dtypes.int64) - values = constant_op.constant([[0, 1], [2, 3], [2, 4], [4, 5]], - dtypes.int64) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=default_value, - empty_key=empty_key, - deleted_key=deleted_key, - name="t1", - checkpoint=True, - initial_num_buckets=32) - - save = saver.Saver() - - self.assertAllEqual(0, table.size().eval()) - table.insert(keys, values).run() - self.assertAllEqual(4, table.size().eval()) - self.assertAllEqual(32, len(table.export()[0].eval())) - - keys2 = constant_op.constant([[12, 13], [16, 17]], dtypes.int64) - table.remove(keys2).run() - self.assertAllEqual(3, table.size().eval()) - self.assertAllEqual(32, len(table.export()[0].eval())) - - val = save.save(sess, save_path) - self.assertTrue(isinstance(val, six.string_types)) - self.assertEqual(save_path, val) - - with self.session(graph=ops.Graph()) as sess: - empty_key = constant_op.constant([11, 13], dtypes.int64) - deleted_key = constant_op.constant([-2, -3], dtypes.int64) - default_value = constant_op.constant([-1, -2], dtypes.int64) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=default_value, - empty_key=empty_key, - deleted_key=deleted_key, - name="t1", - checkpoint=True, - initial_num_buckets=64) - table.insert( - constant_op.constant([[11, 12], [13, 15]], dtypes.int64), - constant_op.constant([[21, 22], [23, 24]], dtypes.int64)).run() - self.assertAllEqual(2, table.size().eval()) - self.assertAllEqual(64, len(table.export()[0].eval())) - - save = saver.Saver() - - # Restore the saved values in the parameter nodes. - save.restore(sess, save_path) - - self.assertAllEqual(3, table.size().eval()) - self.assertAllEqual(32, len(table.export()[0].eval())) - - input_string = constant_op.constant( - [[11, 12], [11, 14], [11, 15], [13, 14], [13, 15]], dtypes.int64) - output = table.lookup(input_string) - self.assertAllEqual([[0, 1], [2, 3], [-1, -2], [4, 5], [-1, -2]], - output.eval()) - - def testVectorScalarSaveRestore(self): - save_dir = os.path.join(self.get_temp_dir(), "vector_scalar_save_restore") - save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash") - - with self.session(graph=ops.Graph()) as sess: - empty_key = constant_op.constant([11, 13], dtypes.int64) - deleted_key = constant_op.constant([-1, -1], dtypes.int64) - default_value = constant_op.constant(-1, dtypes.int64) - keys = constant_op.constant([[11, 12], [11, 14], [12, 13], [13, 14]], - dtypes.int64) - values = constant_op.constant([0, 1, 2, 3], dtypes.int64) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=default_value, - empty_key=empty_key, - deleted_key=deleted_key, - name="t2", - checkpoint=True, - initial_num_buckets=32) - - save = saver.Saver() - - self.assertAllEqual(0, table.size().eval()) - table.insert(keys, values).run() - self.assertAllEqual(4, table.size().eval()) - self.assertAllEqual(32, len(table.export()[0].eval())) - - keys2 = constant_op.constant([[12, 13], [15, 16]], dtypes.int64) - table.remove(keys2).run() - self.assertAllEqual(3, table.size().eval()) - self.assertAllEqual(32, len(table.export()[0].eval())) - - val = save.save(sess, save_path) - self.assertTrue(isinstance(val, six.string_types)) - self.assertEqual(save_path, val) - - with self.session(graph=ops.Graph()) as sess: - empty_key = constant_op.constant([11, 13], dtypes.int64) - deleted_key = constant_op.constant([-1, -1], dtypes.int64) - default_value = constant_op.constant(-1, dtypes.int64) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=default_value, - empty_key=empty_key, - deleted_key=deleted_key, - name="t2", - checkpoint=True, - initial_num_buckets=64) - table.insert( - constant_op.constant([[11, 12], [13, 15]], dtypes.int64), - constant_op.constant([3, 4], dtypes.int64)).run() - self.assertAllEqual(2, table.size().eval()) - self.assertAllEqual(64, len(table.export()[0].eval())) - - save = saver.Saver() - - # Restore the saved values in the parameter nodes. - save.restore(sess, save_path) - - self.assertAllEqual(3, table.size().eval()) - self.assertAllEqual(32, len(table.export()[0].eval())) - - input_string = constant_op.constant( - [[11, 12], [11, 14], [11, 15], [13, 14], [13, 15]], dtypes.int64) - output = table.lookup(input_string) - self.assertAllEqual([0, 1, -1, 3, -1], output.eval()) - - def testReprobe(self): - with self.cached_session(): - # Insert 6 keys into a table with 8 buckets. - # The values are chosen to make sure collisions occur when using GCC STL - keys = constant_op.constant([11, 12, 13, 19, 20, 21], dtypes.int64) - values = constant_op.constant([51, 52, 53, 54, 55, 56], dtypes.int64) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=-1, - empty_key=0, - deleted_key=-1, - initial_num_buckets=8) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(6, table.size().eval()) - - input_string = constant_op.constant([10, 11, 12, 13, 14, 19, 20, 21, 22], - dtypes.int64) - output = table.lookup(input_string) - self.assertAllEqual([9], output.get_shape()) - - result = output.eval() - self.assertAllEqual([-1, 51, 52, 53, -1, 54, 55, 56, -1], result) - - def testCustomEmptyKey(self): - with self.cached_session(): - keys = constant_op.constant([11, 0, 13], dtypes.int64) - values = constant_op.constant([0, 1, 2], dtypes.int64) - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=-1, - empty_key=12, - deleted_key=-1) - self.assertAllEqual(0, table.size().eval()) - - table.insert(keys, values).run() - self.assertAllEqual(3, table.size().eval()) - - input_string = constant_op.constant([11, 0, 15], dtypes.int64) - output = table.lookup(input_string) - self.assertAllEqual([3], output.get_shape()) - - result = output.eval() - self.assertAllEqual([0, 1, -1], result) - - def testErrors(self): - with self.cached_session(): - table = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=-1, - empty_key=0, - deleted_key=-1) - - # Inserting the empty key returns an error - keys1 = constant_op.constant([11, 0], dtypes.int64) - values1 = constant_op.constant([0, 1], dtypes.int64) - with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, - "empty_key"): - table.insert(keys1, values1).run() - - # Looking up the empty key returns an error - with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, - "empty_key"): - table.lookup(keys1).eval() - - # Inserting the deleted key returns an error - keys2 = constant_op.constant([11, -1], dtypes.int64) - values2 = constant_op.constant([0, 1], dtypes.int64) - with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, - "deleted_key"): - table.insert(keys2, values2).run() - - # Looking up the empty key returns an error - with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, - "deleted_key"): - table.lookup(keys2).eval() - - # Arbitrary tensors of keys are not supported - keys = constant_op.constant([[11, 0], [12, 1]], dtypes.int64) - values = constant_op.constant([[11, 0], [12, 1]], dtypes.int64) - with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, - "Expected key shape"): - table.lookup(keys).eval() - with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, - "Expected key shape"): - table.insert(keys, values).run() - - table2 = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=-1, - empty_key=17, - deleted_key=-1, - initial_num_buckets=12) - with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, - "Number of buckets must be"): - self.assertAllEqual(0, table2.size().eval()) - - with self.assertRaisesRegexp( - errors_impl.InvalidArgumentError, - "Empty and deleted keys must have same shape"): - table3 = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=-1, - empty_key=42, - deleted_key=[1, 2]) - self.assertAllEqual(0, table3.size().eval()) - - with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, - "Empty and deleted keys cannot be equal"): - table4 = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=-1, - empty_key=42, - deleted_key=42) - self.assertAllEqual(0, table4.size().eval()) - - with self.assertRaisesRegexp(errors_impl.InvalidArgumentError, - "Empty and deleted keys cannot be equal"): - table5 = lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.int64, - default_value=-1, - empty_key=[1, 2, 3], - deleted_key=[1, 2, 3]) - self.assertAllEqual(0, table5.size().eval()) - - class IndexTableFromFile(test.TestCase): def _createVocabFile(self, basename, values=("brain", "salad", "surgery")): @@ -2721,64 +1481,5 @@ class IdTableWithHashBucketsTest(test.TestCase): hasher_spec=lookup.StrongHashSpec([None, 2])) -class MutableHashTableBenchmark(test.Benchmark): - - def _create_table(self): - return lookup.MutableHashTable(dtypes.int64, dtypes.float32, 0.0) - - def benchmark_single_repeated_scalar_insert_scalar(self): - table = self._create_table() - value = variables.Variable(1.0) - insert = table.insert(0, value) - size = table.size() - with session.Session() as sess: - sess.run(value.initializer) - self.run_op_benchmark(sess, insert, burn_iters=10, min_iters=10000) - assert sess.run(size) == 1 - - def benchmark_many_repeated_scalar_insert_scalar(self): - table = self._create_table() - c = dataset_ops.make_one_shot_iterator(counter.Counter()).get_next() - value = variables.Variable(1.0) - insert = table.insert(c, value) - size = table.size() - with session.Session() as sess: - sess.run(value.initializer) - self.run_op_benchmark(sess, insert, burn_iters=10, min_iters=10000) - assert sess.run(size) >= 10000 - - def benchmark_single_repeated_batch_32_insert_scalar(self): - table = self._create_table() - value = variables.Variable([1.0] * 32) - insert = table.insert(list(range(32)), value) - size = table.size() - with session.Session() as sess: - sess.run(value.initializer) - self.run_op_benchmark(sess, insert, burn_iters=10, min_iters=1000) - assert sess.run(size) == 32 - - def benchmark_many_repeated_batch_32_insert_scalar(self): - table = self._create_table() - c = dataset_ops.make_one_shot_iterator(counter.Counter()).get_next() - value = variables.Variable([1.0] * 32) - insert = table.insert(32 * c + list(range(32)), value) - size = table.size() - with session.Session() as sess: - sess.run(value.initializer) - self.run_op_benchmark(sess, insert, burn_iters=10, min_iters=1000) - assert sess.run(size) >= 1000*32 - - -class MutableDenseHashTableBenchmark(MutableHashTableBenchmark): - - def _create_table(self): - return lookup.MutableDenseHashTable( - dtypes.int64, - dtypes.float32, - default_value=0.0, - empty_key=-1, - deleted_key=-2) - - if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/losses/BUILD b/tensorflow/contrib/losses/BUILD index 728f75f8ef1eb3b107dbd0ab4ffbecd63787bf3e..f4ebbdeee883ddeef0d47cb561901c16e2195bb2 100644 --- a/tensorflow/contrib/losses/BUILD +++ b/tensorflow/contrib/losses/BUILD @@ -82,10 +82,11 @@ py_library( py_test( name = "metric_loss_ops_test", - size = "large", + size = "medium", srcs = [ "python/metric_learning/metric_loss_ops_test.py", ], + shard_count = 4, srcs_version = "PY2AND3", deps = [ ":metric_learning_py", diff --git a/tensorflow/contrib/losses/python/losses/loss_ops.py b/tensorflow/contrib/losses/python/losses/loss_ops.py index 709a042bbcefb89125f7e4cd14a0d7ecd2b53281..5ebdd0b8b50063c99e6b747c594eb99c306b4efb 100644 --- a/tensorflow/contrib/losses/python/losses/loss_ops.py +++ b/tensorflow/contrib/losses/python/losses/loss_ops.py @@ -511,7 +511,7 @@ def mean_squared_error(predictions, labels=None, weights=1.0, scope=None): predictions.get_shape().assert_is_compatible_with(labels.get_shape()) predictions = math_ops.to_float(predictions) labels = math_ops.to_float(labels) - losses = math_ops.square(math_ops.subtract(predictions, labels)) + losses = math_ops.squared_difference(predictions, labels) return compute_weighted_loss(losses, weights, scope=scope) diff --git a/tensorflow/contrib/losses/python/metric_learning/metric_loss_ops.py b/tensorflow/contrib/losses/python/metric_learning/metric_loss_ops.py index de76acb51ffe985162a66c617b266f47c5216b19..f3b0e77740ff1d940fcd6d00b3482e90f6ebf952 100644 --- a/tensorflow/contrib/losses/python/metric_learning/metric_loss_ops.py +++ b/tensorflow/contrib/losses/python/metric_learning/metric_loss_ops.py @@ -105,7 +105,8 @@ def contrastive_loss(labels, embeddings_anchor, embeddings_positive, # Get per pair distances distances = math_ops.sqrt( math_ops.reduce_sum( - math_ops.square(embeddings_anchor - embeddings_positive), 1)) + math_ops.squared_difference(embeddings_anchor, embeddings_positive), + 1)) # Add contrastive loss for the siamese network. # label here is {0,1} for neg, pos. diff --git a/tensorflow/contrib/makefile/download_dependencies.sh b/tensorflow/contrib/makefile/download_dependencies.sh index 2a5232b476712a96f84be0f4725beb78bc138297..af3c541dc214c30e9e59fdcca995ffc53b028df4 100755 --- a/tensorflow/contrib/makefile/download_dependencies.sh +++ b/tensorflow/contrib/makefile/download_dependencies.sh @@ -142,5 +142,6 @@ replace_by_sed 's#static uint64x2_t p2ul_CONJ_XOR = vld1q_u64( p2ul_conj_XOR_DAT # TODO(satok): Remove this once protobuf/autogen.sh is fixed. replace_by_sed 's#https://googlemock.googlecode.com/files/gmock-1.7.0.zip#http://download.tensorflow.org/deps/gmock-1.7.0.zip#' \ "${DOWNLOADS_DIR}/protobuf/autogen.sh" +cat "third_party/eigen3/gebp_neon.patch" | patch "${DOWNLOADS_DIR}/eigen/Eigen/src/Core/products/GeneralBlockPanelKernel.h" echo "download_dependencies.sh completed successfully." >&2 diff --git a/tensorflow/contrib/makefile/proto_text_pb_cc_files.txt b/tensorflow/contrib/makefile/proto_text_pb_cc_files.txt index 87c73ec1ca610cac6d63468887bc350bada5910b..8330c45cc16ffa536107e25699379bb5d9e8993b 100644 --- a/tensorflow/contrib/makefile/proto_text_pb_cc_files.txt +++ b/tensorflow/contrib/makefile/proto_text_pb_cc_files.txt @@ -36,6 +36,7 @@ tensorflow/core/protobuf/queue_runner.pb.cc tensorflow/core/protobuf/rewriter_config.pb.cc tensorflow/core/protobuf/saver.pb.cc tensorflow/core/protobuf/tensorflow_server.pb.cc +tensorflow/core/protobuf/verifier_config.pb.cc tensorflow/core/util/event.pb.cc tensorflow/core/util/memmapped_file_system.pb.cc tensorflow/core/util/saved_tensor_slice.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 4120ea52ec5255b1efce7a6ce6890fc79c1e4831..7257ac8feedfb8ed18c4d691cd85766e70a48ae8 100644 --- a/tensorflow/contrib/makefile/proto_text_pb_h_files.txt +++ b/tensorflow/contrib/makefile/proto_text_pb_h_files.txt @@ -37,6 +37,7 @@ tensorflow/core/protobuf/rewriter_config.pb.h tensorflow/core/protobuf/saver.pb.h tensorflow/core/protobuf/tensor_bundle.pb.h tensorflow/core/protobuf/tensorflow_server.pb.h +tensorflow/core/protobuf/verifier_config.pb.h tensorflow/core/util/event.pb.h tensorflow/core/util/memmapped_file_system.pb.h tensorflow/core/util/saved_tensor_slice.pb.h diff --git a/tensorflow/contrib/makefile/tf_pb_text_files.txt b/tensorflow/contrib/makefile/tf_pb_text_files.txt index f94d70db9046cec43073ab1406762aea1f28c8e3..13e3b6422d1989b0d499d8d20901d919554c630e 100644 --- a/tensorflow/contrib/makefile/tf_pb_text_files.txt +++ b/tensorflow/contrib/makefile/tf_pb_text_files.txt @@ -29,5 +29,6 @@ tensorflow/core/protobuf/debug.pb_text.cc tensorflow/core/protobuf/rewriter_config.pb_text.cc tensorflow/core/protobuf/saver.pb_text.cc tensorflow/core/protobuf/tensor_bundle.pb_text.cc +tensorflow/core/protobuf/verifier_config.pb_text.cc tensorflow/core/util/memmapped_file_system.pb_text.cc tensorflow/core/util/saved_tensor_slice.pb_text.cc diff --git a/tensorflow/contrib/makefile/tf_proto_files.txt b/tensorflow/contrib/makefile/tf_proto_files.txt index 2712e906d719e72dacb60e213205ad68895f905f..24d86d313b76343ed9450a33cf185d9c426696bb 100644 --- a/tensorflow/contrib/makefile/tf_proto_files.txt +++ b/tensorflow/contrib/makefile/tf_proto_files.txt @@ -43,6 +43,7 @@ tensorflow/core/protobuf/rewriter_config.proto tensorflow/core/protobuf/saver.proto tensorflow/core/protobuf/tensor_bundle.proto tensorflow/core/protobuf/tensorflow_server.proto +tensorflow/core/protobuf/verifier_config.proto tensorflow/core/util/event.proto tensorflow/core/util/memmapped_file_system.proto tensorflow/core/util/saved_tensor_slice.proto diff --git a/tensorflow/contrib/metrics/python/ops/metric_ops.py b/tensorflow/contrib/metrics/python/ops/metric_ops.py index 7b432f8bd20989c6d95310bcaca88d44ce3e0d1f..ece246b7c28569a551f7733daf16ee1507f9c95d 100644 --- a/tensorflow/contrib/metrics/python/ops/metric_ops.py +++ b/tensorflow/contrib/metrics/python/ops/metric_ops.py @@ -1356,9 +1356,8 @@ def _compute_placement_auc(labels, predictions, weights, alpha, weights_0 * math_ops.square(1. - placement_values_0 - auc_0)) / (total_0 - 1. + _EPSILON)) var_1 = ( - math_ops.reduce_sum( - weights_1 * math_ops.square(placement_values_1 - auc_1)) / - (total_1 - 1. + _EPSILON)) + math_ops.reduce_sum(weights_1 * math_ops.squared_difference( + placement_values_1, auc_1)) / (total_1 - 1. + _EPSILON)) auc_std_err = math_ops.sqrt( (var_0 / (total_0 + _EPSILON)) + (var_1 / (total_1 + _EPSILON))) diff --git a/tensorflow/contrib/mpi/mpi_server_lib.cc b/tensorflow/contrib/mpi/mpi_server_lib.cc index a31fa9ce0b3110d875689d74a41ca9f9cc85f532..e44e10af0814ba8d6d964dfc34a0470ce45c0b40 100644 --- a/tensorflow/contrib/mpi/mpi_server_lib.cc +++ b/tensorflow/contrib/mpi/mpi_server_lib.cc @@ -54,7 +54,10 @@ MPIServer::~MPIServer() { Status MPIServer::Init(ServiceInitFunction service_func, RendezvousMgrCreationFunction rendezvous_mgr_func) { - Status s = GrpcServer::Init(service_func, rendezvous_mgr_func); + GrpcServerOptions opts; + opts.service_func = service_func; + opts.rendezvous_mgr_func = rendezvous_mgr_func; + Status s = GrpcServer::Init(opts); return s; } diff --git a/tensorflow/contrib/mpi_collectives/BUILD b/tensorflow/contrib/mpi_collectives/BUILD index ecac06354d2ce796f2a6021cdf2370d7c30ccab7..a7be92a35e0d62a61f7923ac61bb2c1267d039c6 100644 --- a/tensorflow/contrib/mpi_collectives/BUILD +++ b/tensorflow/contrib/mpi_collectives/BUILD @@ -52,7 +52,6 @@ tf_custom_op_library( deps = [ ":mpi_defines", ":mpi_message_proto_cc", - "//tensorflow/stream_executor:stream_executor_headers_lib", "//third_party/mpi", ], ) diff --git a/tensorflow/contrib/mpi_collectives/ring.h b/tensorflow/contrib/mpi_collectives/ring.h index cae57ce60eb09509af69f8ccab9eacedea361548..9b5d52e1b648e62af93d5420885e4f22796e3ea1 100644 --- a/tensorflow/contrib/mpi_collectives/ring.h +++ b/tensorflow/contrib/mpi_collectives/ring.h @@ -129,7 +129,7 @@ cudaStream_t CudaStreamForMPI(); * has the fully accumulated Segment 1; and so on. The scatter-reduce is * complete. * - * Next, the allgather distributes these fully accumululated chunks across all + * Next, the allgather distributes these fully accumulated chunks across all * nodes. Communication proceeds in the same ring, once again in N-1 steps. At * the ith step, node j will send chunk (j - i + 1) and receive chunk (j - i). * For example, at the first iteration, the following transfers will occur: diff --git a/tensorflow/contrib/opt/BUILD b/tensorflow/contrib/opt/BUILD index 0446e823d95f8ecbed6a0c34a83ade009e68448b..f30643cf3059754daaeee4093938ac47b26f76ea 100644 --- a/tensorflow/contrib/opt/BUILD +++ b/tensorflow/contrib/opt/BUILD @@ -319,6 +319,9 @@ tf_py_test( "//tensorflow/python:framework_for_generated_wrappers", "//third_party/py/numpy", ], + tags = [ + "oss_serial", + ], ) tf_py_test( @@ -410,8 +413,9 @@ py_test( py_test( name = "shampoo_test", - size = "large", + size = "medium", srcs = ["python/training/shampoo_test.py"], + shard_count = 4, srcs_version = "PY2AND3", deps = [ ":opt_py", diff --git a/tensorflow/contrib/opt/python/training/adam_gs_optimizer.py b/tensorflow/contrib/opt/python/training/adam_gs_optimizer.py index 3fb649ea82e79b3bc78a2da6d5c3e9a071adec6d..0b149ed17533adff3bd7cd8fd8ff94d171f72911 100644 --- a/tensorflow/contrib/opt/python/training/adam_gs_optimizer.py +++ b/tensorflow/contrib/opt/python/training/adam_gs_optimizer.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== - """Adam rewrite to use global step for computing beta1 & beta2 accumulation.""" from __future__ import absolute_import from __future__ import division @@ -38,10 +37,15 @@ class AdamGSOptimizer(optimizer.Optimizer): ([pdf](http://arxiv.org/pdf/1412.6980.pdf)). """ - def __init__(self, global_step=0, learning_rate=0.001, - beta1=0.9, beta2=0.999, epsilon=1e-8, - use_locking=False, name="Adam"): - """Construct a new Adam optimizer. + def __init__(self, + global_step=0, + learning_rate=0.001, + beta1=0.9, + beta2=0.999, + epsilon=1e-8, + use_locking=False, + name="Adam"): + r"""Construct a new Adam optimizer. Branched from tf.train.AdamOptimizer. The only difference is to pass global step for computing beta1 and beta2 accumulators, instead of having @@ -83,23 +87,20 @@ class AdamGSOptimizer(optimizer.Optimizer): Args: global_step: tensorflow variable indicating the step. learning_rate: A Tensor or a floating point value. The learning rate. - beta1: A float value or a constant float tensor. - The exponential decay rate for the 1st moment estimates. - beta2: A float value or a constant float tensor. - The exponential decay rate for the 2nd moment estimates. + beta1: A float value or a constant float tensor. The exponential decay + rate for the 1st moment estimates. + beta2: A float value or a constant float tensor. The exponential decay + rate for the 2nd moment estimates. epsilon: A small constant for numerical stability. This epsilon is "epsilon hat" in the Kingma and Ba paper (in the formula just before Section 2.1), not the epsilon in Algorithm 1 of the paper. use_locking: If True use locks for update operations. name: Optional name for the operations created when applying gradients. - Defaults to "Adam". - - @compatibility(eager) - When eager execution is enabled, `learning_rate`, `beta1`, `beta2`, and - `epsilon` can each be a callable that takes no arguments and returns the - actual value to use. This can be useful for changing these values across - different invocations of optimizer functions. - @end_compatibility + Defaults to "Adam". @compatibility(eager) When eager execution is + enabled, `learning_rate`, `beta1`, `beta2`, and `epsilon` can each be a + callable that takes no arguments and returns the actual value to use. + This can be useful for changing these values across different + invocations of optimizer functions. @end_compatibility """ super(AdamGSOptimizer, self).__init__(use_locking, name) self._lr = learning_rate @@ -115,9 +116,6 @@ class AdamGSOptimizer(optimizer.Optimizer): self._beta2_t = None self._epsilon_t = None - # Created in SparseApply if needed. - self._updated_lr = None - def _get_beta_accumulators(self): return (math_ops.pow(self._beta1_t, self._global_step_on_worker), math_ops.pow(self._beta2_t, self._global_step_on_worker)) @@ -149,28 +147,34 @@ class AdamGSOptimizer(optimizer.Optimizer): v = self.get_slot(var, "v") beta1_power, beta2_power = self._get_beta_accumulators() return training_ops.apply_adam( - var, m, v, + var, + m, + v, math_ops.cast(beta1_power, var.dtype.base_dtype), math_ops.cast(beta2_power, var.dtype.base_dtype), math_ops.cast(self._lr_t, var.dtype.base_dtype), math_ops.cast(self._beta1_t, var.dtype.base_dtype), math_ops.cast(self._beta2_t, var.dtype.base_dtype), math_ops.cast(self._epsilon_t, var.dtype.base_dtype), - grad, use_locking=self._use_locking).op + grad, + use_locking=self._use_locking).op def _resource_apply_dense(self, grad, var): m = self.get_slot(var, "m") v = self.get_slot(var, "v") beta1_power, beta2_power = self._get_beta_accumulators() return training_ops.resource_apply_adam( - var.handle, m.handle, v.handle, + var.handle, + m.handle, + v.handle, math_ops.cast(beta1_power, grad.dtype.base_dtype), math_ops.cast(beta2_power, grad.dtype.base_dtype), math_ops.cast(self._lr_t, grad.dtype.base_dtype), math_ops.cast(self._beta1_t, grad.dtype.base_dtype), math_ops.cast(self._beta2_t, grad.dtype.base_dtype), math_ops.cast(self._epsilon_t, grad.dtype.base_dtype), - grad, use_locking=self._use_locking) + grad, + use_locking=self._use_locking) def _apply_sparse_shared(self, grad, var, indices, scatter_add): beta1_power, beta2_power = self._get_beta_accumulators() @@ -184,8 +188,7 @@ class AdamGSOptimizer(optimizer.Optimizer): # m_t = beta1 * m + (1 - beta1) * g_t m = self.get_slot(var, "m") m_scaled_g_values = grad * (1 - beta1_t) - m_t = state_ops.assign(m, m * beta1_t, - use_locking=self._use_locking) + m_t = state_ops.assign(m, m * beta1_t, use_locking=self._use_locking) with ops.control_dependencies([m_t]): m_t = scatter_add(m, indices, m_scaled_g_values) # v_t = beta2 * v + (1 - beta2) * (g_t * g_t) @@ -195,23 +198,26 @@ class AdamGSOptimizer(optimizer.Optimizer): with ops.control_dependencies([v_t]): v_t = scatter_add(v, indices, v_scaled_g_values) v_sqrt = math_ops.sqrt(v_t) - var_update = state_ops.assign_sub(var, - lr * m_t / (v_sqrt + epsilon_t), - use_locking=self._use_locking) + var_update = state_ops.assign_sub( + var, lr * m_t / (v_sqrt + epsilon_t), use_locking=self._use_locking) return control_flow_ops.group(*[var_update, m_t, v_t]) def _apply_sparse(self, grad, var): return self._apply_sparse_shared( - grad.values, var, grad.indices, + grad.values, + var, + grad.indices, lambda x, i, v: state_ops.scatter_add( # pylint: disable=g-long-lambda - x, i, v, use_locking=self._use_locking)) + x, + i, + v, + use_locking=self._use_locking)) def _resource_scatter_add(self, x, i, v): with ops.control_dependencies( - [resource_variable_ops.resource_scatter_add( - x.handle, i, v)]): + [resource_variable_ops.resource_scatter_add(x.handle, i, v)]): return x.value() def _resource_apply_sparse(self, grad, var, indices): - return self._apply_sparse_shared( - grad, var, indices, self._resource_scatter_add) + return self._apply_sparse_shared(grad, var, indices, + self._resource_scatter_add) diff --git a/tensorflow/contrib/optimizer_v2/adam.py b/tensorflow/contrib/optimizer_v2/adam.py index 248ffb1f7eb5dc27112ddf9b8670344904065ed0..1b7800f324b908e3c88fe90d31a2a08cbbd5ccf2 100644 --- a/tensorflow/contrib/optimizer_v2/adam.py +++ b/tensorflow/contrib/optimizer_v2/adam.py @@ -36,7 +36,7 @@ class AdamOptimizer(optimizer_v2.OptimizerV2): def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8, use_locking=False, name="Adam"): - """Construct a new Adam optimizer. + r"""Construct a new Adam optimizer. Initialization: diff --git a/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py b/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py index 72019b31540a943582ebb4699013d9dcfc10769f..b5de726a4cf833d23d668968b8080e7e484cd496 100644 --- a/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py +++ b/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py @@ -44,11 +44,12 @@ from tensorflow.python.ops import variable_scope from tensorflow.python.training import checkpoint_management from tensorflow.python.training import saver as core_saver from tensorflow.python.training import training_util +from tensorflow.python.training.checkpointable import graph_view from tensorflow.python.training.checkpointable import tracking from tensorflow.python.training.checkpointable import util -class NonLayerCheckpointable(tracking.Checkpointable): +class NonLayerCheckpointable(tracking.AutoCheckpointable): def __init__(self): super(NonLayerCheckpointable, self).__init__() @@ -118,9 +119,8 @@ class CheckpointingTests(test.TestCase): self.evaluate(util.gather_initializers( root_checkpointable)) self.evaluate(train_op) - named_variables, serialized_graph, _ = ( - util._serialize_object_graph( - root_checkpointable, saveables_cache=None)) + named_variables, serialized_graph, _ = graph_view.ObjectGraphView( + root_checkpointable).serialize_object_graph() expected_checkpoint_names = ( # Created in the root node, so no prefix. "optimizer_step", @@ -440,7 +440,7 @@ class CheckpointingTests(test.TestCase): def testDeferredSlotRestoration(self): checkpoint_directory = self.get_temp_dir() - root = tracking.Checkpointable() + root = util.Checkpoint() root.var = util.add_variable( root, name="var", initializer=0.) optimizer = adam.AdamOptimizer(0.1) @@ -455,21 +455,17 @@ class CheckpointingTests(test.TestCase): util.Checkpoint(root=root, optimizer=optimizer))) self.evaluate(train_op) self.evaluate(state_ops.assign(root.var, 12.)) - no_slots_path = util.CheckpointableSaver(root).save( - os.path.join(checkpoint_directory, "no_slots")) + no_slots_path = root.save(os.path.join(checkpoint_directory, "no_slots")) root.optimizer = optimizer self.evaluate(state_ops.assign(root.var, 13.)) self.evaluate(state_ops.assign(optimizer.get_slot(name="m", var=root.var), 14.)) - slots_path = util.CheckpointableSaver(root).save( - os.path.join(checkpoint_directory, "with_slots")) - new_root = tracking.Checkpointable() + slots_path = root.save(os.path.join(checkpoint_directory, "with_slots")) + new_root = util.Checkpoint() # Load the slot-containing checkpoint (deferred), then immediately overwrite # the non-slot variable (also deferred). - slot_status = util.CheckpointableSaver( - new_root).restore(slots_path) - no_slot_status = util.CheckpointableSaver( - new_root).restore(no_slots_path) + slot_status = new_root.restore(slots_path) + no_slot_status = new_root.restore(no_slots_path) with self.assertRaises(AssertionError): no_slot_status.assert_consumed() new_root.var = util.add_variable( @@ -508,15 +504,14 @@ class CheckpointingTests(test.TestCase): with graph.as_default(), self.session(graph): checkpoint_directory = self.get_temp_dir() checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - obj = tracking.Checkpointable() + obj = util.Checkpoint() obj.var = variable_scope.get_variable(name="v", initializer=0.) obj.opt = adam.AdamOptimizer(0.1) obj.opt.minimize(obj.var.read_value()) self.evaluate(util.gather_initializers(obj)) - saver = util.CheckpointableSaver(obj) - saver.save(checkpoint_prefix) + obj.save(checkpoint_prefix) before_ops = graph.get_operations() - saver.save(checkpoint_prefix) + obj.save(checkpoint_prefix) self.assertEqual(before_ops, graph.get_operations()) def testManyRestoresGraph(self): @@ -526,16 +521,15 @@ class CheckpointingTests(test.TestCase): with graph.as_default(), self.session(graph): checkpoint_directory = self.get_temp_dir() checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - obj = tracking.Checkpointable() + obj = util.Checkpoint() obj.var = variable_scope.get_variable(name="v", initializer=0.) obj.opt = adam.AdamOptimizer(0.1) obj.opt.minimize(obj.var.read_value()) self.evaluate(util.gather_initializers(obj)) - saver = util.CheckpointableSaver(obj) - save_path = saver.save(checkpoint_prefix) - saver.restore(save_path) + save_path = obj.save(checkpoint_prefix) + obj.restore(save_path) before_ops = graph.get_operations() - saver.restore(save_path) + obj.restore(save_path) self.assertEqual(before_ops, graph.get_operations()) def testMultipleGraphsNonSlotVariables(self): @@ -704,7 +698,7 @@ class CheckpointCompatibilityTests(test.TestCase): self._set_sentinels(root) with self.assertRaises(AssertionError): self._check_sentinels(root) - object_saver = util.CheckpointableSaver(root) + object_saver = util.CheckpointableSaver(graph_view.ObjectGraphView(root)) self._set_sentinels(root) status = object_saver.restore(save_path) if context.executing_eagerly(): diff --git a/tensorflow/contrib/optimizer_v2/optimizer_v2.py b/tensorflow/contrib/optimizer_v2/optimizer_v2.py index 7fb23abc38d9dc101204ed83808aebe5a8ef1e78..a49149e592f72b1977aae67078d7f41ca6f03d94 100644 --- a/tensorflow/contrib/optimizer_v2/optimizer_v2.py +++ b/tensorflow/contrib/optimizer_v2/optimizer_v2.py @@ -24,7 +24,6 @@ import abc import six -from tensorflow.python.distribute import distribute_lib from tensorflow.python.distribute import distribution_strategy_context as distribute_ctx from tensorflow.python.distribute import reduce_util as ds_reduce_util from tensorflow.python.eager import backprop @@ -661,7 +660,7 @@ class OptimizerV2(optimizer_v1.Optimizer): name=None, grad_loss=None, stop_gradients=None, - scale_loss_by_num_replicas=None): + scale_loss_by_num_replicas=False): """Add operations to minimize `loss` by updating `var_list`. This method simply combines calls `compute_gradients()` and @@ -685,8 +684,7 @@ class OptimizerV2(optimizer_v1.Optimizer): stop_gradients: Optional. A Tensor or list of tensors not to differentiate through. scale_loss_by_num_replicas: Optional boolean. If true, scale the loss down - by the number of replicas. By default, auto-detects whether this is - needed. + by the number of replicas. DEPRECATED and generally no longer needed. Returns: An Operation that updates the variables in `var_list`. If `global_step` @@ -732,7 +730,7 @@ class OptimizerV2(optimizer_v1.Optimizer): aggregation_method=None, grad_loss=None, stop_gradients=None, - scale_loss_by_num_replicas=None): + scale_loss_by_num_replicas=False): """Compute gradients of `loss` for the variables in `var_list`. This is the first part of `minimize()`. It returns a list @@ -756,8 +754,7 @@ class OptimizerV2(optimizer_v1.Optimizer): stop_gradients: Optional. A Tensor or list of tensors not to differentiate through. scale_loss_by_num_replicas: Optional boolean. If true, scale the loss down - by the number of replicas. By default, auto-detects whether this is - needed. + by the number of replicas. DEPRECATED and generally no longer needed. Returns: A list of (gradient, variable) pairs. Variable is always present, but @@ -781,9 +778,7 @@ class OptimizerV2(optimizer_v1.Optimizer): tape.watch(var_list) loss_value = loss() - # Scale loss for number of replicas (callable-loss case). In this case, - # we have to be careful to call distribute_lib.get_loss_reduction() - # *after* loss() is evaluated, so we know what loss reduction it uses. + # Scale loss for number of replicas (callable-loss case). loss_value = self._scale_loss(loss_value, scale_loss_by_num_replicas) if var_list is None: @@ -839,12 +834,8 @@ class OptimizerV2(optimizer_v1.Optimizer): @staticmethod def _scale_loss(loss_value, scale_loss_by_num_replicas): """Scale loss for the number of replicas.""" - if scale_loss_by_num_replicas is None: - scale_loss_by_num_replicas = ( - distribute_lib.get_loss_reduction() == ds_reduce_util.ReduceOp.MEAN) if scale_loss_by_num_replicas: - num_replicas = \ - distribute_ctx.get_distribution_strategy().num_replicas_in_sync + num_replicas = distribute_ctx.get_strategy().num_replicas_in_sync if num_replicas > 1: loss_value *= 1. / num_replicas return loss_value diff --git a/tensorflow/contrib/optimizer_v2/optimizer_v2_test.py b/tensorflow/contrib/optimizer_v2/optimizer_v2_test.py index dd7f2f44055a2e48e8a48d01c1da3a8e7513255d..2fc0b5ea4de2332ff3bf32f9a12a15eee566d5c4 100644 --- a/tensorflow/contrib/optimizer_v2/optimizer_v2_test.py +++ b/tensorflow/contrib/optimizer_v2/optimizer_v2_test.py @@ -26,7 +26,7 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import clip_ops -from tensorflow.python.ops import gradients_impl +from tensorflow.python.ops import gradients_util from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import state_ops from tensorflow.python.ops import variables @@ -71,7 +71,7 @@ class OptimizerTest(test.TestCase): opt_op = sgd_op.minimize( cost, global_step, [var0, var1], - aggregation_method=gradients_impl.AggregationMethod. + aggregation_method=gradients_util.AggregationMethod. EXPERIMENTAL_ACCUMULATE_N) variables.global_variables_initializer().run() diff --git a/tensorflow/contrib/proto/python/kernel_tests/decode_proto_op_test_base.py b/tensorflow/contrib/proto/python/kernel_tests/decode_proto_op_test_base.py index 17b69c7b35dce130c45ab0aadb28be330b4bfb88..c8524e9871864e0b4fffbd549d1fe347714f072a 100644 --- a/tensorflow/contrib/proto/python/kernel_tests/decode_proto_op_test_base.py +++ b/tensorflow/contrib/proto/python/kernel_tests/decode_proto_op_test_base.py @@ -84,7 +84,10 @@ class DecodeProtoOpTestBase(test_base.ProtoOpTestBase, parameterized.TestCase): values = field_dict[field.name] self.assertEqual(dtypes.as_dtype(values.dtype), field.dtype) - fd = field.value.DESCRIPTOR.fields_by_name[field.name] + if 'ext_value' in field.name: + fd = test_example_pb2.PrimitiveValue() + else: + fd = field.value.DESCRIPTOR.fields_by_name[field.name] # Values has the same shape as the input plus an extra # dimension for repeats. @@ -92,13 +95,16 @@ class DecodeProtoOpTestBase(test_base.ProtoOpTestBase, parameterized.TestCase): # Nested messages are represented as TF strings, requiring # some special handling. - if field.name == 'message_value': + if field.name == 'message_value' or 'ext_value' in field.name: vs = [] for buf in values.flat: msg = test_example_pb2.PrimitiveValue() msg.ParseFromString(buf) vs.append(msg) - evs = getattr(field.value, field.name) + if 'ext_value' in field.name: + evs = field.value.Extensions[test_example_pb2.ext_value] + else: + evs = getattr(field.value, field.name) if len(vs) != len(evs): self.fail('Field %s decoded %d outputs, expected %d' % (fd.name, len(vs), len(evs))) @@ -223,7 +229,8 @@ class DecodeProtoOpTestBase(test_base.ProtoOpTestBase, parameterized.TestCase): sanitize=False, force_disordered=True) - @parameterized.named_parameters(*test_base.ProtoOpTestBase.named_parameters()) + @parameterized.named_parameters( + *test_base.ProtoOpTestBase.named_parameters(extension=False)) def testPacked(self, case): # Now try with the packed serialization. # @@ -235,8 +242,7 @@ class DecodeProtoOpTestBase(test_base.ProtoOpTestBase, parameterized.TestCase): # Note: float_format='.17g' is necessary to ensure preservation of # doubles and floats in text format. text_format.Parse( - text_format.MessageToString( - value, float_format='.17g'), + text_format.MessageToString(value, float_format='.17g'), test_example_pb2.PackedTestValue()).SerializeToString() for value in case.values ] diff --git a/tensorflow/contrib/proto/python/kernel_tests/encode_proto_op_test_base.py b/tensorflow/contrib/proto/python/kernel_tests/encode_proto_op_test_base.py index 01b3ccc7fd3918c4ff910281289e31177e5a8097..5ec681ff55dbd18580761bb23e7017cfc9767b89 100644 --- a/tensorflow/contrib/proto/python/kernel_tests/encode_proto_op_test_base.py +++ b/tensorflow/contrib/proto/python/kernel_tests/encode_proto_op_test_base.py @@ -15,9 +15,6 @@ # ============================================================================= """Table-driven test for encode_proto op. -This test is run once with each of the *.TestCase.pbtxt files -in the test directory. - It tests that encode_proto is a lossless inverse of decode_proto (for the specified fields). """ @@ -145,7 +142,8 @@ class EncodeProtoOpTestBase(test_base.ProtoOpTestBase, parameterized.TestCase): # loss of packing in the encoding). self.assertEqual(in_buf, out_buf) - @parameterized.named_parameters(*test_base.ProtoOpTestBase.named_parameters()) + @parameterized.named_parameters( + *test_base.ProtoOpTestBase.named_parameters(extension=False)) def testRoundtrip(self, case): in_bufs = [value.SerializeToString() for value in case.values] @@ -154,7 +152,8 @@ class EncodeProtoOpTestBase(test_base.ProtoOpTestBase, parameterized.TestCase): return self._testRoundtrip( in_bufs, 'tensorflow.contrib.proto.TestValue', case.fields) - @parameterized.named_parameters(*test_base.ProtoOpTestBase.named_parameters()) + @parameterized.named_parameters( + *test_base.ProtoOpTestBase.named_parameters(extension=False)) def testRoundtripPacked(self, case): # Now try with the packed serialization. # We test the packed representations by loading the same test cases using diff --git a/tensorflow/contrib/proto/python/kernel_tests/proto_op_test_base.py b/tensorflow/contrib/proto/python/kernel_tests/proto_op_test_base.py index 2950c7dfdc59a11ba7d2c07d8406bd4af26b5bd9..1a636486a1765ad9544b5cb5e52961cc47f92950 100644 --- a/tensorflow/contrib/proto/python/kernel_tests/proto_op_test_base.py +++ b/tensorflow/contrib/proto/python/kernel_tests/proto_op_test_base.py @@ -38,17 +38,18 @@ class ProtoOpTestBase(test.TestCase): ct.cdll.LoadLibrary(lib) @staticmethod - def named_parameters(): - return ( - ("defaults", ProtoOpTestBase.defaults_test_case()), - ("minmax", ProtoOpTestBase.minmax_test_case()), - ("nested", ProtoOpTestBase.nested_test_case()), - ("optional", ProtoOpTestBase.optional_test_case()), - ("promote", ProtoOpTestBase.promote_test_case()), - ("ragged", ProtoOpTestBase.ragged_test_case()), - ("shaped_batch", ProtoOpTestBase.shaped_batch_test_case()), - ("simple", ProtoOpTestBase.simple_test_case()), - ) + def named_parameters(extension=True): + parameters = [("defaults", ProtoOpTestBase.defaults_test_case()), + ("minmax", ProtoOpTestBase.minmax_test_case()), + ("nested", ProtoOpTestBase.nested_test_case()), + ("optional", ProtoOpTestBase.optional_test_case()), + ("promote", ProtoOpTestBase.promote_test_case()), + ("ragged", ProtoOpTestBase.ragged_test_case()), + ("shaped_batch", ProtoOpTestBase.shaped_batch_test_case()), + ("simple", ProtoOpTestBase.simple_test_case())] + if extension: + parameters.append(("extension", ProtoOpTestBase.extension_test_case())) + return parameters @staticmethod def defaults_test_case(): @@ -399,6 +400,21 @@ class ProtoOpTestBase(test.TestCase): field.value.bool_value.append(True) return test_case + @staticmethod + def extension_test_case(): + test_case = test_example_pb2.TestCase() + value = test_case.values.add() + message_value = value.Extensions[test_example_pb2.ext_value].add() + message_value.double_value = 23.5 + test_case.shapes.append(1) + test_case.sizes.append(1) + field = test_case.fields.add() + field.name = test_example_pb2.ext_value.full_name + field.dtype = types_pb2.DT_STRING + message_value = field.value.Extensions[test_example_pb2.ext_value].add() + message_value.double_value = 23.5 + return test_case + @staticmethod def simple_test_case(): test_case = test_example_pb2.TestCase() diff --git a/tensorflow/contrib/proto/python/kernel_tests/test_example.proto b/tensorflow/contrib/proto/python/kernel_tests/test_example.proto index 674d881220a1113631def47c5111e3ef401b99f3..b1ce66de4feb9c6666ca9ccf39403b4e12840fcf 100644 --- a/tensorflow/contrib/proto/python/kernel_tests/test_example.proto +++ b/tensorflow/contrib/proto/python/kernel_tests/test_example.proto @@ -61,6 +61,8 @@ message TestValue { optional sfixed64 sfixed64_value_with_default = 32 [default = 11]; optional sint32 sint32_value_with_default = 33 [default = 12]; optional sint64 sint64_value_with_default = 34 [default = 13]; + + extensions 100 to 199; } // A PackedTestValue looks exactly the same as a TestValue in the text format, @@ -68,7 +70,7 @@ message TestValue { // by loading the same test cases using this definition instead of TestValue. // // NOTE: This definition must be kept in sync with TestValue in every way except -// the packed=true declaration. +// the packed=true declaration and the lack of extensions. message PackedTestValue { repeated double double_value = 1 [packed = true]; repeated float float_value = 2 [packed = true]; @@ -132,6 +134,10 @@ message ExtraFields { optional bool bool_value = 1777; } +extend TestValue { + repeated PrimitiveValue ext_value = 100; +} + // The messages below are for yet-to-be created tests. message EnumValue { diff --git a/tensorflow/contrib/quantize/BUILD b/tensorflow/contrib/quantize/BUILD index b35c4fde1a2c704880e023a0c3ac1e0766493514..b67e68ea96a15f94e62050c92405eec4fe4be70f 100644 --- a/tensorflow/contrib/quantize/BUILD +++ b/tensorflow/contrib/quantize/BUILD @@ -202,8 +202,9 @@ py_test( py_test( name = "quantize_parameterized_test", - size = "large", + size = "medium", srcs = ["python/quantize_parameterized_test.py"], + shard_count = 4, srcs_version = "PY2AND3", # TODO(b/118839526): Re-enable msan test. tags = [ diff --git a/tensorflow/contrib/quantize/README.md b/tensorflow/contrib/quantize/README.md index 9085d9fa719520ac84ef6f8e07d7fa335bef5605..b335e1af69b7b2e6020f8e745c43bb1bdc95a62d 100644 --- a/tensorflow/contrib/quantize/README.md +++ b/tensorflow/contrib/quantize/README.md @@ -8,9 +8,9 @@ for both training and inference. There are two aspects to this: For efficient inference, TensorFlow combines batch normalization with the preceding convolutional and fully-connected layers prior to quantization by -[folding batch norm layers](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/quantize/python/fold_batch_norms.py){:.external}. +[folding batch norm layers](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/quantize/python/fold_batch_norms.py){:.external}. -The quantization error is modeled using [fake quantization](../api_guides/python/array_ops.md#Fake_quantization) +The quantization error is modeled using [fake quantization](../../api_guides/python/array_ops.md#Fake_quantization) nodes to simulate the effect of quantization in the forward and backward passes. The forward-pass models quantization, while the backward-pass models quantization as a straight-through estimator. Both the forward- and backward-pass simulate the quantization @@ -105,12 +105,12 @@ toco \ --std_value=127.5 --mean_value=127.5 ``` -See the documentation for `tf.contrib.quantize` and [TensorFlow Lite](../lite/). +See the documentation for `tf.contrib.quantize` and [TensorFlow Lite](../../lite/). ## Quantized accuracy results -The following are results of trainiing some popular CNN models (Mobilenet-v1, +The following are results of training some popular CNN models (Mobilenet-v1, Mobilenet-v2, and Inception-v3) using this tool:

diff --git a/tensorflow/contrib/quantize/python/fold_batch_norms.py b/tensorflow/contrib/quantize/python/fold_batch_norms.py index e0c6da00d86fe4c5f881bcab7b444182da092b8f..a70f748fad60c6467946225ad5035caaf89c2aaf 100644 --- a/tensorflow/contrib/quantize/python/fold_batch_norms.py +++ b/tensorflow/contrib/quantize/python/fold_batch_norms.py @@ -454,7 +454,7 @@ def _CloneWithNewOperands(layer_op, input_tensor, weight_tensor, strides=layer_op.get_attr('strides'), padding=layer_op.get_attr('padding'), use_cudnn_on_gpu=layer_op.get_attr('use_cudnn_on_gpu'), - data_format=layer_op.get_attr('data_format'), + data_format=layer_op.get_attr('data_format').decode(), name=new_layer_name) elif layer_op.type == 'MatMul': return math_ops.matmul( @@ -867,7 +867,7 @@ class _OpCloner(object): strides=op.get_attr('strides'), padding=op.get_attr('padding'), use_cudnn_on_gpu=op.get_attr('use_cudnn_on_gpu'), - data_format=op.get_attr('data_format'), + data_format=op.get_attr('data_format').decode(), name=new_name).op def _CloneDepthwiseConv2d(self, op, inputs, new_name): diff --git a/tensorflow/contrib/quantize/python/quant_ops.py b/tensorflow/contrib/quantize/python/quant_ops.py index 8619708cdaecd78bcc7de0e8e0cbf2baa11bf6a2..39082cacf9770619cf5fb529ac9a0aad6e955c6d 100644 --- a/tensorflow/contrib/quantize/python/quant_ops.py +++ b/tensorflow/contrib/quantize/python/quant_ops.py @@ -224,8 +224,8 @@ def MovingAvgQuantize(inputs, None, default_name=name_prefix, values=[inputs], reuse=reuse) as scope: scope.set_partitioner(None) input_shape = inputs.get_shape() - input_dim = len(input_shape) if per_channel: + input_dim = len(input_shape) # Only support quantizing 1-, 2- and 4-dimensional tensors. assert input_dim in [1, 2, 4], ('Expected 1D, 2D or 4D input, was: %s in ' ' scope: %s' % (input_shape, name_prefix)) diff --git a/tensorflow/contrib/quantize/python/quant_ops_test.py b/tensorflow/contrib/quantize/python/quant_ops_test.py index 36d2af94e059cdc75b758bbf607d26c4e1ee73e9..c636c90d23a0f5a6de9d14085c824283cb41f6ca 100644 --- a/tensorflow/contrib/quantize/python/quant_ops_test.py +++ b/tensorflow/contrib/quantize/python/quant_ops_test.py @@ -63,6 +63,12 @@ class QuantOpsTest(googletest.TestCase): self.assertAlmostEqual(min_value, -0.5, delta=1e-3) self.assertAlmostEqual(max_value, 0.5, delta=1e-3) + def testMovingAvgQuantizeTrainingAssignNoShape(self): + min_value, max_value = self._GetMinMaxValues( + quant_ops.MovingAvgQuantize, [[-1, 1], [0, 0]], shape=None) + self.assertAlmostEqual(min_value, -0.5, delta=1e-3) + self.assertAlmostEqual(max_value, 0.5, delta=1e-3) + def testMovingAvgSymmetricQuantizeTrainingAssign(self): min_value, max_value = self._GetMinMaxValues( quant_ops.MovingAvgQuantize, [[-1, 0.5], [0, 0]], symmetric=True) @@ -109,10 +115,10 @@ class QuantOpsTest(googletest.TestCase): is_training=True, vars_collection=_MIN_MAX_VARS) - def _GetMinMaxValues(self, quantize_fn, input_values, **kwds): + def _GetMinMaxValues(self, quantize_fn, input_values, shape=(2), **kwds): g = ops.Graph() with session.Session(graph=g) as sess: - x = array_ops.placeholder(dtypes.float32, shape=[2]) + x = array_ops.placeholder(dtypes.float32, shape=shape) y = quantize_fn( x, init_min=0.0, diff --git a/tensorflow/contrib/receptive_field/python/util/parse_layer_parameters.py b/tensorflow/contrib/receptive_field/python/util/parse_layer_parameters.py index 0e3c46f17d2e2a277418d39e31927db73a509670..92ae1021bc8f8fbf19ca7f7cbe208ecea18128e8 100644 --- a/tensorflow/contrib/receptive_field/python/util/parse_layer_parameters.py +++ b/tensorflow/contrib/receptive_field/python/util/parse_layer_parameters.py @@ -27,7 +27,8 @@ from tensorflow.python.platform import tf_logging as logging _UNCHANGED_RF_LAYER_OPS = [ "Add", "BiasAdd", "Cast", "Ceil", "ConcatV2", "Const", "Floor", "FusedBatchNorm", "Identity", "Log", "Mul", "Pow", "RealDiv", "Relu", - "Relu6", "Round", "Rsqrt", "Softplus", "Sub", "VariableV2", "LRN" + "Relu6", "Round", "Rsqrt", "Softplus", "Sub", "VariableV2", "LRN", + "GreaterEqual" ] # Different ways in which padding modes may be spelled. @@ -276,11 +277,11 @@ def get_layer_params(node, name_to_node, input_resolution=None, force=False): kernel_size_x, kernel_size_y = _conv_kernel_size(node, name_to_node) # Compute the padding for this node separately for each direction. total_padding_x, padding_x = _padding_size_conv_pool( - node, kernel_size_x, stride_x, input_resolution[1] - if input_resolution is not None else None) + node, kernel_size_x, stride_x, + input_resolution[1] if input_resolution is not None else None) total_padding_y, padding_y = _padding_size_conv_pool( - node, kernel_size_y, stride_y, input_resolution[0] - if input_resolution is not None else None) + node, kernel_size_y, stride_y, + input_resolution[0] if input_resolution is not None else None) elif node.op == "Pad": # Kernel and stride are simply 1 in this case. kernel_size_x = 1 @@ -294,11 +295,11 @@ def get_layer_params(node, name_to_node, input_resolution=None, force=False): kernel_size_x, kernel_size_y = _pool_kernel_size(node, name_to_node) # Compute the padding for this node separately for each direction. total_padding_x, padding_x = _padding_size_conv_pool( - node, kernel_size_x, stride_x, input_resolution[1] - if input_resolution is not None else None) + node, kernel_size_x, stride_x, + input_resolution[1] if input_resolution is not None else None) total_padding_y, padding_y = _padding_size_conv_pool( - node, kernel_size_y, stride_y, input_resolution[0] - if input_resolution is not None else None) + node, kernel_size_y, stride_y, + input_resolution[0] if input_resolution is not None else None) elif node.op in _UNCHANGED_RF_LAYER_OPS: # These nodes do not modify the RF parameters. kernel_size_x = 1 @@ -320,7 +321,7 @@ def get_layer_params(node, name_to_node, input_resolution=None, force=False): total_padding_y = None padding_y = None else: - raise ValueError("Unknown layer for operation '%s': %s" % (node.name, - node.op)) + raise ValueError( + "Unknown layer for operation '%s': %s" % (node.name, node.op)) return (kernel_size_x, kernel_size_y, stride_x, stride_y, padding_x, padding_y, total_padding_x, total_padding_y) diff --git a/tensorflow/contrib/remote_fused_graph/pylib/python/ops/remote_fused_graph_ops.py b/tensorflow/contrib/remote_fused_graph/pylib/python/ops/remote_fused_graph_ops.py index 2054367f0d1461c8868e3332d82322a8a3dd38af..7e79785d2867de586f0730373d4864602ef770ae 100644 --- a/tensorflow/contrib/remote_fused_graph/pylib/python/ops/remote_fused_graph_ops.py +++ b/tensorflow/contrib/remote_fused_graph/pylib/python/ops/remote_fused_graph_ops.py @@ -50,13 +50,13 @@ def remote_fused_graph_execute(inputs, if default_graph_input_tensor_type_shapes: for type_shape in default_graph_input_tensor_type_shapes: type_shape_proto = info_proto.default_graph_input_tensor_shape.add() - type_shape_proto.dtype = int(dtypes.as_dtype(type_shape[0])) + type_shape_proto.dtype = dtypes.as_dtype(type_shape[0]).as_datatype_enum for dim in type_shape[1]: type_shape_proto.shape.dim.add().size = dim if default_graph_output_tensor_type_shapes: for type_shape in default_graph_output_tensor_type_shapes: type_shape_proto = info_proto.default_graph_output_tensor_shape.add() - type_shape_proto.dtype = int(dtypes.as_dtype(type_shape[0])) + type_shape_proto.dtype = dtypes.as_dtype(type_shape[0]).as_datatype_enum for dim in type_shape[1]: type_shape_proto.shape.dim.add().size = dim diff --git a/tensorflow/contrib/rnn/BUILD b/tensorflow/contrib/rnn/BUILD index 39b688596875ab1b208d97a5d6f9a5ee811674cb..24fa740d24502a28cb42c994715d09180ee99899 100644 --- a/tensorflow/contrib/rnn/BUILD +++ b/tensorflow/contrib/rnn/BUILD @@ -102,26 +102,6 @@ cuda_py_tests( xla_enabled = True, ) -cuda_py_tests( - name = "core_rnn_cell_test", - size = "medium", - srcs = ["python/kernel_tests/core_rnn_cell_test.py"], - additional_deps = [ - ":rnn_py", - "//third_party/py/numpy", - "//tensorflow/python:array_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:init_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python:rnn", - "//tensorflow/python:rnn_cell", - "//tensorflow/python:variable_scope", - "//tensorflow/python:variables", - "@absl_py//absl/testing:parameterized", - ], -) - cuda_py_tests( name = "rnn_test", size = "medium", @@ -144,32 +124,6 @@ cuda_py_tests( ], ) -cuda_py_tests( - name = "core_rnn_test", - size = "medium", - srcs = ["python/kernel_tests/core_rnn_test.py"], - additional_deps = [ - ":rnn_py", - "//third_party/py/numpy", - "//tensorflow/python:array_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:control_flow_ops", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:gradients", - "//tensorflow/python:init_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python:platform", - "//tensorflow/python:rnn", - "//tensorflow/python:tensor_array_ops", - "//tensorflow/python:util", - "//tensorflow/python:variable_scope", - "//tensorflow/python:variables", - "//tensorflow/python/eager:context", - ], - shard_count = 10, -) - tf_py_test( name = "fused_rnn_cell_test", size = "medium", @@ -388,6 +342,13 @@ py_binary( name = "checkpoint_convert", srcs = ["python/tools/checkpoint_convert.py"], srcs_version = "PY2AND3", + deps = [":checkpoint_convert_lib"], +) + +py_library( + name = "checkpoint_convert_lib", + srcs = ["python/tools/checkpoint_convert.py"], + srcs_version = "PY2AND3", deps = [ "//tensorflow/core:protos_all_py", "//tensorflow/python:framework_ops", @@ -406,7 +367,7 @@ py_test( srcs_version = "PY2AND3", tags = ["no_pip"], deps = [ - ":checkpoint_convert", + ":checkpoint_convert_lib", "//tensorflow/python:client_testlib", "//tensorflow/python:framework_ops", "//tensorflow/python:session", diff --git a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py deleted file mode 100644 index a0d013c618ea56077098b15b7eed5f9110239516..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py +++ /dev/null @@ -1,1181 +0,0 @@ -# Copyright 2015 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 RNN cells.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import os - -from absl.testing import parameterized -import numpy as np - -from tensorflow.contrib import rnn as contrib_rnn -from tensorflow.contrib.rnn.python.ops import core_rnn_cell -from tensorflow.contrib.rnn.python.ops import rnn_cell as contrib_rnn_cell -from tensorflow.core.protobuf import config_pb2 -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops -from tensorflow.python.framework import random_seed -from tensorflow.python.framework import test_util -from tensorflow.python.keras import layers as keras_layers -from tensorflow.python.layers import base as base_layer -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import init_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import rnn -from tensorflow.python.ops import rnn_cell_impl -from tensorflow.python.ops import variable_scope -from tensorflow.python.ops import variables as variables_lib -from tensorflow.python.platform import test -from tensorflow.python.training.checkpointable import util as checkpointable_utils - -# pylint: enable=protected-access -Linear = core_rnn_cell._Linear # pylint: disable=invalid-name - - -class RNNCellTest(test.TestCase): - - def testLinear(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(1.0)): - x = array_ops.zeros([1, 2]) - l = Linear([x], 2, False)([x]) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([l], {x.name: np.array([[1., 2.]])}) - self.assertAllClose(res[0], [[3.0, 3.0]]) - - # Checks prevent you from accidentally creating a shared function. - with self.assertRaises(ValueError): - l1 = Linear([x], 2, False)([x]) - - # But you can create a new one in a new scope and share the variables. - with variable_scope.variable_scope("l1") as new_scope: - l1 = Linear([x], 2, False)([x]) - with variable_scope.variable_scope(new_scope, reuse=True): - Linear([l1], 2, False)([l1]) - self.assertEqual(len(variables_lib.trainable_variables()), 2) - - def testBasicRNNCell(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 2]) - m = array_ops.zeros([1, 2]) - cell = rnn_cell_impl.BasicRNNCell(2) - g, _ = cell(x, m) - self.assertEqual([ - "root/basic_rnn_cell/%s:0" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME, - "root/basic_rnn_cell/%s:0" % rnn_cell_impl._BIAS_VARIABLE_NAME - ], [v.name for v in cell.trainable_variables]) - self.assertFalse(cell.non_trainable_variables) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g], { - x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1]]) - }) - self.assertEqual(res[0].shape, (1, 2)) - - def testBasicRNNCellNotTrainable(self): - with self.cached_session() as sess: - - def not_trainable_getter(getter, *args, **kwargs): - kwargs["trainable"] = False - return getter(*args, **kwargs) - - with variable_scope.variable_scope( - "root", - initializer=init_ops.constant_initializer(0.5), - custom_getter=not_trainable_getter): - x = array_ops.zeros([1, 2]) - m = array_ops.zeros([1, 2]) - cell = rnn_cell_impl.BasicRNNCell(2) - g, _ = cell(x, m) - self.assertFalse(cell.trainable_variables) - self.assertEqual([ - "root/basic_rnn_cell/%s:0" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME, - "root/basic_rnn_cell/%s:0" % rnn_cell_impl._BIAS_VARIABLE_NAME - ], [v.name for v in cell.non_trainable_variables]) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g], { - x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1]]) - }) - self.assertEqual(res[0].shape, (1, 2)) - - def testIndRNNCell(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 2]) - m = array_ops.zeros([1, 2]) - cell = contrib_rnn_cell.IndRNNCell(2) - g, _ = cell(x, m) - self.assertEqual([ - "root/ind_rnn_cell/%s_w:0" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME, - "root/ind_rnn_cell/%s_u:0" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME, - "root/ind_rnn_cell/%s:0" % rnn_cell_impl._BIAS_VARIABLE_NAME - ], [v.name for v in cell.trainable_variables]) - self.assertFalse(cell.non_trainable_variables) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g], { - x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1]]) - }) - self.assertEqual(res[0].shape, (1, 2)) - - def testGRUCell(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 2]) - m = array_ops.zeros([1, 2]) - g, _ = rnn_cell_impl.GRUCell(2)(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g], { - x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1]]) - }) - # Smoke test - self.assertAllClose(res[0], [[0.175991, 0.175991]]) - with variable_scope.variable_scope( - "other", initializer=init_ops.constant_initializer(0.5)): - # Test GRUCell with input_size != num_units. - x = array_ops.zeros([1, 3]) - m = array_ops.zeros([1, 2]) - g, _ = rnn_cell_impl.GRUCell(2)(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g], { - x.name: np.array([[1., 1., 1.]]), - m.name: np.array([[0.1, 0.1]]) - }) - # Smoke test - self.assertAllClose(res[0], [[0.156736, 0.156736]]) - - def testIndyGRUCell(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 2]) - m = array_ops.zeros([1, 2]) - g, _ = contrib_rnn_cell.IndyGRUCell(2)(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g], { - x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1]]) - }) - # Smoke test - self.assertAllClose(res[0], [[0.185265, 0.17704]]) - with variable_scope.variable_scope( - "other", initializer=init_ops.constant_initializer(0.5)): - # Test IndyGRUCell with input_size != num_units. - x = array_ops.zeros([1, 3]) - m = array_ops.zeros([1, 2]) - g, _ = contrib_rnn_cell.IndyGRUCell(2)(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g], { - x.name: np.array([[1., 1., 1.]]), - m.name: np.array([[0.1, 0.1]]) - }) - # Smoke test - self.assertAllClose(res[0], [[0.155127, 0.157328]]) - - def testSRUCell(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 2]) - m = array_ops.zeros([1, 2]) - g, _ = contrib_rnn_cell.SRUCell(2)(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g], { - x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1]]) - }) - # Smoke test - self.assertAllClose(res[0], [[0.509682, 0.509682]]) - - def testSRUCellWithDiffSize(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 3]) - m = array_ops.zeros([1, 2]) - g, _ = contrib_rnn_cell.SRUCell(2)(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g], { - x.name: np.array([[1., 1., 1.]]), - m.name: np.array([[0.1, 0.1]]) - }) - # Smoke test - self.assertAllClose(res[0], [[0.55255556, 0.55255556]]) - - def testBasicLSTMCell(self): - for dtype in [dtypes.float16, dtypes.float32]: - np_dtype = dtype.as_numpy_dtype - with self.session(graph=ops.Graph()) as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 2], dtype=dtype) - m = array_ops.zeros([1, 8], dtype=dtype) - cell = rnn_cell_impl.MultiRNNCell( - [ - rnn_cell_impl.BasicLSTMCell(2, state_is_tuple=False) - for _ in range(2) - ], - state_is_tuple=False) - self.assertEqual(cell.dtype, None) - self.assertEqual("cell-0", cell._checkpoint_dependencies[0].name) - self.assertEqual("cell-1", cell._checkpoint_dependencies[1].name) - cell.get_config() # Should not throw an error - g, out_m = cell(x, m) - # Layer infers the input type. - self.assertEqual(cell.dtype, dtype.name) - expected_variable_names = [ - "root/multi_rnn_cell/cell_0/basic_lstm_cell/%s:0" % - rnn_cell_impl._WEIGHTS_VARIABLE_NAME, - "root/multi_rnn_cell/cell_0/basic_lstm_cell/%s:0" % - rnn_cell_impl._BIAS_VARIABLE_NAME, - "root/multi_rnn_cell/cell_1/basic_lstm_cell/%s:0" % - rnn_cell_impl._WEIGHTS_VARIABLE_NAME, - "root/multi_rnn_cell/cell_1/basic_lstm_cell/%s:0" % - rnn_cell_impl._BIAS_VARIABLE_NAME - ] - self.assertEqual(expected_variable_names, - [v.name for v in cell.trainable_variables]) - self.assertFalse(cell.non_trainable_variables) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g, out_m], { - x.name: np.array([[1., 1.]]), - m.name: 0.1 * np.ones([1, 8]) - }) - self.assertEqual(len(res), 2) - variables = variables_lib.global_variables() - self.assertEqual(expected_variable_names, [v.name for v in variables]) - # The numbers in results were not calculated, this is just a - # smoke test. - self.assertAllClose(res[0], np.array( - [[0.240, 0.240]], dtype=np_dtype), 1e-2) - expected_mem = np.array( - [[0.689, 0.689, 0.448, 0.448, 0.398, 0.398, 0.240, 0.240]], - dtype=np_dtype) - self.assertAllClose(res[1], expected_mem, 1e-2) - with variable_scope.variable_scope( - "other", initializer=init_ops.constant_initializer(0.5)): - # Test BasicLSTMCell with input_size != num_units. - x = array_ops.zeros([1, 3], dtype=dtype) - m = array_ops.zeros([1, 4], dtype=dtype) - g, out_m = rnn_cell_impl.BasicLSTMCell(2, state_is_tuple=False)(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g, out_m], { - x.name: np.array([[1., 1., 1.]], dtype=np_dtype), - m.name: 0.1 * np.ones([1, 4], dtype=np_dtype) - }) - self.assertEqual(len(res), 2) - - def testBasicLSTMCellDimension0Error(self): - """Tests that dimension 0 in both(x and m) shape must be equal.""" - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - num_units = 2 - state_size = num_units * 2 - batch_size = 3 - input_size = 4 - x = array_ops.zeros([batch_size, input_size]) - m = array_ops.zeros([batch_size - 1, state_size]) - with self.assertRaises(ValueError): - g, out_m = rnn_cell_impl.BasicLSTMCell( - num_units, state_is_tuple=False)(x, m) - sess.run([variables_lib.global_variables_initializer()]) - sess.run( - [g, out_m], { - x.name: 1 * np.ones([batch_size, input_size]), - m.name: 0.1 * np.ones([batch_size - 1, state_size]) - }) - - def testBasicLSTMCellStateSizeError(self): - """Tests that state_size must be num_units * 2.""" - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - num_units = 2 - state_size = num_units * 3 # state_size must be num_units * 2 - batch_size = 3 - input_size = 4 - x = array_ops.zeros([batch_size, input_size]) - m = array_ops.zeros([batch_size, state_size]) - with self.assertRaises(ValueError): - g, out_m = rnn_cell_impl.BasicLSTMCell( - num_units, state_is_tuple=False)(x, m) - sess.run([variables_lib.global_variables_initializer()]) - sess.run( - [g, out_m], { - x.name: 1 * np.ones([batch_size, input_size]), - m.name: 0.1 * np.ones([batch_size, state_size]) - }) - - def testBasicLSTMCellStateTupleType(self): - with self.cached_session(): - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 2]) - m0 = (array_ops.zeros([1, 2]),) * 2 - m1 = (array_ops.zeros([1, 2]),) * 2 - cell = rnn_cell_impl.MultiRNNCell( - [rnn_cell_impl.BasicLSTMCell(2) for _ in range(2)], - state_is_tuple=True) - self.assertTrue(isinstance(cell.state_size, tuple)) - self.assertTrue( - isinstance(cell.state_size[0], rnn_cell_impl.LSTMStateTuple)) - self.assertTrue( - isinstance(cell.state_size[1], rnn_cell_impl.LSTMStateTuple)) - - # Pass in regular tuples - _, (out_m0, out_m1) = cell(x, (m0, m1)) - self.assertTrue(isinstance(out_m0, rnn_cell_impl.LSTMStateTuple)) - self.assertTrue(isinstance(out_m1, rnn_cell_impl.LSTMStateTuple)) - - # Pass in LSTMStateTuples - variable_scope.get_variable_scope().reuse_variables() - zero_state = cell.zero_state(1, dtypes.float32) - self.assertTrue(isinstance(zero_state, tuple)) - self.assertTrue(isinstance(zero_state[0], rnn_cell_impl.LSTMStateTuple)) - self.assertTrue(isinstance(zero_state[1], rnn_cell_impl.LSTMStateTuple)) - _, (out_m0, out_m1) = cell(x, zero_state) - self.assertTrue(isinstance(out_m0, rnn_cell_impl.LSTMStateTuple)) - self.assertTrue(isinstance(out_m1, rnn_cell_impl.LSTMStateTuple)) - - def testBasicLSTMCellWithStateTuple(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 2]) - m0 = array_ops.zeros([1, 4]) - m1 = array_ops.zeros([1, 4]) - cell = rnn_cell_impl.MultiRNNCell( - [ - rnn_cell_impl.BasicLSTMCell(2, state_is_tuple=False) - for _ in range(2) - ], - state_is_tuple=True) - g, (out_m0, out_m1) = cell(x, (m0, m1)) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g, out_m0, out_m1], { - x.name: np.array([[1., 1.]]), - m0.name: 0.1 * np.ones([1, 4]), - m1.name: 0.1 * np.ones([1, 4]) - }) - self.assertEqual(len(res), 3) - # The numbers in results were not calculated, this is just a smoke test. - # Note, however, these values should match the original - # version having state_is_tuple=False. - self.assertAllClose(res[0], [[0.24024698, 0.24024698]]) - expected_mem0 = np.array( - [[0.68967271, 0.68967271, 0.44848421, 0.44848421]]) - expected_mem1 = np.array( - [[0.39897051, 0.39897051, 0.24024698, 0.24024698]]) - self.assertAllClose(res[1], expected_mem0) - self.assertAllClose(res[2], expected_mem1) - - def testIndyLSTMCell(self): - for dtype in [dtypes.float16, dtypes.float32]: - np_dtype = dtype.as_numpy_dtype - with self.session(graph=ops.Graph()) as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 2], dtype=dtype) - state_0 = (array_ops.zeros([1, 2], dtype=dtype),) * 2 - state_1 = (array_ops.zeros([1, 2], dtype=dtype),) * 2 - cell = rnn_cell_impl.MultiRNNCell( - [contrib_rnn_cell.IndyLSTMCell(2) for _ in range(2)]) - self.assertEqual(cell.dtype, None) - self.assertEqual("cell-0", cell._checkpoint_dependencies[0].name) - self.assertEqual("cell-1", cell._checkpoint_dependencies[1].name) - cell.get_config() # Should not throw an error - g, (out_state_0, out_state_1) = cell(x, (state_0, state_1)) - # Layer infers the input type. - self.assertEqual(cell.dtype, dtype.name) - expected_variable_names = [ - "root/multi_rnn_cell/cell_0/indy_lstm_cell/%s_w:0" % - rnn_cell_impl._WEIGHTS_VARIABLE_NAME, - "root/multi_rnn_cell/cell_0/indy_lstm_cell/%s_u:0" % - rnn_cell_impl._WEIGHTS_VARIABLE_NAME, - "root/multi_rnn_cell/cell_0/indy_lstm_cell/%s:0" % - rnn_cell_impl._BIAS_VARIABLE_NAME, - "root/multi_rnn_cell/cell_1/indy_lstm_cell/%s_w:0" % - rnn_cell_impl._WEIGHTS_VARIABLE_NAME, - "root/multi_rnn_cell/cell_1/indy_lstm_cell/%s_u:0" % - rnn_cell_impl._WEIGHTS_VARIABLE_NAME, - "root/multi_rnn_cell/cell_1/indy_lstm_cell/%s:0" % - rnn_cell_impl._BIAS_VARIABLE_NAME - ] - self.assertEqual(expected_variable_names, - [v.name for v in cell.trainable_variables]) - self.assertFalse(cell.non_trainable_variables) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g, out_state_0, out_state_1], { - x.name: np.array([[1., 1.]]), - state_0[0].name: 0.1 * np.ones([1, 2]), - state_0[1].name: 0.1 * np.ones([1, 2]), - state_1[0].name: 0.1 * np.ones([1, 2]), - state_1[1].name: 0.1 * np.ones([1, 2]), - }) - self.assertEqual(len(res), 3) - variables = variables_lib.global_variables() - self.assertEqual(expected_variable_names, [v.name for v in variables]) - # Only check the range of outputs as this is just a smoke test. - self.assertAllInRange(res[0], -1.0, 1.0) - self.assertAllInRange(res[1], -1.0, 1.0) - self.assertAllInRange(res[2], -1.0, 1.0) - with variable_scope.variable_scope( - "other", initializer=init_ops.constant_initializer(0.5)): - # Test IndyLSTMCell with input_size != num_units. - x = array_ops.zeros([1, 3], dtype=dtype) - state = (array_ops.zeros([1, 2], dtype=dtype),) * 2 - g, out_state = contrib_rnn_cell.IndyLSTMCell(2)(x, state) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g, out_state], { - x.name: np.array([[1., 1., 1.]], dtype=np_dtype), - state[0].name: 0.1 * np.ones([1, 2], dtype=np_dtype), - state[1].name: 0.1 * np.ones([1, 2], dtype=np_dtype), - }) - self.assertEqual(len(res), 2) - - def testLSTMCell(self): - with self.cached_session() as sess: - num_units = 8 - num_proj = 6 - state_size = num_units + num_proj - batch_size = 3 - input_size = 2 - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([batch_size, input_size]) - m = array_ops.zeros([batch_size, state_size]) - cell = rnn_cell_impl.LSTMCell( - num_units=num_units, - num_proj=num_proj, - forget_bias=1.0, - state_is_tuple=False) - output, state = cell(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [output, state], { - x.name: np.array([[1., 1.], [2., 2.], [3., 3.]]), - m.name: 0.1 * np.ones((batch_size, state_size)) - }) - self.assertEqual(len(res), 2) - # The numbers in results were not calculated, this is mostly just a - # smoke test. - self.assertEqual(res[0].shape, (batch_size, num_proj)) - self.assertEqual(res[1].shape, (batch_size, state_size)) - # Different inputs so different outputs and states - for i in range(1, batch_size): - self.assertTrue( - float(np.linalg.norm((res[0][0, :] - res[0][i, :]))) > 1e-6) - self.assertTrue( - float(np.linalg.norm((res[1][0, :] - res[1][i, :]))) > 1e-6) - - def testLSTMCellVariables(self): - with self.cached_session(): - num_units = 8 - num_proj = 6 - state_size = num_units + num_proj - batch_size = 3 - input_size = 2 - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([batch_size, input_size]) - m = array_ops.zeros([batch_size, state_size]) - cell = rnn_cell_impl.LSTMCell( - num_units=num_units, - num_proj=num_proj, - forget_bias=1.0, - state_is_tuple=False) - cell(x, m) # Execute to create variables - variables = variables_lib.global_variables() - self.assertEquals(variables[0].op.name, "root/lstm_cell/kernel") - self.assertEquals(variables[1].op.name, "root/lstm_cell/bias") - self.assertEquals(variables[2].op.name, - "root/lstm_cell/projection/kernel") - - def testLSTMCellLayerNorm(self): - with self.cached_session() as sess: - num_units = 2 - num_proj = 3 - batch_size = 1 - input_size = 4 - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([batch_size, input_size]) - c = array_ops.zeros([batch_size, num_units]) - h = array_ops.zeros([batch_size, num_proj]) - state = rnn_cell_impl.LSTMStateTuple(c, h) - cell = contrib_rnn_cell.LayerNormLSTMCell( - num_units=num_units, - num_proj=num_proj, - forget_bias=1.0, - layer_norm=True, - norm_gain=1.0, - norm_shift=0.0) - g, out_m = cell(x, state) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - [g, out_m], { - x.name: np.ones((batch_size, input_size)), - c.name: 0.1 * np.ones((batch_size, num_units)), - h.name: 0.1 * np.ones((batch_size, num_proj)) - }) - self.assertEqual(len(res), 2) - # The numbers in results were not calculated, this is mostly just a - # smoke test. - self.assertEqual(res[0].shape, (batch_size, num_proj)) - self.assertEqual(res[1][0].shape, (batch_size, num_units)) - self.assertEqual(res[1][1].shape, (batch_size, num_proj)) - # Different inputs so different outputs and states - for i in range(1, batch_size): - self.assertTrue( - float(np.linalg.norm((res[0][0, :] - res[0][i, :]))) < 1e-6) - self.assertTrue( - float(np.linalg.norm((res[1][0, :] - res[1][i, :]))) < 1e-6) - - @test_util.run_in_graph_and_eager_modes - def testWrapperCheckpointing(self): - for wrapper_type in [ - rnn_cell_impl.DropoutWrapper, - rnn_cell_impl.ResidualWrapper, - lambda cell: rnn_cell_impl.MultiRNNCell([cell])]: - cell = rnn_cell_impl.BasicRNNCell(1) - wrapper = wrapper_type(cell) - wrapper(array_ops.ones([1, 1]), - state=wrapper.zero_state(batch_size=1, dtype=dtypes.float32)) - self.evaluate([v.initializer for v in cell.variables]) - checkpoint = checkpointable_utils.Checkpoint(wrapper=wrapper) - prefix = os.path.join(self.get_temp_dir(), "ckpt") - self.evaluate(cell._bias.assign([40.])) - save_path = checkpoint.save(prefix) - self.evaluate(cell._bias.assign([0.])) - checkpoint.restore(save_path).assert_consumed().run_restore_ops() - self.assertAllEqual([40.], self.evaluate(cell._bias)) - - def testOutputProjectionWrapper(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 3]) - m = array_ops.zeros([1, 3]) - cell = contrib_rnn.OutputProjectionWrapper(rnn_cell_impl.GRUCell(3), 2) - g, new_m = cell(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g, new_m], { - x.name: np.array([[1., 1., 1.]]), - m.name: np.array([[0.1, 0.1, 0.1]]) - }) - self.assertEqual(res[1].shape, (1, 3)) - # The numbers in results were not calculated, this is just a smoke test. - self.assertAllClose(res[0], [[0.231907, 0.231907]]) - - def testInputProjectionWrapper(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 2]) - m = array_ops.zeros([1, 3]) - cell = contrib_rnn.InputProjectionWrapper( - rnn_cell_impl.GRUCell(3), num_proj=3) - g, new_m = cell(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g, new_m], { - x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1, 0.1]]) - }) - self.assertEqual(res[1].shape, (1, 3)) - # The numbers in results were not calculated, this is just a smoke test. - self.assertAllClose(res[0], [[0.154605, 0.154605, 0.154605]]) - - def testResidualWrapper(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 3]) - m = array_ops.zeros([1, 3]) - base_cell = rnn_cell_impl.GRUCell(3) - g, m_new = base_cell(x, m) - variable_scope.get_variable_scope().reuse_variables() - wrapper_object = rnn_cell_impl.ResidualWrapper(base_cell) - (name, dep), = wrapper_object._checkpoint_dependencies - wrapper_object.get_config() # Should not throw an error - self.assertIs(dep, base_cell) - self.assertEqual("cell", name) - - g_res, m_new_res = wrapper_object(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g, g_res, m_new, m_new_res], { - x: np.array([[1., 1., 1.]]), - m: np.array([[0.1, 0.1, 0.1]]) - }) - # Residual connections - self.assertAllClose(res[1], res[0] + [1., 1., 1.]) - # States are left untouched - self.assertAllClose(res[2], res[3]) - - def testResidualWrapperWithSlice(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 5]) - m = array_ops.zeros([1, 3]) - base_cell = rnn_cell_impl.GRUCell(3) - g, m_new = base_cell(x, m) - variable_scope.get_variable_scope().reuse_variables() - - def residual_with_slice_fn(inp, out): - inp_sliced = array_ops.slice(inp, [0, 0], [-1, 3]) - return inp_sliced + out - - g_res, m_new_res = rnn_cell_impl.ResidualWrapper( - base_cell, residual_with_slice_fn)(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res_g, res_g_res, res_m_new, res_m_new_res = sess.run( - [g, g_res, m_new, m_new_res], { - x: np.array([[1., 1., 1., 1., 1.]]), - m: np.array([[0.1, 0.1, 0.1]]) - }) - # Residual connections - self.assertAllClose(res_g_res, res_g + [1., 1., 1.]) - # States are left untouched - self.assertAllClose(res_m_new, res_m_new_res) - - def testDeviceWrapper(self): - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 3]) - m = array_ops.zeros([1, 3]) - wrapped = rnn_cell_impl.GRUCell(3) - cell = rnn_cell_impl.DeviceWrapper(wrapped, "/cpu:14159") - (name, dep), = cell._checkpoint_dependencies - cell.get_config() # Should not throw an error - self.assertIs(dep, wrapped) - self.assertEqual("cell", name) - - outputs, _ = cell(x, m) - self.assertTrue("cpu:14159" in outputs.device.lower()) - - def _retrieve_cpu_gpu_stats(self, run_metadata): - cpu_stats = None - gpu_stats = None - step_stats = run_metadata.step_stats - for ds in step_stats.dev_stats: - if "cpu:0" in ds.device[-5:].lower(): - cpu_stats = ds.node_stats - if "gpu:0" == ds.device[-5:].lower(): - gpu_stats = ds.node_stats - return cpu_stats, gpu_stats - - def testDeviceWrapperDynamicExecutionNodesAreAllProperlyLocated(self): - if not test.is_gpu_available(): - # Can't perform this test w/o a GPU - return - - gpu_dev = test.gpu_device_name() - with self.session(use_gpu=True) as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 1, 3]) - cell = rnn_cell_impl.DeviceWrapper(rnn_cell_impl.GRUCell(3), gpu_dev) - with ops.device("/cpu:0"): - outputs, _ = rnn.dynamic_rnn( - cell=cell, inputs=x, dtype=dtypes.float32) - run_metadata = config_pb2.RunMetadata() - opts = config_pb2.RunOptions( - trace_level=config_pb2.RunOptions.FULL_TRACE) - - sess.run([variables_lib.global_variables_initializer()]) - _ = sess.run(outputs, options=opts, run_metadata=run_metadata) - - cpu_stats, gpu_stats = self._retrieve_cpu_gpu_stats(run_metadata) - self.assertFalse([s for s in cpu_stats if "gru_cell" in s.node_name]) - self.assertTrue([s for s in gpu_stats if "gru_cell" in s.node_name]) - - def testEmbeddingWrapper(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 1], dtype=dtypes.int32) - m = array_ops.zeros([1, 2]) - embedding_cell = contrib_rnn.EmbeddingWrapper( - rnn_cell_impl.GRUCell(2), embedding_classes=3, embedding_size=2) - self.assertEqual(embedding_cell.output_size, 2) - g, new_m = embedding_cell(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([g, new_m], { - x.name: np.array([[1]]), - m.name: np.array([[0.1, 0.1]]) - }) - self.assertEqual(res[1].shape, (1, 2)) - # The numbers in results were not calculated, this is just a smoke test. - self.assertAllClose(res[0], [[0.17139, 0.17139]]) - - def testEmbeddingWrapperWithDynamicRnn(self): - with self.cached_session() as sess: - with variable_scope.variable_scope("root"): - inputs = ops.convert_to_tensor([[[0], [0]]], dtype=dtypes.int64) - input_lengths = ops.convert_to_tensor([2], dtype=dtypes.int64) - embedding_cell = contrib_rnn.EmbeddingWrapper( - rnn_cell_impl.BasicLSTMCell(1, state_is_tuple=True), - embedding_classes=1, - embedding_size=2) - outputs, _ = rnn.dynamic_rnn( - cell=embedding_cell, - inputs=inputs, - sequence_length=input_lengths, - dtype=dtypes.float32) - sess.run([variables_lib.global_variables_initializer()]) - # This will fail if output's dtype is inferred from input's. - sess.run(outputs) - - def testMultiRNNCell(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 2]) - m = array_ops.zeros([1, 4]) - multi_rnn_cell = rnn_cell_impl.MultiRNNCell( - [rnn_cell_impl.GRUCell(2) for _ in range(2)], - state_is_tuple=False) - _, ml = multi_rnn_cell(x, m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run(ml, { - x.name: np.array([[1., 1.]]), - m.name: np.array([[0.1, 0.1, 0.1, 0.1]]) - }) - # The numbers in results were not calculated, this is just a smoke test. - self.assertAllClose(res, [[0.175991, 0.175991, 0.13248, 0.13248]]) - self.assertEqual(len(multi_rnn_cell.weights), 2 * 4) - self.assertTrue( - [x.dtype == dtypes.float32 for x in multi_rnn_cell.weights]) - - def testMultiRNNCellWithStateTuple(self): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - x = array_ops.zeros([1, 2]) - m_bad = array_ops.zeros([1, 4]) - m_good = (array_ops.zeros([1, 2]), array_ops.zeros([1, 2])) - - # Test incorrectness of state - with self.assertRaisesRegexp(ValueError, "Expected state .* a tuple"): - rnn_cell_impl.MultiRNNCell( - [rnn_cell_impl.GRUCell(2) for _ in range(2)], - state_is_tuple=True)(x, m_bad) - - _, ml = rnn_cell_impl.MultiRNNCell( - [rnn_cell_impl.GRUCell(2) for _ in range(2)], - state_is_tuple=True)(x, m_good) - - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run( - ml, { - x.name: np.array([[1., 1.]]), - m_good[0].name: np.array([[0.1, 0.1]]), - m_good[1].name: np.array([[0.1, 0.1]]) - }) - - # The numbers in results were not calculated, this is just a - # smoke test. However, these numbers should match those of - # the test testMultiRNNCell. - self.assertAllClose(res[0], [[0.175991, 0.175991]]) - self.assertAllClose(res[1], [[0.13248, 0.13248]]) - - -class DropoutWrapperTest(test.TestCase, parameterized.TestCase): - - def _testDropoutWrapper(self, - batch_size=None, - time_steps=None, - parallel_iterations=None, - wrapper_type=None, - **kwargs): - with self.cached_session() as sess: - with variable_scope.variable_scope( - "root", initializer=init_ops.constant_initializer(0.5)): - if batch_size is None and time_steps is None: - # 2 time steps, batch size 1, depth 3 - batch_size = 1 - time_steps = 2 - x = constant_op.constant( - [[[2., 2., 2.]], [[1., 1., 1.]]], dtype=dtypes.float32) - m = rnn_cell_impl.LSTMStateTuple( - *[constant_op.constant([[0.1, 0.1, 0.1]], dtype=dtypes.float32 - )] * 2) - else: - x = constant_op.constant( - np.random.randn(time_steps, batch_size, 3).astype(np.float32)) - m = rnn_cell_impl.LSTMStateTuple(*[ - constant_op. - constant([[0.1, 0.1, 0.1]] * batch_size, dtype=dtypes.float32) - ] * 2) - outputs, final_state = rnn.dynamic_rnn( - cell=wrapper_type( - rnn_cell_impl.LSTMCell(3), dtype=x.dtype, **kwargs), - time_major=True, - parallel_iterations=parallel_iterations, - inputs=x, - initial_state=m) - sess.run([variables_lib.global_variables_initializer()]) - res = sess.run([outputs, final_state]) - self.assertEqual(res[0].shape, (time_steps, batch_size, 3)) - self.assertEqual(res[1].c.shape, (batch_size, 3)) - self.assertEqual(res[1].h.shape, (batch_size, 3)) - return res - - @parameterized.parameters( - [rnn_cell_impl.DropoutWrapper, rnn_cell_impl.DropoutWrapperV2]) - def testDropoutWrapperProperties(self, wrapper_type): - cell = rnn_cell_impl.BasicRNNCell(10) - wrapper = wrapper_type(cell) - # Github issue 15810 - self.assertEqual(wrapper.wrapped_cell, cell) - self.assertEqual(wrapper.state_size, 10) - self.assertEqual(wrapper.output_size, 10) - - @parameterized.parameters( - [rnn_cell_impl.DropoutWrapper, rnn_cell_impl.DropoutWrapperV2]) - def testDropoutWrapperZeroState(self, wrapper_type): - class _Cell(rnn_cell_impl.BasicRNNCell): - - def zero_state(self, batch_size=None, dtype=None): - return "wrapped_cell_zero_state" - wrapper = wrapper_type(_Cell(10)) - self.assertEqual(wrapper.zero_state(10, dtypes.float32), - "wrapped_cell_zero_state") - - @parameterized.parameters( - [rnn_cell_impl.DropoutWrapper, rnn_cell_impl.DropoutWrapperV2]) - def testDropoutWrapperKeepAllConstantInput(self, wrapper_type): - keep = array_ops.ones([]) - res = self._testDropoutWrapper( - input_keep_prob=keep, output_keep_prob=keep, state_keep_prob=keep, - wrapper_type=wrapper_type) - true_full_output = np.array( - [[[0.751109, 0.751109, 0.751109]], [[0.895509, 0.895509, 0.895509]]], - dtype=np.float32) - true_full_final_c = np.array( - [[1.949385, 1.949385, 1.949385]], dtype=np.float32) - self.assertAllClose(true_full_output, res[0]) - self.assertAllClose(true_full_output[1], res[1].h) - self.assertAllClose(true_full_final_c, res[1].c) - - @parameterized.parameters( - [rnn_cell_impl.DropoutWrapper, rnn_cell_impl.DropoutWrapperV2]) - def testDropoutWrapperKeepAll(self, wrapper_type): - keep = variable_scope.get_variable("all", initializer=1.0) - res = self._testDropoutWrapper( - input_keep_prob=keep, output_keep_prob=keep, state_keep_prob=keep, - wrapper_type=wrapper_type) - true_full_output = np.array( - [[[0.751109, 0.751109, 0.751109]], [[0.895509, 0.895509, 0.895509]]], - dtype=np.float32) - true_full_final_c = np.array( - [[1.949385, 1.949385, 1.949385]], dtype=np.float32) - self.assertAllClose(true_full_output, res[0]) - self.assertAllClose(true_full_output[1], res[1].h) - self.assertAllClose(true_full_final_c, res[1].c) - - @parameterized.parameters( - [rnn_cell_impl.DropoutWrapper, rnn_cell_impl.DropoutWrapperV2]) - def testDropoutWrapperWithSeed(self, wrapper_type): - keep_some = 0.5 - random_seed.set_random_seed(2) - ## Use parallel_iterations = 1 in both calls to - ## _testDropoutWrapper to ensure the (per-time step) dropout is - ## consistent across both calls. Otherwise the seed may not end - ## up being munged consistently across both graphs. - res_standard_1 = self._testDropoutWrapper( - input_keep_prob=keep_some, - output_keep_prob=keep_some, - state_keep_prob=keep_some, - seed=10, - parallel_iterations=1, - wrapper_type=wrapper_type) - # Clear away the graph and the test session (which keeps variables around) - ops.reset_default_graph() - self._ClearCachedSession() - random_seed.set_random_seed(2) - res_standard_2 = self._testDropoutWrapper( - input_keep_prob=keep_some, - output_keep_prob=keep_some, - state_keep_prob=keep_some, - seed=10, - parallel_iterations=1, - wrapper_type=wrapper_type) - self.assertAllClose(res_standard_1[0], res_standard_2[0]) - self.assertAllClose(res_standard_1[1].c, res_standard_2[1].c) - self.assertAllClose(res_standard_1[1].h, res_standard_2[1].h) - - @parameterized.parameters( - [rnn_cell_impl.DropoutWrapper, rnn_cell_impl.DropoutWrapperV2]) - def testDropoutWrapperKeepNoOutput(self, wrapper_type): - keep_all = variable_scope.get_variable("all", initializer=1.0) - keep_none = variable_scope.get_variable("none", initializer=1e-6) - res = self._testDropoutWrapper( - input_keep_prob=keep_all, - output_keep_prob=keep_none, - state_keep_prob=keep_all, - wrapper_type=wrapper_type) - true_full_output = np.array( - [[[0.751109, 0.751109, 0.751109]], [[0.895509, 0.895509, 0.895509]]], - dtype=np.float32) - true_full_final_c = np.array( - [[1.949385, 1.949385, 1.949385]], dtype=np.float32) - self.assertAllClose(np.zeros(res[0].shape), res[0]) - self.assertAllClose(true_full_output[1], res[1].h) - self.assertAllClose(true_full_final_c, res[1].c) - - @parameterized.parameters( - [rnn_cell_impl.DropoutWrapper, rnn_cell_impl.DropoutWrapperV2]) - def testDropoutWrapperKeepNoStateExceptLSTMCellMemory(self, wrapper_type): - keep_all = variable_scope.get_variable("all", initializer=1.0) - keep_none = variable_scope.get_variable("none", initializer=1e-6) - # Even though we dropout state, by default DropoutWrapper never - # drops out the memory ("c") term of an LSTMStateTuple. - res = self._testDropoutWrapper( - input_keep_prob=keep_all, - output_keep_prob=keep_all, - state_keep_prob=keep_none, - wrapper_type=wrapper_type) - true_c_state = np.array([[1.713925, 1.713925, 1.713925]], dtype=np.float32) - true_full_output = np.array( - [[[0.751109, 0.751109, 0.751109]], [[0.895509, 0.895509, 0.895509]]], - dtype=np.float32) - self.assertAllClose(true_full_output[0], res[0][0]) - # Second output is modified by zero input state - self.assertGreater(np.linalg.norm(true_full_output[1] - res[0][1]), 1e-4) - # h state has been set to zero - self.assertAllClose(np.zeros(res[1].h.shape), res[1].h) - # c state of an LSTMStateTuple is NEVER modified. - self.assertAllClose(true_c_state, res[1].c) - - @parameterized.parameters( - [rnn_cell_impl.DropoutWrapper, rnn_cell_impl.DropoutWrapperV2]) - def testDropoutWrapperKeepNoInput(self, wrapper_type): - keep_all = variable_scope.get_variable("all", initializer=1.0) - keep_none = variable_scope.get_variable("none", initializer=1e-6) - true_full_output = np.array( - [[[0.751109, 0.751109, 0.751109]], [[0.895509, 0.895509, 0.895509]]], - dtype=np.float32) - true_full_final_c = np.array( - [[1.949385, 1.949385, 1.949385]], dtype=np.float32) - # All outputs are different because inputs are zeroed out - res = self._testDropoutWrapper( - input_keep_prob=keep_none, - output_keep_prob=keep_all, - state_keep_prob=keep_all, - wrapper_type=wrapper_type) - self.assertGreater(np.linalg.norm(res[0] - true_full_output), 1e-4) - self.assertGreater(np.linalg.norm(res[1].h - true_full_output[1]), 1e-4) - self.assertGreater(np.linalg.norm(res[1].c - true_full_final_c), 1e-4) - - @parameterized.parameters( - [rnn_cell_impl.DropoutWrapper, rnn_cell_impl.DropoutWrapperV2]) - def testDropoutWrapperRecurrentOutput(self, wrapper_type): - keep_some = 0.8 - keep_all = variable_scope.get_variable("all", initializer=1.0) - res = self._testDropoutWrapper( - input_keep_prob=keep_all, - output_keep_prob=keep_some, - state_keep_prob=keep_all, - variational_recurrent=True, - wrapper_type=wrapper_type, - input_size=3, - batch_size=5, - time_steps=7) - # Ensure the same dropout pattern for all time steps - output_mask = np.abs(res[0]) > 1e-6 - for m in output_mask[1:]: - self.assertAllClose(output_mask[0], m) - - @parameterized.parameters( - [rnn_cell_impl.DropoutWrapper, rnn_cell_impl.DropoutWrapperV2]) - def testDropoutWrapperRecurrentStateInputAndOutput(self, wrapper_type): - keep_some = 0.9 - res = self._testDropoutWrapper( - input_keep_prob=keep_some, - output_keep_prob=keep_some, - state_keep_prob=keep_some, - variational_recurrent=True, - wrapper_type=wrapper_type, - input_size=3, - batch_size=5, - time_steps=7) - - # Smoke test for the state/input masks. - output_mask = np.abs(res[0]) > 1e-6 - for time_step in output_mask: - # Ensure the same dropout output pattern for all time steps - self.assertAllClose(output_mask[0], time_step) - for batch_entry in time_step: - # Assert all batch entries get the same mask - self.assertAllClose(batch_entry, time_step[0]) - - # For state, ensure all batch entries have the same mask - state_c_mask = np.abs(res[1].c) > 1e-6 - state_h_mask = np.abs(res[1].h) > 1e-6 - for batch_entry in state_c_mask: - self.assertAllClose(batch_entry, state_c_mask[0]) - for batch_entry in state_h_mask: - self.assertAllClose(batch_entry, state_h_mask[0]) - - @parameterized.parameters( - [rnn_cell_impl.DropoutWrapper, rnn_cell_impl.DropoutWrapperV2]) - def testDropoutWrapperRecurrentStateInputAndOutputWithSeed( - self, wrapper_type): - keep_some = 0.9 - random_seed.set_random_seed(2347) - np.random.seed(23487) - res0 = self._testDropoutWrapper( - input_keep_prob=keep_some, - output_keep_prob=keep_some, - state_keep_prob=keep_some, - variational_recurrent=True, - wrapper_type=wrapper_type, - input_size=3, - batch_size=5, - time_steps=7, - seed=-234987) - ops.reset_default_graph() - self._ClearCachedSession() - random_seed.set_random_seed(2347) - np.random.seed(23487) - res1 = self._testDropoutWrapper( - input_keep_prob=keep_some, - output_keep_prob=keep_some, - state_keep_prob=keep_some, - variational_recurrent=True, - wrapper_type=wrapper_type, - input_size=3, - batch_size=5, - time_steps=7, - seed=-234987) - - output_mask = np.abs(res0[0]) > 1e-6 - for time_step in output_mask: - # Ensure the same dropout output pattern for all time steps - self.assertAllClose(output_mask[0], time_step) - for batch_entry in time_step: - # Assert all batch entries get the same mask - self.assertAllClose(batch_entry, time_step[0]) - - # For state, ensure all batch entries have the same mask - state_c_mask = np.abs(res0[1].c) > 1e-6 - state_h_mask = np.abs(res0[1].h) > 1e-6 - for batch_entry in state_c_mask: - self.assertAllClose(batch_entry, state_c_mask[0]) - for batch_entry in state_h_mask: - self.assertAllClose(batch_entry, state_h_mask[0]) - - # Ensure seeded calculation is identical. - self.assertAllClose(res0[0], res1[0]) - self.assertAllClose(res0[1].c, res1[1].c) - self.assertAllClose(res0[1].h, res1[1].h) - - def testDropoutWrapperKerasStyle(self): - """Tests if DropoutWrapperV2 cell is instantiated in keras style scope.""" - wrapped_cell_v2 = rnn_cell_impl.DropoutWrapperV2( - rnn_cell_impl.BasicRNNCell(1)) - self.assertTrue(wrapped_cell_v2._keras_style) - - wrapped_cell = rnn_cell_impl.DropoutWrapper(rnn_cell_impl.BasicRNNCell(1)) - self.assertFalse(wrapped_cell._keras_style) - - def testDropoutWrapperV2VariableNames(self): - """Tests that variables names do not depend on wrapper in RNN layer.""" - - def _rnn_input(apply_wrapper): - """Creates a RNN layer with/without wrapper and returns built rnn cell.""" - with base_layer.keras_style_scope(): - base_cell = rnn_cell_impl.MultiRNNCell( - [rnn_cell_impl.BasicRNNCell(1) for _ in range(2)]) - if apply_wrapper: - rnn_cell = rnn_cell_impl.DropoutWrapperV2(base_cell) - else: - rnn_cell = base_cell - rnn_layer = keras_layers.RNN(rnn_cell) - inputs = ops.convert_to_tensor([[[1]]], dtype=dtypes.float32) - _ = rnn_layer(inputs) - return base_cell._cells[0] - - rnn_1 = _rnn_input(True) - ops.reset_default_graph() - rnn_2 = _rnn_input(False) - - self.assertLen(rnn_1.weights, expected_len=2) - self.assertCountEqual([v.name for v in rnn_1.weights], - [v.name for v in rnn_2.weights]) - - def testDropoutWrapperV2Caller(self): - """Tests that DropoutWrapperV2 is using the LayerRNNCell's caller.""" - - with base_layer.keras_style_scope(): - base_cell = rnn_cell_impl.MultiRNNCell( - [rnn_cell_impl.BasicRNNCell(1) for _ in range(2)]) - rnn_cell = rnn_cell_impl.DropoutWrapperV2(base_cell) - inputs = ops.convert_to_tensor([[1]], dtype=dtypes.float32) - state = ops.convert_to_tensor([[1]], dtype=dtypes.float32) - _ = rnn_cell(inputs, [state, state]) - weights = base_cell._cells[0].weights - self.assertLen(weights, expected_len=2) - self.assertTrue(all(["dropout_wrapper" in v.name for v in weights])) - - def testDropoutWrapperV2Build(self): - cell = rnn_cell_impl.LSTMCell(10) - wrapper = rnn_cell_impl.DropoutWrapperV2(cell) - wrapper.build((1,)) - self.assertTrue(cell.built) - - -def basic_rnn_cell(inputs, state, num_units, scope=None): - if state is None: - if inputs is not None: - batch_size = inputs.get_shape()[0] - dtype = inputs.dtype - else: - batch_size = 0 - dtype = dtypes.float32 - init_output = array_ops.zeros( - array_ops.stack([batch_size, num_units]), dtype=dtype) - init_state = array_ops.zeros( - array_ops.stack([batch_size, num_units]), dtype=dtype) - init_output.set_shape([batch_size, num_units]) - init_state.set_shape([batch_size, num_units]) - return init_output, init_state - else: - with variable_scope.variable_scope(scope, "basic_rnn_cell", - [inputs, state]): - output = math_ops.tanh( - Linear([inputs, state], num_units, True)([inputs, state])) - return output, output - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py b/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py index aa1d7d2b01b4595bbb03ba8e867e93db759cbd52..dfac2df6a0d4143106ad0f090805597c26659280 100644 --- a/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py +++ b/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py @@ -22,6 +22,7 @@ import itertools import numpy as np +from tensorflow.contrib.rnn.python.ops import core_rnn_cell as legacy_rnn_cell from tensorflow.contrib.rnn.python.ops import rnn_cell as contrib_rnn_cell from tensorflow.core.protobuf import config_pb2 from tensorflow.python.client import session @@ -29,7 +30,9 @@ from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import random_seed +from tensorflow.python.framework import test_util from tensorflow.python.keras import initializers +from tensorflow.python.keras import layers as keras_layers from tensorflow.python.keras import testing_utils from tensorflow.python.keras import utils from tensorflow.python.ops import array_ops @@ -51,6 +54,294 @@ from tensorflow.python.util import nest class RNNCellTest(test.TestCase): + def testIndRNNCell(self): + with self.cached_session() as sess: + with variable_scope.variable_scope( + "root", initializer=init_ops.constant_initializer(0.5)): + x = array_ops.zeros([1, 2]) + m = array_ops.zeros([1, 2]) + cell = contrib_rnn_cell.IndRNNCell(2) + g, _ = cell(x, m) + self.assertEqual([ + "root/ind_rnn_cell/%s_w:0" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME, + "root/ind_rnn_cell/%s_u:0" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME, + "root/ind_rnn_cell/%s:0" % rnn_cell_impl._BIAS_VARIABLE_NAME + ], [v.name for v in cell.trainable_variables]) + self.assertFalse(cell.non_trainable_variables) + sess.run([variables.global_variables_initializer()]) + res = sess.run([g], { + x.name: np.array([[1., 1.]]), + m.name: np.array([[0.1, 0.1]]) + }) + self.assertEqual(res[0].shape, (1, 2)) + + def testIndyGRUCell(self): + with self.cached_session() as sess: + with variable_scope.variable_scope( + "root", initializer=init_ops.constant_initializer(0.5)): + x = array_ops.zeros([1, 2]) + m = array_ops.zeros([1, 2]) + g, _ = contrib_rnn_cell.IndyGRUCell(2)(x, m) + sess.run([variables.global_variables_initializer()]) + res = sess.run([g], { + x.name: np.array([[1., 1.]]), + m.name: np.array([[0.1, 0.1]]) + }) + # Smoke test + self.assertAllClose(res[0], [[0.185265, 0.17704]]) + with variable_scope.variable_scope( + "other", initializer=init_ops.constant_initializer(0.5)): + # Test IndyGRUCell with input_size != num_units. + x = array_ops.zeros([1, 3]) + m = array_ops.zeros([1, 2]) + g, _ = contrib_rnn_cell.IndyGRUCell(2)(x, m) + sess.run([variables.global_variables_initializer()]) + res = sess.run([g], { + x.name: np.array([[1., 1., 1.]]), + m.name: np.array([[0.1, 0.1]]) + }) + # Smoke test + self.assertAllClose(res[0], [[0.155127, 0.157328]]) + + def testIndyLSTMCell(self): + for dtype in [dtypes.float16, dtypes.float32]: + np_dtype = dtype.as_numpy_dtype + with self.session(graph=ops.Graph()) as sess: + with variable_scope.variable_scope( + "root", initializer=init_ops.constant_initializer(0.5)): + x = array_ops.zeros([1, 2], dtype=dtype) + state_0 = (array_ops.zeros([1, 2], dtype=dtype),) * 2 + state_1 = (array_ops.zeros([1, 2], dtype=dtype),) * 2 + cell = rnn_cell_impl.MultiRNNCell( + [contrib_rnn_cell.IndyLSTMCell(2) for _ in range(2)]) + self.assertEqual(cell.dtype, None) + self.assertEqual("cell-0", cell._checkpoint_dependencies[0].name) + self.assertEqual("cell-1", cell._checkpoint_dependencies[1].name) + cell.get_config() # Should not throw an error + g, (out_state_0, out_state_1) = cell(x, (state_0, state_1)) + # Layer infers the input type. + self.assertEqual(cell.dtype, dtype.name) + expected_variable_names = [ + "root/multi_rnn_cell/cell_0/indy_lstm_cell/%s_w:0" % + rnn_cell_impl._WEIGHTS_VARIABLE_NAME, + "root/multi_rnn_cell/cell_0/indy_lstm_cell/%s_u:0" % + rnn_cell_impl._WEIGHTS_VARIABLE_NAME, + "root/multi_rnn_cell/cell_0/indy_lstm_cell/%s:0" % + rnn_cell_impl._BIAS_VARIABLE_NAME, + "root/multi_rnn_cell/cell_1/indy_lstm_cell/%s_w:0" % + rnn_cell_impl._WEIGHTS_VARIABLE_NAME, + "root/multi_rnn_cell/cell_1/indy_lstm_cell/%s_u:0" % + rnn_cell_impl._WEIGHTS_VARIABLE_NAME, + "root/multi_rnn_cell/cell_1/indy_lstm_cell/%s:0" % + rnn_cell_impl._BIAS_VARIABLE_NAME + ] + self.assertEqual(expected_variable_names, + [v.name for v in cell.trainable_variables]) + self.assertFalse(cell.non_trainable_variables) + sess.run([variables.global_variables_initializer()]) + res = sess.run( + [g, out_state_0, out_state_1], { + x.name: np.array([[1., 1.]]), + state_0[0].name: 0.1 * np.ones([1, 2]), + state_0[1].name: 0.1 * np.ones([1, 2]), + state_1[0].name: 0.1 * np.ones([1, 2]), + state_1[1].name: 0.1 * np.ones([1, 2]), + }) + self.assertEqual(len(res), 3) + global_variables = variables.global_variables() + self.assertEqual(expected_variable_names, + [v.name for v in global_variables]) + # Only check the range of outputs as this is just a smoke test. + self.assertAllInRange(res[0], -1.0, 1.0) + self.assertAllInRange(res[1], -1.0, 1.0) + self.assertAllInRange(res[2], -1.0, 1.0) + with variable_scope.variable_scope( + "other", initializer=init_ops.constant_initializer(0.5)): + # Test IndyLSTMCell with input_size != num_units. + x = array_ops.zeros([1, 3], dtype=dtype) + state = (array_ops.zeros([1, 2], dtype=dtype),) * 2 + g, out_state = contrib_rnn_cell.IndyLSTMCell(2)(x, state) + sess.run([variables.global_variables_initializer()]) + res = sess.run( + [g, out_state], { + x.name: np.array([[1., 1., 1.]], dtype=np_dtype), + state[0].name: 0.1 * np.ones([1, 2], dtype=np_dtype), + state[1].name: 0.1 * np.ones([1, 2], dtype=np_dtype), + }) + self.assertEqual(len(res), 2) + + def testLSTMCellLayerNorm(self): + with self.cached_session() as sess: + num_units = 2 + num_proj = 3 + batch_size = 1 + input_size = 4 + with variable_scope.variable_scope( + "root", initializer=init_ops.constant_initializer(0.5)): + x = array_ops.zeros([batch_size, input_size]) + c = array_ops.zeros([batch_size, num_units]) + h = array_ops.zeros([batch_size, num_proj]) + state = rnn_cell_impl.LSTMStateTuple(c, h) + cell = contrib_rnn_cell.LayerNormLSTMCell( + num_units=num_units, + num_proj=num_proj, + forget_bias=1.0, + layer_norm=True, + norm_gain=1.0, + norm_shift=0.0) + g, out_m = cell(x, state) + sess.run([variables.global_variables_initializer()]) + res = sess.run( + [g, out_m], { + x.name: np.ones((batch_size, input_size)), + c.name: 0.1 * np.ones((batch_size, num_units)), + h.name: 0.1 * np.ones((batch_size, num_proj)) + }) + self.assertEqual(len(res), 2) + # The numbers in results were not calculated, this is mostly just a + # smoke test. + self.assertEqual(res[0].shape, (batch_size, num_proj)) + self.assertEqual(res[1][0].shape, (batch_size, num_units)) + self.assertEqual(res[1][1].shape, (batch_size, num_proj)) + # Different inputs so different outputs and states + for i in range(1, batch_size): + self.assertTrue( + float(np.linalg.norm((res[0][0, :] - res[0][i, :]))) < 1e-6) + self.assertTrue( + float(np.linalg.norm((res[1][0, :] - res[1][i, :]))) < 1e-6) + + def testOutputProjectionWrapper(self): + with self.cached_session() as sess: + with variable_scope.variable_scope( + "root", initializer=init_ops.constant_initializer(0.5)): + x = array_ops.zeros([1, 3]) + m = array_ops.zeros([1, 3]) + cell = legacy_rnn_cell.OutputProjectionWrapper( + rnn_cell_impl.GRUCell(3), 2) + g, new_m = cell(x, m) + sess.run([variables.global_variables_initializer()]) + res = sess.run([g, new_m], { + x.name: np.array([[1., 1., 1.]]), + m.name: np.array([[0.1, 0.1, 0.1]]) + }) + self.assertEqual(res[1].shape, (1, 3)) + # The numbers in results were not calculated, this is just a smoke test. + self.assertAllClose(res[0], [[0.231907, 0.231907]]) + + def testInputProjectionWrapper(self): + with self.cached_session() as sess: + with variable_scope.variable_scope( + "root", initializer=init_ops.constant_initializer(0.5)): + x = array_ops.zeros([1, 2]) + m = array_ops.zeros([1, 3]) + cell = legacy_rnn_cell.InputProjectionWrapper( + rnn_cell_impl.GRUCell(3), num_proj=3) + g, new_m = cell(x, m) + sess.run([variables.global_variables_initializer()]) + res = sess.run([g, new_m], { + x.name: np.array([[1., 1.]]), + m.name: np.array([[0.1, 0.1, 0.1]]) + }) + self.assertEqual(res[1].shape, (1, 3)) + # The numbers in results were not calculated, this is just a smoke test. + self.assertAllClose(res[0], [[0.154605, 0.154605, 0.154605]]) + + def testEmbeddingWrapper(self): + with self.cached_session() as sess: + with variable_scope.variable_scope( + "root", initializer=init_ops.constant_initializer(0.5)): + x = array_ops.zeros([1, 1], dtype=dtypes.int32) + m = array_ops.zeros([1, 2]) + embedding_cell = legacy_rnn_cell.EmbeddingWrapper( + rnn_cell_impl.GRUCell(2), embedding_classes=3, embedding_size=2) + self.assertEqual(embedding_cell.output_size, 2) + g, new_m = embedding_cell(x, m) + sess.run([variables.global_variables_initializer()]) + res = sess.run([g, new_m], { + x.name: np.array([[1]]), + m.name: np.array([[0.1, 0.1]]) + }) + self.assertEqual(res[1].shape, (1, 2)) + # The numbers in results were not calculated, this is just a smoke test. + self.assertAllClose(res[0], [[0.17139, 0.17139]]) + + def testEmbeddingWrapperWithDynamicRnn(self): + with self.cached_session() as sess: + with variable_scope.variable_scope("root"): + inputs = ops.convert_to_tensor([[[0], [0]]], dtype=dtypes.int64) + input_lengths = ops.convert_to_tensor([2], dtype=dtypes.int64) + embedding_cell = legacy_rnn_cell.EmbeddingWrapper( + rnn_cell_impl.BasicLSTMCell(1, state_is_tuple=True), + embedding_classes=1, + embedding_size=2) + outputs, _ = rnn.dynamic_rnn( + cell=embedding_cell, + inputs=inputs, + sequence_length=input_lengths, + dtype=dtypes.float32) + sess.run([variables.global_variables_initializer()]) + # This will fail if output's dtype is inferred from input's. + sess.run(outputs) + + def testSRUCell(self): + with self.cached_session() as sess: + with variable_scope.variable_scope( + "root", initializer=init_ops.constant_initializer(0.5)): + x = array_ops.zeros([1, 2]) + m = array_ops.zeros([1, 2]) + g, _ = contrib_rnn_cell.SRUCell(2)(x, m) + sess.run([variables.global_variables_initializer()]) + res = sess.run([g], { + x.name: np.array([[1., 1.]]), + m.name: np.array([[0.1, 0.1]]) + }) + # Smoke test + self.assertAllClose(res[0], [[0.509682, 0.509682]]) + + def testSRUCellKerasRNN(self): + """Tests that SRUCell works with keras RNN layer.""" + cell = contrib_rnn_cell.SRUCell(10) + seq_input = ops.convert_to_tensor( + np.random.rand(2, 3, 5), name="seq_input", dtype=dtypes.float32) + rnn_layer = keras_layers.RNN(cell=cell) + rnn_outputs_keras = rnn_layer(seq_input) + with self.cached_session() as sess: + sess.run([variables.global_variables_initializer()]) + self.assertEqual(sess.run(rnn_outputs_keras).shape, (2, 10)) + + def testSRUCellBiasType(self): + """Tests that the bias' dtype is properly set.""" + cell = contrib_rnn_cell.SRUCell(10) + cell.build((2, 3, 5)) + self.assertEqual(cell._bias.dtype, dtypes.float32_ref) + + cell = contrib_rnn_cell.SRUCell(10, dtype=dtypes.int32) + cell.build((2, 3, 5)) + self.assertEqual(cell._bias.dtype, dtypes.int32_ref) + + cell_input = ops.convert_to_tensor( + np.random.rand(2, 5), name="cell_input", dtype=dtypes.float16) + cell_state = ops.convert_to_tensor( + np.random.rand(2, 10), name="cell_state", dtype=dtypes.float16) + cell = contrib_rnn_cell.SRUCell(10) + cell(cell_input, [cell_state]) + self.assertEqual(cell._bias.dtype, dtypes.float16_ref) + + def testSRUCellWithDiffSize(self): + with self.cached_session() as sess: + with variable_scope.variable_scope( + "root", initializer=init_ops.constant_initializer(0.5)): + x = array_ops.zeros([1, 3]) + m = array_ops.zeros([1, 2]) + g, _ = contrib_rnn_cell.SRUCell(2)(x, m) + sess.run([variables.global_variables_initializer()]) + res = sess.run([g], { + x.name: np.array([[1., 1., 1.]]), + m.name: np.array([[0.1, 0.1]]) + }) + # Smoke test + self.assertAllClose(res[0], [[0.55255556, 0.55255556]]) + def testCoupledInputForgetGateLSTMCell(self): with self.cached_session() as sess: num_units = 2 @@ -763,6 +1054,17 @@ class RNNCellTest(test.TestCase): self.assertEqual(new_h.shape[1], num_proj) self.assertAllClose(np.concatenate(res[1], axis=1), expected_state) + @test_util.run_in_graph_and_eager_modes + def testNASCellKerasRNN(self): + """Tests that NASCell works with keras RNN layer.""" + cell = contrib_rnn_cell.NASCell(10) + seq_input = ops.convert_to_tensor( + np.random.rand(2, 3, 5), name="seq_input", dtype=dtypes.float32) + rnn_layer = keras_layers.RNN(cell=cell) + rnn_outputs = rnn_layer(seq_input) + self.evaluate([variables.global_variables_initializer()]) + self.assertEqual(self.evaluate(rnn_outputs).shape, (2, 10)) + def testUGRNNCell(self): num_units = 2 batch_size = 3 diff --git a/tensorflow/contrib/rnn/python/ops/rnn.py b/tensorflow/contrib/rnn/python/ops/rnn.py index 0266b72dcb15e4aba01a9a31b4be75c5b84d44da..41b1698321e20f4360d75fa2db79f9bd8a806cea 100644 --- a/tensorflow/contrib/rnn/python/ops/rnn.py +++ b/tensorflow/contrib/rnn/python/ops/rnn.py @@ -131,7 +131,8 @@ def stack_bidirectional_dynamic_rnn(cells_fw, sequence_length=None, parallel_iterations=None, time_major=False, - scope=None): + scope=None, + swap_memory=False): """Creates a dynamic bidirectional recurrent neural network. Stacks several bidirectional rnn layers. The combined forward and backward @@ -171,6 +172,10 @@ def stack_bidirectional_dynamic_rnn(cells_fw, data is batch-major, so by default this function accepts input and emits output in batch-major form. scope: VariableScope for the created subgraph; defaults to None. + swap_memory: Transparently swap the tensors produced in forward inference + but needed for back prop from GPU to CPU. This allows training RNNs + which would typically not fit on a single GPU, with very minimal (or no) + performance penalty. Returns: A tuple (outputs, output_state_fw, output_state_bw) where: @@ -230,6 +235,7 @@ def stack_bidirectional_dynamic_rnn(cells_fw, sequence_length=sequence_length, parallel_iterations=parallel_iterations, dtype=dtype, + swap_memory=swap_memory, time_major=time_major) # Concat the outputs to create the new input. prev_layer = array_ops.concat(outputs, 2) diff --git a/tensorflow/contrib/rnn/python/ops/rnn_cell.py b/tensorflow/contrib/rnn/python/ops/rnn_cell.py index 8a1c09f171e6108174671e3122d5ff4c0b236003..482e547a16be85804beec88a91fa03b053d09b27 100644 --- a/tensorflow/contrib/rnn/python/ops/rnn_cell.py +++ b/tensorflow/contrib/rnn/python/ops/rnn_cell.py @@ -1462,7 +1462,7 @@ class LayerNormBasicLSTMCell(rnn_cell_impl.RNNCell): return new_h, new_state -class NASCell(rnn_cell_impl.RNNCell): +class NASCell(rnn_cell_impl.LayerRNNCell): """Neural Architecture Search (NAS) recurrent network cell. This implements the recurrent cell from the paper: @@ -1475,23 +1475,28 @@ class NASCell(rnn_cell_impl.RNNCell): The class uses an optional projection layer. """ - def __init__(self, num_units, num_proj=None, use_biases=False, reuse=None): + # NAS cell's architecture base. + _NAS_BASE = 8 + + def __init__(self, num_units, num_proj=None, use_bias=False, reuse=None, + **kwargs): """Initialize the parameters for a NAS cell. Args: - num_units: int, The number of units in the NAS cell + num_units: int, The number of units in the NAS cell. num_proj: (optional) int, The output dimensionality for the projection matrices. If None, no projection is performed. - use_biases: (optional) bool, If True then use biases within the cell. This + use_bias: (optional) bool, If True then use biases within the cell. This is False by default. reuse: (optional) Python boolean describing whether to reuse variables in an existing scope. If not `True`, and the existing scope already has the given variables, an error is raised. + **kwargs: Additional keyword arguments. """ - super(NASCell, self).__init__(_reuse=reuse) + super(NASCell, self).__init__(_reuse=reuse, **kwargs) self._num_units = num_units self._num_proj = num_proj - self._use_biases = use_biases + self._use_bias = use_bias self._reuse = reuse if num_proj is not None: @@ -1509,6 +1514,33 @@ class NASCell(rnn_cell_impl.RNNCell): def output_size(self): return self._output_size + def build(self, inputs_shape): + input_size = tensor_shape.dimension_value( + tensor_shape.TensorShape(inputs_shape).with_rank(2)[1]) + if input_size is None: + raise ValueError("Could not infer input size from inputs.get_shape()[-1]") + + num_proj = self._num_units if self._num_proj is None else self._num_proj + + # Variables for the NAS cell. `recurrent_kernel` is all matrices multiplying + # the hiddenstate and `kernel` is all matrices multiplying the inputs. + self.recurrent_kernel = self.add_variable( + "recurrent_kernel", [num_proj, self._NAS_BASE * self._num_units]) + self.kernel = self.add_variable( + "kernel", [input_size, self._NAS_BASE * self._num_units]) + + if self._use_bias: + self.bias = self.add_variable("bias", + shape=[self._NAS_BASE * self._num_units], + initializer=init_ops.zeros_initializer) + + # Projection layer if specified + if self._num_proj is not None: + self.projection_weights = self.add_variable( + "projection_weights", [self._num_units, self._num_proj]) + + self.built = True + def call(self, inputs, state): """Run one step of NAS Cell. @@ -1535,38 +1567,20 @@ class NASCell(rnn_cell_impl.RNNCell): tanh = math_ops.tanh relu = nn_ops.relu - num_proj = self._num_units if self._num_proj is None else self._num_proj - (c_prev, m_prev) = state - dtype = inputs.dtype - input_size = inputs.get_shape().with_rank(2).dims[1] - if input_size.value is None: - raise ValueError("Could not infer input size from inputs.get_shape()[-1]") - # Variables for the NAS cell. W_m is all matrices multiplying the - # hiddenstate and W_inputs is all matrices multiplying the inputs. - concat_w_m = vs.get_variable("recurrent_kernel", - [num_proj, 8 * self._num_units], dtype) - concat_w_inputs = vs.get_variable( - "kernel", [input_size.value, 8 * self._num_units], dtype) - - m_matrix = math_ops.matmul(m_prev, concat_w_m) - inputs_matrix = math_ops.matmul(inputs, concat_w_inputs) - - if self._use_biases: - b = vs.get_variable( - "bias", - shape=[8 * self._num_units], - initializer=init_ops.zeros_initializer(), - dtype=dtype) - m_matrix = nn_ops.bias_add(m_matrix, b) + m_matrix = math_ops.matmul(m_prev, self.recurrent_kernel) + inputs_matrix = math_ops.matmul(inputs, self.kernel) + + if self._use_bias: + m_matrix = nn_ops.bias_add(m_matrix, self.bias) # The NAS cell branches into 8 different splits for both the hiddenstate # and the input m_matrix_splits = array_ops.split( - axis=1, num_or_size_splits=8, value=m_matrix) + axis=1, num_or_size_splits=self._NAS_BASE, value=m_matrix) inputs_matrix_splits = array_ops.split( - axis=1, num_or_size_splits=8, value=inputs_matrix) + axis=1, num_or_size_splits=self._NAS_BASE, value=inputs_matrix) # First layer layer1_0 = sigmoid(inputs_matrix_splits[0] + m_matrix_splits[0]) @@ -1598,9 +1612,7 @@ class NASCell(rnn_cell_impl.RNNCell): # Projection layer if specified if self._num_proj is not None: - concat_w_proj = vs.get_variable("projection_weights", - [self._num_units, self._num_proj], dtype) - new_m = math_ops.matmul(new_m, concat_w_proj) + new_m = math_ops.matmul(new_m, self.projection_weights) new_state = rnn_cell_impl.LSTMStateTuple(new_c, new_m) return new_m, new_state @@ -2071,7 +2083,7 @@ class ConvLSTMCell(rnn_cell_impl.RNNCell): conv_ndims: Convolution dimensionality (1, 2 or 3). input_shape: Shape of the input as int tuple, excluding the batch size. output_channels: int, number of output channels of the conv LSTM. - kernel_shape: Shape of kernel as in tuple (of size 1,2 or 3). + kernel_shape: Shape of kernel as an int tuple (of size 1, 2 or 3). use_bias: (bool) Use bias in convolutions. skip_connection: If set to `True`, concatenate the input to the output of the conv LSTM. Default: `False`. @@ -2092,7 +2104,7 @@ class ConvLSTMCell(rnn_cell_impl.RNNCell): self._conv_ndims = conv_ndims self._input_shape = input_shape self._output_channels = output_channels - self._kernel_shape = kernel_shape + self._kernel_shape = list(kernel_shape) self._use_bias = use_bias self._forget_bias = forget_bias self._skip_connection = skip_connection @@ -2172,7 +2184,7 @@ def _conv(args, filter_size, num_features, bias, bias_start=0.0): Args: args: a Tensor or a list of Tensors of dimension 3D, 4D or 5D, batch x n, Tensors. - filter_size: int tuple of filter height and width. + filter_size: int tuple of filter shape (of size 1, 2 or 3). num_features: int, number of features. bias: Whether to use biases in the convolution layer. bias_start: starting value to initialize the bias; 0 by default. @@ -2744,10 +2756,12 @@ class SRUCell(rnn_cell_impl.LayerRNNCell): name: (optional) String, the name of the layer. Layers with the same name will share weights, but to avoid mistakes we require reuse=True in such cases. + **kwargs: Additional keyword arguments. """ - def __init__(self, num_units, activation=None, reuse=None, name=None): - super(SRUCell, self).__init__(_reuse=reuse, name=name) + def __init__(self, num_units, activation=None, reuse=None, name=None, + **kwargs): + super(SRUCell, self).__init__(_reuse=reuse, name=name, **kwargs) self._num_units = num_units self._activation = activation or math_ops.tanh @@ -2777,7 +2791,7 @@ class SRUCell(rnn_cell_impl.LayerRNNCell): self._bias = self.add_variable( rnn_cell_impl._BIAS_VARIABLE_NAME, # pylint: disable=protected-access shape=[2 * self._num_units], - initializer=init_ops.constant_initializer(0.0, dtype=self.dtype)) + initializer=init_ops.zeros_initializer) self._built = True diff --git a/tensorflow/contrib/rpc/python/kernel_tests/rpc_op_test.py b/tensorflow/contrib/rpc/python/kernel_tests/rpc_op_test.py index 3fc6bfbb4d03a39906d4441e48b2788423caa234..d8ab9eba7049e468b373a1641f92dc781aa22558 100644 --- a/tensorflow/contrib/rpc/python/kernel_tests/rpc_op_test.py +++ b/tensorflow/contrib/rpc/python/kernel_tests/rpc_op_test.py @@ -61,10 +61,7 @@ class RpcOpTest(test.TestCase, rpc_op_test_base.RpcOpTestBase): self._server = server def tearDown(self): - # TODO(ebrevdo): Figure out why this sometimes times out. - # self._service.ExitLoop() - # self._service_thread.join() - # self._server.stop() + self._server.stop(grace=None) super(RpcOpTest, self).tearDown() diff --git a/tensorflow/contrib/saved_model/BUILD b/tensorflow/contrib/saved_model/BUILD index 269443b2c6508bb618d30f64487b1a6a84e8646f..f0242a3b40fd566ec0f477d462426d5f550d1620 100644 --- a/tensorflow/contrib/saved_model/BUILD +++ b/tensorflow/contrib/saved_model/BUILD @@ -84,35 +84,6 @@ py_library( srcs_version = "PY2AND3", visibility = ["//visibility:public"], deps = [ - "//tensorflow/python:array_ops", - "//tensorflow/python:framework_ops", - "//tensorflow/python:lib", - "//tensorflow/python:metrics", - "//tensorflow/python:platform", - "//tensorflow/python:saver", - "//tensorflow/python:util", - "//tensorflow/python/estimator:estimator_py", - "//tensorflow/python/keras:engine", - "//tensorflow/python/saved_model", - ], -) - -py_test( - name = "keras_saved_model_test", - size = "medium", - srcs = ["python/saved_model/keras_saved_model_test.py"], - srcs_version = "PY2AND3", - tags = [ - "no_oss", # TODO(b/119349471): Re-enable - "no_windows", - ], - deps = [ - ":keras_saved_model", - "//tensorflow/python:client_testlib", - "//tensorflow/python:training", - "//tensorflow/python/estimator:estimator_py", "//tensorflow/python/keras", - "//third_party/py/numpy", - "@absl_py//absl/testing:parameterized", ], ) 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 2a4b6eae367fe617e9a19d80f16eb3fda9ade1c0..0392ed9eee79391c60318faf68d8dfd6eb64a994 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 @@ -18,398 +18,9 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import os -import six +from tensorflow.python.keras import saving -from tensorflow.python.client import session -from tensorflow.python.framework import ops -from tensorflow.python.keras import backend as K -from tensorflow.python.keras import models as models_lib -from tensorflow.python.keras import optimizers -from tensorflow.python.keras.engine import sequential -from tensorflow.python.keras.engine import training_utils -from tensorflow.python.keras.metrics import Metric -from tensorflow.python.keras.models import model_from_json -from tensorflow.python.lib.io import file_io -from tensorflow.python.ops import variables -from tensorflow.python.platform import tf_logging as logging -from tensorflow.python.saved_model import builder as saved_model_builder -from tensorflow.python.saved_model import constants -from tensorflow.python.saved_model import save as save_lib -from tensorflow.python.saved_model import utils_impl as saved_model_utils -from tensorflow.python.training import saver as saver_lib -from tensorflow.python.training.checkpointable import util as checkpointable_utils -from tensorflow.python.util import compat -from tensorflow.python.util import nest -from tensorflow_estimator.python.estimator import keras as estimator_keras_util -from tensorflow_estimator.python.estimator import model_fn as model_fn_lib -from tensorflow_estimator.python.estimator.export import export as export_helpers - -def save_keras_model( - model, saved_model_path, custom_objects=None, as_text=None, - input_signature=None, serving_only=False): - """Saves a `tf.keras.Model` into Tensorflow SavedModel format. - - `save_model` generates new files/folders under the `saved_model_path` folder: - 1) a checkpoint containing the model weights. - 2) a saved_model.pb file containing the model's MetaGraphs. The prediction - graph is always exported. The evaluaton and training graphs are exported - if the following conditions are met: - - Evaluation: model loss is defined. - - Training: model is compiled with an optimizer defined under `tf.train`. - This is because `tf.keras.optimizers.Optimizer` instances cannot be - saved to checkpoints. - 3) Model's json configuration, if model.get_config() has been implemented. - This file can be used to reload the model using - tf.keras.models.model_from_json(). Note that if any custom objects were - used, they should be passed to the `custom_object` argument when loading - the model. - - Model limitations: - - Sequential and functional models can always be saved. - - Subclassed models can only be saved when `serving_only=True`. This is due to - the current implementation copying the model in order to export the training - and evaluation graphs. Because the topology of subclassed models cannot be - determined, the subclassed models cannot be cloned. Subclassed models will - be entirely exportable in the future. - - Note that each mode is exported in separate graphs, so different modes do not - share variables. To use the train graph with evaluation or prediction graphs, - create a new checkpoint if variable values have been updated. - - Example: - - ```python - import tensorflow as tf - - # Create a tf.keras model. - model = tf.keras.Sequential() - model.add(tf.keras.layers.Dense(1, input_shape=[10])) - model.summary() - - # Save the tf.keras model in the SavedModel format. - saved_to_path = tf.contrib.saved_model.save_keras_model( - model, '/tmp/my_simple_tf_keras_saved_model') - - # Load the saved keras model back. - model_prime = tf.contrib.saved_model.load_keras_model(saved_to_path) - model_prime.summary() - ``` - - Args: - model: A `tf.keras.Model` to be saved. If the model is subclassed, the flag - `serving_only` must be set to True. - saved_model_path: a string specifying the path to the SavedModel directory. - The SavedModel will be saved to a timestamped folder created within this - directory. - custom_objects: Optional dictionary mapping string names to custom classes - or functions (e.g. custom loss functions). - as_text: whether to write the `SavedModel` proto in text format. Currently - unavailable in serving-only mode. - input_signature: A possibly nested sequence of `tf.TensorSpec` objects, used - to specify the expected model inputs. `input_signature`'s nested structure - should match the expected nested structure of the inputs to the model. If - this is not set, this function will attempt to infer the input shapes and - dtypes from the model. Note that if the model is subclassed, the tensor - inputs to the call function should be nested in the first argument (this - is a general requirement for using subclassed models with Keras functions - .fit(), .predict(), etc.). - serving_only: Export only the outputs produced from calling the model in - predict mode. The losses, optimizer, and other training configurations are - not saved. If the SavedModel will only be used for serving (rather than - retraining), or if the model is subclassed, this can be set to True. - - Returns: - String path to the SavedModel folder, a subdirectory of `saved_model_path`. - - Raises: - NotImplementedError: If the model is a subclassed model, and serving_only is - False. - ValueError: If the input signature cannot be inferred from the model. - """ - export_dir = export_helpers.get_timestamped_export_dir(saved_model_path) - - if serving_only: - save_lib.save( - model, export_dir, - signatures=training_utils.trace_model_call(model, input_signature)) - else: - _save_v1_format(model, export_dir, custom_objects, as_text, input_signature) - - try: - _export_model_json(model, export_dir) - except NotImplementedError: - logging.warning('Skipped saving model JSON, subclassed model does not have ' - 'get_config() defined.') - - return export_dir - - -def _export_model_json(model, saved_model_path): - """Saves model configuration as a json string under assets folder.""" - model_json = model.to_json() - model_json_filepath = os.path.join( - saved_model_utils.get_or_create_assets_dir(saved_model_path), - compat.as_text(constants.SAVED_MODEL_FILENAME_JSON)) - file_io.write_string_to_file(model_json_filepath, model_json) - - -def _export_model_variables(model, saved_model_path): - """Saves model weights in checkpoint format under variables folder.""" - saved_model_utils.get_or_create_variables_dir(saved_model_path) - checkpoint_prefix = saved_model_utils.get_variables_path(saved_model_path) - model.save_weights(checkpoint_prefix, save_format='tf', overwrite=True) - return checkpoint_prefix - - -def _save_v1_format(model, path, custom_objects, as_text, input_signature): - """Exports model to v1 SavedModel format.""" - if not model._is_graph_network: - if isinstance(model, sequential.Sequential): - # If input shape is not directly set in the model, the exported model - # will infer the expected shapes of the input from the model. - if not model.built and input_signature is None: - raise ValueError( - 'Sequential model\'s input shape is unknown. Please build the ' - 'model, or use the input_signature argument to specify the ' - 'model inputs.') - else: - raise NotImplementedError( - 'Subclassed models can only be exported for serving. Please set ' - 'argument serving_only=True.') - - builder = saved_model_builder._SavedModelBuilder(path) - - # Manually save variables to export them in an object-based checkpoint. This - # skips the `builder.add_meta_graph_and_variables()` step, which saves a - # named-based checkpoint. - # TODO(b/113134168): Add fn to Builder to save with object-based saver. - # TODO(b/113178242): This should only export the model json structure. Only - # one save is needed once the weights can be copied from the model to clone. - checkpoint_path = _export_model_variables(model, path) - - # Export each mode. Use ModeKeys enums defined for `Estimator` to ensure that - # Keras models and `Estimator`s are exported with the same format. - # Every time a mode is exported, the code checks to see if new variables have - # been created (e.g. optimizer slot variables). If that is the case, the - # checkpoint is re-saved to include the new variables. - export_args = {'builder': builder, - 'model': model, - 'custom_objects': custom_objects, - 'checkpoint_path': checkpoint_path, - 'input_signature': input_signature} - - has_saved_vars = False - if model.optimizer: - # TODO(kathywu): Verify this works with v2 optimizer. - if isinstance(model.optimizer, optimizers.TFOptimizer): - _export_mode(model_fn_lib.ModeKeys.TRAIN, has_saved_vars, **export_args) - has_saved_vars = True - _export_mode(model_fn_lib.ModeKeys.EVAL, has_saved_vars, **export_args) - else: - logging.warning( - 'Model was compiled with an optimizer, but the optimizer is not from ' - '`tf.train` (e.g. `tf.train.AdagradOptimizer`). Only the serving ' - 'graph was exported. The train and evaluate graphs were not added to ' - 'the SavedModel.') - _export_mode(model_fn_lib.ModeKeys.PREDICT, has_saved_vars, **export_args) - - builder.save(as_text) - - -def _get_var_list(model): - """Returns list of all checkpointed saveable objects in the model.""" - return checkpointable_utils.named_saveables(model) - - -def create_placeholder(spec): - return K.placeholder(shape=spec.shape, dtype=spec.dtype, name=spec.name) - - -def _export_mode( - mode, has_saved_vars, builder, model, custom_objects, checkpoint_path, - input_signature): - """Exports a model, and optionally saves new vars from the clone model. - - Args: - mode: A `tf.estimator.ModeKeys` string. - has_saved_vars: A `boolean` indicating whether the SavedModel has already - exported variables. - builder: A `SavedModelBuilder` object. - model: A `tf.keras.Model` object. - custom_objects: A dictionary mapping string names to custom classes - or functions. - checkpoint_path: String path to checkpoint. - input_signature: Nested TensorSpec containing the expected inputs. Can be - `None`, in which case the signature will be inferred from the model. - - Raises: - ValueError: If the train/eval mode is being exported, but the model does - not have an optimizer. - """ - compile_clone = (mode != model_fn_lib.ModeKeys.PREDICT) - if compile_clone and not model.optimizer: - raise ValueError( - 'Model does not have an optimizer. Cannot export mode %s' % mode) - - model_graph = ops.get_default_graph() - with ops.Graph().as_default() as g: - - K.set_learning_phase(mode == model_fn_lib.ModeKeys.TRAIN) - - if input_signature is None: - input_tensors = None - else: - input_tensors = nest.map_structure(create_placeholder, input_signature) - - # Clone the model into blank graph. This will create placeholders for inputs - # and targets. - clone = models_lib.clone_and_build_model( - model, input_tensors=input_tensors, custom_objects=custom_objects, - compile_clone=compile_clone) - - # Make sure that iterations variable is added to the global step collection, - # to ensure that, when the SavedModel graph is loaded, the iterations - # variable is returned by `tf.train.get_global_step()`. This is required for - # compatibility with the SavedModelEstimator. - if compile_clone: - g.add_to_collection(ops.GraphKeys.GLOBAL_STEP, clone.optimizer.iterations) - - # Extract update and train ops from train/test/predict functions. - train_op = None - if mode == model_fn_lib.ModeKeys.TRAIN: - clone._make_train_function() - train_op = clone.train_function.updates_op - elif mode == model_fn_lib.ModeKeys.EVAL: - clone._make_test_function() - else: - clone._make_predict_function() - g.get_collection_ref(ops.GraphKeys.UPDATE_OPS).extend(clone.state_updates) - - clone_var_list = checkpointable_utils.named_saveables(clone) - - with session.Session().as_default(): - if has_saved_vars: - # Confirm all variables in the clone have an entry in the checkpoint. - status = clone.load_weights(checkpoint_path) - status.assert_existing_objects_matched() - else: - # Confirm that variables between the clone and model match up exactly, - # not counting optimizer objects. Optimizer objects are ignored because - # if the model has not trained, the slot variables will not have been - # created yet. - # TODO(b/113179535): Replace with checkpointable equivalence. - _assert_same_non_optimizer_objects(model, model_graph, clone, g) - - # TODO(b/113178242): Use value transfer for checkpointable objects. - clone.load_weights(checkpoint_path) - - # Add graph and variables to SavedModel. - # TODO(b/113134168): Switch to add_meta_graph_and_variables. - clone.save_weights(checkpoint_path, save_format='tf', overwrite=True) - builder._has_saved_variables = True - - # Add graph to the SavedModel builder. - builder.add_meta_graph( - model_fn_lib.EXPORT_TAG_MAP[mode], - signature_def_map=_create_signature_def_map(clone, mode), - saver=saver_lib.Saver(clone_var_list), - init_op=variables.local_variables_initializer(), - train_op=train_op) - return None - - -def _create_signature_def_map(model, mode): - """Creates a SignatureDef map from a Keras model.""" - inputs_dict = {name: x for name, x in zip(model.input_names, model.inputs)} - if model.optimizer: - targets_dict = {x.name.split(':')[0]: x - for x in model.targets if x is not None} - inputs_dict.update(targets_dict) - outputs_dict = {name: x - for name, x in zip(model.output_names, model.outputs)} - metrics = estimator_keras_util._convert_keras_metrics_to_estimator(model) - - # Add metric variables to the `LOCAL_VARIABLES` collection. Metric variables - # are by default not added to any collections. We are doing this here, so - # that metric variables get initialized. - local_vars = set(ops.get_collection(ops.GraphKeys.LOCAL_VARIABLES)) - vars_to_add = set() - if metrics is not None: - for key, value in six.iteritems(metrics): - if isinstance(value, Metric): - vars_to_add.update(value.variables) - # Convert Metric instances to (value_tensor, update_op) tuple. - metrics[key] = (value.result(), value.updates[0]) - # Remove variables that are in the local variables collection already. - vars_to_add = vars_to_add.difference(local_vars) - for v in vars_to_add: - ops.add_to_collection(ops.GraphKeys.LOCAL_VARIABLES, v) - - export_outputs = model_fn_lib.export_outputs_for_mode( - mode, - predictions=outputs_dict, - loss=model.total_loss if model.optimizer else None, - metrics=metrics) - return export_helpers.build_all_signature_defs( - inputs_dict, - export_outputs=export_outputs, - serving_only=(mode == model_fn_lib.ModeKeys.PREDICT)) - - -def _assert_same_non_optimizer_objects(model, model_graph, clone, clone_graph): # pylint: disable=unused-argument - """Asserts model and clone contain the same checkpointable objects.""" - - # TODO(fchollet, kathywu): make sure this works in eager mode. - return True - - -def load_keras_model(saved_model_path): - """Loads a keras.Model from SavedModel. - - load_model reinstantiates model state by: - 1) loading model topology from json (this will eventually come - from metagraph). - 2) loading model weights from checkpoint. - - Example: - - ```python - import tensorflow as tf - - # Create a tf.keras model. - model = tf.keras.Sequential() - model.add(tf.keras.layers.Dense(1, input_shape=[10])) - model.summary() - - # Save the tf.keras model in the SavedModel format. - saved_to_path = tf.contrib.saved_model.save_keras_model( - model, '/tmp/my_simple_tf_keras_saved_model') - - # Load the saved keras model back. - model_prime = tf.contrib.saved_model.load_keras_model(saved_to_path) - model_prime.summary() - ``` - - Args: - saved_model_path: a string specifying the path to an existing SavedModel. - - Returns: - a keras.Model instance. - """ - # restore model topology from json string - model_json_filepath = os.path.join( - compat.as_bytes(saved_model_path), - compat.as_bytes(constants.ASSETS_DIRECTORY), - compat.as_bytes(constants.SAVED_MODEL_FILENAME_JSON)) - model_json = file_io.read_file_to_string(model_json_filepath) - model = model_from_json(model_json) - - # restore model weights - checkpoint_prefix = os.path.join( - compat.as_text(saved_model_path), - compat.as_text(constants.VARIABLES_DIRECTORY), - compat.as_text(constants.VARIABLES_FILENAME)) - model.load_weights(checkpoint_prefix) - return model +# TODO(kathywu): Remove all contrib callers, switch to tf.keras. +save_keras_model = saving.export +load_keras_model = saving.load_from_saved_model diff --git a/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model_test.py b/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model_test.py deleted file mode 100644 index fbf8138493362d4a3c8a75e1ee1bb2fbe8096499..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model_test.py +++ /dev/null @@ -1,538 +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. -# ============================================================================== -# pylint: disable=protected-access -"""Tests for saving/loading function for keras Model.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import os -import shutil - -from absl.testing import parameterized -import numpy as np - -from tensorflow.contrib.saved_model.python.saved_model import keras_saved_model -from tensorflow.python import keras -from tensorflow.python.client import session -from tensorflow.python.eager import context -from tensorflow.python.estimator import model_fn as model_fn_lib -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_spec -from tensorflow.python.framework import test_util -from tensorflow.python.keras.engine import training -from tensorflow.python.keras.utils import tf_utils -from tensorflow.python.ops import array_ops -from tensorflow.python.platform import test -from tensorflow.python.saved_model import loader_impl -from tensorflow.python.saved_model import signature_constants -from tensorflow.python.training import training as training_module - - -class TestModelSavingandLoading(test.TestCase): - - def _save_model_dir(self, dirname='saved_model'): - temp_dir = self.get_temp_dir() - self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True) - return os.path.join(temp_dir, dirname) - - def test_saving_sequential_model(self): - with self.cached_session(): - model = keras.models.Sequential() - model.add(keras.layers.Dense(2, input_shape=(3,))) - model.add(keras.layers.RepeatVector(3)) - model.add(keras.layers.TimeDistributed(keras.layers.Dense(3))) - model.compile( - loss=keras.losses.MSE, - optimizer=keras.optimizers.RMSprop(lr=0.0001), - metrics=[keras.metrics.categorical_accuracy], - sample_weight_mode='temporal') - x = np.random.random((1, 3)) - y = np.random.random((1, 3, 3)) - model.train_on_batch(x, y) - - ref_y = model.predict(x) - - temp_saved_model = self._save_model_dir() - output_path = keras_saved_model.save_keras_model(model, temp_saved_model) - - loaded_model = keras_saved_model.load_keras_model(output_path) - y = loaded_model.predict(x) - self.assertAllClose(ref_y, y, atol=1e-05) - - @test_util.run_in_graph_and_eager_modes - def test_saving_sequential_model_without_compile(self): - with self.cached_session(): - model = keras.models.Sequential() - model.add(keras.layers.Dense(2, input_shape=(3,))) - model.add(keras.layers.RepeatVector(3)) - model.add(keras.layers.TimeDistributed(keras.layers.Dense(3))) - - x = np.random.random((1, 3)) - ref_y = model.predict(x) - - temp_saved_model = self._save_model_dir() - output_path = keras_saved_model.save_keras_model(model, temp_saved_model) - loaded_model = keras_saved_model.load_keras_model(output_path) - - y = loaded_model.predict(x) - self.assertAllClose(ref_y, y, atol=1e-05) - - def test_saving_functional_model(self): - with self.cached_session(): - inputs = keras.layers.Input(shape=(3,)) - x = keras.layers.Dense(2)(inputs) - output = keras.layers.Dense(3)(x) - - model = keras.models.Model(inputs, output) - model.compile( - loss=keras.losses.MSE, - optimizer=keras.optimizers.RMSprop(lr=0.0001), - metrics=[keras.metrics.categorical_accuracy]) - x = np.random.random((1, 3)) - y = np.random.random((1, 3)) - model.train_on_batch(x, y) - - ref_y = model.predict(x) - - temp_saved_model = self._save_model_dir() - output_path = keras_saved_model.save_keras_model(model, temp_saved_model) - loaded_model = keras_saved_model.load_keras_model(output_path) - - y = loaded_model.predict(x) - self.assertAllClose(ref_y, y, atol=1e-05) - - @test_util.run_in_graph_and_eager_modes - def test_saving_functional_model_without_compile(self): - with self.cached_session(): - inputs = keras.layers.Input(shape=(3,)) - x = keras.layers.Dense(2)(inputs) - output = keras.layers.Dense(3)(x) - - model = keras.models.Model(inputs, output) - - x = np.random.random((1, 3)) - y = np.random.random((1, 3)) - - ref_y = model.predict(x) - - temp_saved_model = self._save_model_dir() - output_path = keras_saved_model.save_keras_model(model, temp_saved_model) - loaded_model = keras_saved_model.load_keras_model(output_path) - - y = loaded_model.predict(x) - self.assertAllClose(ref_y, y, atol=1e-05) - - @test_util.run_in_graph_and_eager_modes - def test_saving_with_tf_optimizer(self): - with self.cached_session(): - model = keras.models.Sequential() - model.add(keras.layers.Dense(2, input_shape=(3,))) - model.add(keras.layers.Dense(3)) - model.compile( - loss='mse', - optimizer=training_module.RMSPropOptimizer(0.1), - metrics=['acc']) - - x = np.random.random((1, 3)) - y = np.random.random((1, 3)) - model.train_on_batch(x, y) - ref_y = model.predict(x) - - temp_saved_model = self._save_model_dir() - output_path = keras_saved_model.save_keras_model(model, temp_saved_model) - loaded_model = keras_saved_model.load_keras_model(output_path) - loaded_model.compile( - loss='mse', - optimizer=training_module.RMSPropOptimizer(0.1), - metrics=['acc']) - y = loaded_model.predict(x) - self.assertAllClose(ref_y, y, atol=1e-05) - - # test that new updates are the same with both models - x = np.random.random((1, 3)) - y = np.random.random((1, 3)) - - ref_loss = model.train_on_batch(x, y) - loss = loaded_model.train_on_batch(x, y) - self.assertAllClose(ref_loss, loss, atol=1e-05) - - ref_y = model.predict(x) - y = loaded_model.predict(x) - self.assertAllClose(ref_y, y, atol=1e-05) - - # test saving/loading again - temp_saved_model2 = self._save_model_dir('saved_model_2') - output_path2 = keras_saved_model.save_keras_model( - loaded_model, temp_saved_model2) - loaded_model = keras_saved_model.load_keras_model(output_path2) - y = loaded_model.predict(x) - self.assertAllClose(ref_y, y, atol=1e-05) - - def test_saving_subclassed_model_raise_error(self): - # For now, saving subclassed model should raise an error. It should be - # avoided later with loading from SavedModel.pb. - - class SubclassedModel(training.Model): - - def __init__(self): - super(SubclassedModel, self).__init__() - self.layer1 = keras.layers.Dense(3) - self.layer2 = keras.layers.Dense(1) - - def call(self, inp): - return self.layer2(self.layer1(inp)) - - model = SubclassedModel() - - temp_saved_model = self._save_model_dir() - with self.assertRaises(NotImplementedError): - keras_saved_model.save_keras_model(model, temp_saved_model) - - -class LayerWithLearningPhase(keras.engine.base_layer.Layer): - - def call(self, x): - phase = keras.backend.learning_phase() - output = tf_utils.smart_cond( - phase, lambda: x * 0, lambda: array_ops.identity(x)) - if not context.executing_eagerly(): - output._uses_learning_phase = True # pylint: disable=protected-access - return output - - def compute_output_shape(self, input_shape): - return input_shape - - -def functional_model(uses_learning_phase=True): - inputs = keras.layers.Input(shape=(3,)) - x = keras.layers.Dense(2)(inputs) - x = keras.layers.Dense(3)(x) - if uses_learning_phase: - x = LayerWithLearningPhase()(x) - return keras.models.Model(inputs, x) - - -def sequential_model(uses_learning_phase=True): - model = keras.models.Sequential() - model.add(keras.layers.Dense(2, input_shape=(3,))) - model.add(keras.layers.Dense(3)) - if uses_learning_phase: - model.add(LayerWithLearningPhase()) - return model - - -def sequential_model_without_input_shape(uses_learning_phase=True): - model = keras.models.Sequential() - model.add(keras.layers.Dense(2)) - model.add(keras.layers.Dense(3)) - if uses_learning_phase: - model.add(LayerWithLearningPhase()) - return model - - -class Subclassed(keras.models.Model): - - def __init__(self): - super(Subclassed, self).__init__() - self.dense1 = keras.layers.Dense(2) - self.dense2 = keras.layers.Dense(3) - - def call(self, inputs): - x = self.dense1(inputs) - x = self.dense2(x) - return x - - -def subclassed_model(): - return Subclassed() - - -def load_model(sess, path, mode): - tags = model_fn_lib.EXPORT_TAG_MAP[mode] - if mode == model_fn_lib.ModeKeys.PREDICT: - sig_def_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY - else: - sig_def_key = mode - - meta_graph_def = loader_impl.load(sess, tags, path) - inputs = { - k: sess.graph.get_tensor_by_name(v.name) - for k, v in meta_graph_def.signature_def[sig_def_key].inputs.items()} - outputs = { - k: sess.graph.get_tensor_by_name(v.name) - for k, v in meta_graph_def.signature_def[sig_def_key].outputs.items()} - return inputs, outputs, meta_graph_def - - -@test_util.run_all_in_graph_and_eager_modes -class TestModelSavedModelExport(test.TestCase, parameterized.TestCase): - - def _save_model_dir(self, dirname='saved_model'): - temp_dir = self.get_temp_dir() - self.addCleanup(shutil.rmtree, temp_dir, ignore_errors=True) - return os.path.join(temp_dir, dirname) - - @parameterized.parameters( - { - 'model_builder': functional_model, - 'uses_learning_phase': True, - 'optimizer': training_module.AdadeltaOptimizer(), - 'train_before_export': True}, - { - 'model_builder': functional_model, - 'uses_learning_phase': True, - 'optimizer': training_module.AdadeltaOptimizer(), - 'train_before_export': False}, - { - 'model_builder': functional_model, - 'uses_learning_phase': False, - 'optimizer': None, - 'train_before_export': False}, - { - 'model_builder': sequential_model, - 'uses_learning_phase': True, - 'optimizer': training_module.AdadeltaOptimizer(), - 'train_before_export': True}, - { - 'model_builder': sequential_model, - 'uses_learning_phase': True, - 'optimizer': training_module.AdadeltaOptimizer(), - 'train_before_export': False}, - { - 'model_builder': sequential_model, - 'uses_learning_phase': False, - 'optimizer': None, - 'train_before_export': False}, - { - 'model_builder': sequential_model_without_input_shape, - 'uses_learning_phase': True, - 'optimizer': training_module.AdadeltaOptimizer(), - 'train_before_export': False}) - def testSaveAndLoadSavedModelExport( - self, model_builder, uses_learning_phase, optimizer, train_before_export): - saved_model_path = self._save_model_dir() - with self.session(graph=ops.Graph()): - np.random.seed(130) - input_arr = np.random.random((1, 3)) - target_arr = np.random.random((1, 3)) - - model = model_builder(uses_learning_phase) - if optimizer is not None: - model.compile( - loss='mse', - optimizer=optimizer, - metrics=['mae']) - if train_before_export: - model.train_on_batch(input_arr, target_arr) - - ref_loss, ref_mae = model.evaluate(input_arr, target_arr) - - ref_predict = model.predict(input_arr) - - # Export SavedModel - output_path = keras_saved_model.save_keras_model(model, saved_model_path) - - input_name = model.input_names[0] - output_name = model.output_names[0] - target_name = output_name + '_target' - - # Load predict graph, and test predictions - with session.Session(graph=ops.Graph()) as sess: - inputs, outputs, _ = load_model(sess, output_path, - model_fn_lib.ModeKeys.PREDICT) - - predictions = sess.run(outputs[output_name], - {inputs[input_name]: input_arr}) - self.assertAllClose(ref_predict, predictions, atol=1e-05) - - if optimizer: - # Load eval graph, and test predictions, loss and metric values - with session.Session(graph=ops.Graph()) as sess: - inputs, outputs, _ = load_model(sess, output_path, - model_fn_lib.ModeKeys.EVAL) - - # First obtain the loss and predictions, and run the metric update op by - # feeding in the inputs and targets. - loss, predictions, _ = sess.run( - (outputs['loss'], outputs['predictions/' + output_name], - outputs['metrics/mean_absolute_error/update_op']), { - inputs[input_name]: input_arr, - inputs[target_name]: target_arr - }) - - # The metric value should be run after the update op, to ensure that it - # reflects the correct value. - metric_value = sess.run(outputs['metrics/mean_absolute_error/value']) - - self.assertEqual(int(train_before_export), - sess.run(training_module.get_global_step())) - self.assertAllClose(ref_loss, loss, atol=1e-05) - self.assertAllClose(ref_mae, metric_value, atol=1e-05) - self.assertAllClose(ref_predict, predictions, atol=1e-05) - - # Load train graph, and check for the train op, and prediction values - with session.Session(graph=ops.Graph()) as sess: - inputs, outputs, meta_graph_def = load_model( - sess, output_path, model_fn_lib.ModeKeys.TRAIN) - self.assertEqual(int(train_before_export), - sess.run(training_module.get_global_step())) - self.assertIn('loss', outputs) - self.assertIn('metrics/mean_absolute_error/update_op', outputs) - self.assertIn('metrics/mean_absolute_error/value', outputs) - self.assertIn('predictions/' + output_name, outputs) - - # Train for a step - train_op = loader_impl.get_train_op(meta_graph_def) - train_outputs, _ = sess.run( - [outputs, train_op], {inputs[input_name]: input_arr, - inputs[target_name]: target_arr}) - self.assertEqual(int(train_before_export) + 1, - sess.run(training_module.get_global_step())) - - if uses_learning_phase: - self.assertAllClose( - [[0, 0, 0]], train_outputs['predictions/' + output_name], - atol=1e-05) - else: - self.assertNotAllClose( - [[0, 0, 0]], train_outputs['predictions/' + output_name], - atol=1e-05) - - def testSaveAndLoadSavedModelWithCustomObject(self): - saved_model_path = self._save_model_dir() - with session.Session(graph=ops.Graph()) as sess: - def relu6(x): - return keras.backend.relu(x, max_value=6) - inputs = keras.layers.Input(shape=(1,)) - outputs = keras.layers.Activation(relu6)(inputs) - model = keras.models.Model(inputs, outputs) - output_path = keras_saved_model.save_keras_model( - model, saved_model_path, custom_objects={'relu6': relu6}) - with session.Session(graph=ops.Graph()) as sess: - inputs, outputs, _ = load_model(sess, output_path, - model_fn_lib.ModeKeys.PREDICT) - input_name = model.input_names[0] - output_name = model.output_names[0] - predictions = sess.run( - outputs[output_name], {inputs[input_name]: [[7], [-3], [4]]}) - self.assertAllEqual([[6], [0], [4]], predictions) - - def testAssertModelCloneSameObjectsIgnoreOptimizer(self): - input_arr = np.random.random((1, 3)) - target_arr = np.random.random((1, 3)) - - model_graph = ops.Graph() - clone_graph = ops.Graph() - - # Create two models with the same layers but different optimizers. - with session.Session(graph=model_graph): - inputs = keras.layers.Input(shape=(3,)) - x = keras.layers.Dense(2)(inputs) - x = keras.layers.Dense(3)(x) - model = keras.models.Model(inputs, x) - - model.compile(loss='mse', optimizer=training_module.AdadeltaOptimizer()) - model.train_on_batch(input_arr, target_arr) - - with session.Session(graph=clone_graph): - inputs = keras.layers.Input(shape=(3,)) - x = keras.layers.Dense(2)(inputs) - x = keras.layers.Dense(3)(x) - clone = keras.models.Model(inputs, x) - clone.compile(loss='mse', optimizer=keras.optimizers.RMSprop(lr=0.0001)) - clone.train_on_batch(input_arr, target_arr) - - keras_saved_model._assert_same_non_optimizer_objects( - model, model_graph, clone, clone_graph) - - def testAssertModelCloneSameObjectsThrowError(self): - input_arr = np.random.random((1, 3)) - target_arr = np.random.random((1, 3)) - - model_graph = ops.Graph() - clone_graph = ops.Graph() - - # Create two models with the same layers but different optimizers. - with session.Session(graph=model_graph): - inputs = keras.layers.Input(shape=(3,)) - x = keras.layers.Dense(2)(inputs) - x = keras.layers.Dense(3)(x) - model = keras.models.Model(inputs, x) - - model.compile(loss='mse', optimizer=training_module.AdadeltaOptimizer()) - model.train_on_batch(input_arr, target_arr) - - with session.Session(graph=clone_graph): - inputs = keras.layers.Input(shape=(3,)) - x = keras.layers.Dense(2)(inputs) - x = keras.layers.Dense(4)(x) - x = keras.layers.Dense(3)(x) - clone = keras.models.Model(inputs, x) - clone.compile(loss='mse', optimizer=keras.optimizers.RMSprop(lr=0.0001)) - clone.train_on_batch(input_arr, target_arr) - - def testSaveSequentialModelWithoutInputShapes(self): - model = sequential_model_without_input_shape(True) - # A Sequential model that hasn't been built should raise an error. - with self.assertRaisesRegexp(ValueError, 'Please build the model'): - keras_saved_model.save_keras_model(model, '') - - saved_model_path = self._save_model_dir() - output_path = keras_saved_model.save_keras_model( - model, saved_model_path, - input_signature=tensor_spec.TensorSpec(shape=(10, 11, 12, 13, 14), - dtype=dtypes.float32, - name='spec_input')) - - with session.Session(graph=ops.Graph()) as sess: - inputs, outputs, _ = load_model(sess, output_path, - model_fn_lib.ModeKeys.PREDICT) - self.assertEqual(5, inputs[next(iter(inputs.keys()))].shape.ndims) - self.assertEqual(5, outputs[next(iter(outputs.keys()))].shape.ndims) - self.assertEqual(3, outputs[next(iter(outputs.keys()))].shape[-1]) - - @test_util.run_v2_only - @parameterized.parameters( - { - 'model_builder': sequential_model_without_input_shape, - 'input_signature': [tensor_spec.TensorSpec(shape=[None, 3], - dtype=dtypes.float32)]}, - { - 'model_builder': subclassed_model, - 'input_signature': [tensor_spec.TensorSpec(shape=[None, 3], - dtype=dtypes.float32)]}) - def testServingOnly(self, model_builder, input_signature): - saved_model_path = self._save_model_dir() - input_arr = np.random.random((5, 3)).astype(np.float32) - model = model_builder() - ref_predict = model.predict(input_arr) - - output_path = keras_saved_model.save_keras_model( - model, saved_model_path, serving_only=True, - input_signature=input_signature) - - # Load predict graph, and test predictions - with session.Session(graph=ops.Graph()) as sess: - inputs, outputs, _ = load_model(sess, output_path, - model_fn_lib.ModeKeys.PREDICT) - predictions = sess.run(outputs[next(iter(outputs.keys()))], - {inputs[next(iter(inputs.keys()))]: input_arr}) - self.assertAllClose(ref_predict, predictions, atol=1e-05) - - -if __name__ == '__main__': - test.main() diff --git a/tensorflow/contrib/seq2seq/BUILD b/tensorflow/contrib/seq2seq/BUILD index 18b56cd21942e28cb0dc3210df0bb04d55c1e16f..2a70b08f5c46e11e7fd83fe134741b9a241153f5 100644 --- a/tensorflow/contrib/seq2seq/BUILD +++ b/tensorflow/contrib/seq2seq/BUILD @@ -33,7 +33,6 @@ tf_custom_op_py_library( srcs_version = "PY2AND3", deps = [ ":beam_search_ops", - "//tensorflow/contrib/distributions:distributions_py", "//tensorflow/contrib/layers:layers_py", "//tensorflow/contrib/rnn:rnn_py", "//tensorflow/contrib/util:util_py", @@ -59,7 +58,6 @@ tf_custom_op_py_library( "//tensorflow/python:tensor_util", "//tensorflow/python:util", "//tensorflow/python:variable_scope", - "//tensorflow/python/ops/distributions", "//third_party/py/numpy", "@six_archive//:six", ], @@ -141,6 +139,27 @@ cuda_py_test( ], ) +cuda_py_test( + name = "basic_decoder_v2_test", + size = "medium", + srcs = ["python/kernel_tests/basic_decoder_v2_test.py"], + additional_deps = [ + ":seq2seq_py", + "@absl_py//absl/testing:parameterized", + "//tensorflow/contrib/layers:layers_py", + "//tensorflow/contrib/rnn:rnn_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:init_ops", + "//tensorflow/python:platform_test", + "//tensorflow/python:rnn", + "//tensorflow/python:variable_scope", + "//tensorflow/python:variables", + ], +) + cuda_py_test( name = "beam_search_ops_test", size = "medium", @@ -175,6 +194,27 @@ cuda_py_test( ], ) +cuda_py_test( + name = "decoder_v2_test", + size = "medium", + srcs = ["python/kernel_tests/decoder_v2_test.py"], + additional_deps = [ + ":seq2seq_py", + "@absl_py//absl/testing:parameterized", + "//tensorflow/contrib/layers:layers_py", + "//tensorflow/contrib/rnn:rnn_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:init_ops", + "//tensorflow/python:platform_test", + "//tensorflow/python:rnn", + "//tensorflow/python:variable_scope", + "//tensorflow/python:variables", + ], +) + cuda_py_test( name = "beam_search_decoder_test", size = "medium", @@ -215,3 +255,18 @@ cuda_py_test( "//tensorflow/python:variables", ], ) + +cuda_py_test( + name = "attention_wrapper_v2_test", + size = "medium", + srcs = ["python/kernel_tests/attention_wrapper_v2_test.py"], + additional_deps = [ + ":seq2seq_py", + "@absl_py//absl/testing:parameterized", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_test_lib", + "//tensorflow/python:init_ops", + "//tensorflow/python:platform_test", + "//tensorflow/python:variables", + ], +) diff --git a/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py b/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py index d815f81f847ad79ddcc6c6ecf5c050598e185d8d..1a5692f7b5be5e87b78dac9d1ae51f280ca089f8 100644 --- a/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py @@ -358,7 +358,7 @@ class AttentionWrapperTest(test.TestCase): rnn_output=ResultSummary( shape=(5, 3, 6), dtype=dtype('float32'), mean=-0.00597103), sample_id=ResultSummary( - shape=(5, 3), dtype=dtype('int32'), mean=1.6)) + shape=(5, 3), dtype=dtype('int32'), mean=1.4)) expected_final_state = AttentionWrapperState( cell_state=LSTMStateTuple( c=ResultSummary( @@ -387,7 +387,7 @@ class AttentionWrapperTest(test.TestCase): rnn_output=ResultSummary( shape=(5, 3, 6), dtype=dtype('float32'), mean=-0.0052615386), sample_id=ResultSummary( - shape=(5, 3), dtype=dtype('int32'), mean=1.3333333333)) + shape=(5, 3), dtype=dtype('int32'), mean=1.4)) expected_final_state = AttentionWrapperState( cell_state=LSTMStateTuple( c=ResultSummary( @@ -454,7 +454,7 @@ class AttentionWrapperTest(test.TestCase): rnn_output=ResultSummary( shape=(5, 3, 6), dtype=dtype('float32'), mean=-0.0052615386), sample_id=ResultSummary( - shape=(5, 3), dtype=dtype('int32'), mean=1.3333333333333333)) + shape=(5, 3), dtype=dtype('int32'), mean=1.4)) expected_final_state = AttentionWrapperState( cell_state=LSTMStateTuple( c=ResultSummary( @@ -696,7 +696,7 @@ class AttentionWrapperTest(test.TestCase): rnn_output=ResultSummary( shape=(5, 3, 6), dtype=dtype('float32'), mean=-0.0025896581), sample_id=ResultSummary( - shape=(5, 3), dtype=dtype('int32'), mean=1.6)) + shape=(5, 3), dtype=dtype('int32'), mean=1.73333333)) expected_final_state = AttentionWrapperState( cell_state=LSTMStateTuple( c=ResultSummary( @@ -707,12 +707,12 @@ class AttentionWrapperTest(test.TestCase): shape=(5, 6), dtype=dtype('float32'), mean=-0.00069823361), time=3, alignments=ResultSummary( - shape=(5, 8), dtype=dtype('float32'), mean=0.028698336), + shape=(5, 8), dtype=dtype('float32'), mean=0.029914695), attention_state=ResultSummary( - shape=(5, 8), dtype=dtype('float32'), mean=0.028698336), + shape=(5, 8), dtype=dtype('float32'), mean=0.029914695), alignment_history=()) expected_final_alignment_history = ResultSummary( - shape=(3, 5, 8), dtype=dtype('float32'), mean=0.04865776002407074) + shape=(3, 5, 8), dtype=dtype('float32'), mean=0.0465225502849) self._testWithAttention( create_attention_mechanism, @@ -921,9 +921,9 @@ class AttentionWrapperTest(test.TestCase): expected_final_output = BasicDecoderOutput( rnn_output=ResultSummary( - shape=(5, 3, 20), dtype=dtype('float32'), mean=0.11723966), + shape=(5, 3, 20), dtype=dtype('float32'), mean=0.115853324533), sample_id=ResultSummary( - shape=(5, 3), dtype=dtype('int32'), mean=7.266666666666667)) + shape=(5, 3), dtype=dtype('int32'), mean=8.6)) expected_final_state = AttentionWrapperState( cell_state=LSTMStateTuple( c=ResultSummary( @@ -931,7 +931,7 @@ class AttentionWrapperTest(test.TestCase): h=ResultSummary( shape=(5, 9), dtype=dtype('float32'), mean=-0.0018327223)), attention=ResultSummary( - shape=(5, 20), dtype=dtype('float32'), mean=0.11601614207), + shape=(5, 20), dtype=dtype('float32'), mean=0.11462739855), time=3, alignments=(ResultSummary( shape=(5, 8), dtype=dtype('float32'), mean=0.125), diff --git a/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_v2_test.py b/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_v2_test.py new file mode 100644 index 0000000000000000000000000000000000000000..5ee01f66f165bd2ac22cae10807f24f6b97f0c64 --- /dev/null +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_v2_test.py @@ -0,0 +1,745 @@ +# 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. +# ============================================================================== +"""Tests for contrib.seq2seq.python.ops.attention_wrapper.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import collections + +from absl.testing import parameterized +import numpy as np + +from tensorflow.contrib.seq2seq.python.ops import attention_wrapper as wrapper +from tensorflow.contrib.seq2seq.python.ops import basic_decoder +from tensorflow.contrib.seq2seq.python.ops import sampler as sampler_py +from tensorflow.python import keras +from tensorflow.python.eager import context +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import test_util +from tensorflow.python.keras import initializers +from tensorflow.python.ops import rnn_cell +from tensorflow.python.ops import variables +from tensorflow.python.platform import test +from tensorflow.python.util import nest + + +@test_util.run_all_in_graph_and_eager_modes +class AttentionMechanismTest(test.TestCase, parameterized.TestCase): + + def setUp(self): + super(AttentionMechanismTest, self).setUp() + self.batch = 10 + self.timestep = 5 + self.memory_size = 6 + self.units = 8 + + self.memory = np.random.randn(self.batch, self.timestep, + self.memory_size).astype(np.float32) + self.query = np.random.randn(self.batch, self.units).astype(np.float32) + self.state = np.random.randn(self.batch, self.timestep).astype(np.float32) + + @parameterized.named_parameters( + ("luong", wrapper.LuongAttentionV2), + ("luong_monotonic", wrapper.LuongMonotonicAttentionV2), + ("bahdanau", wrapper.BahdanauAttentionV2), + ("bahdanau_monotonic", wrapper.BahdanauMonotonicAttentionV2), + ) + def test_attention_shape_inference(self, attention_cls): + attention = attention_cls(self.units, self.memory) + attention_score = attention([self.query, self.state]) + self.assertLen(attention_score, 2) + self.assertEqual(attention_score[0].shape, (self.batch, self.timestep)) + self.assertEqual(attention_score[1].shape, (self.batch, self.timestep)) + + @parameterized.named_parameters( + ("luong", wrapper.LuongAttentionV2), + ("luong_monotonic", wrapper.LuongMonotonicAttentionV2), + ("bahdanau", wrapper.BahdanauAttentionV2), + ("bahdanau_monotonic", wrapper.BahdanauMonotonicAttentionV2), + ) + def test_get_config(self, attention_cls): + attention = attention_cls(self.units, self.memory) + config = attention.get_config() + + attention_from_config = attention_cls.from_config(config) + config_from_clone = attention_from_config.get_config() + + self.assertDictEqual(config, config_from_clone) + + @parameterized.named_parameters( + ("luong", wrapper.LuongAttentionV2), + ("luong_monotonic", wrapper.LuongMonotonicAttentionV2), + ("bahdanau", wrapper.BahdanauAttentionV2), + ("bahdanau_monotonic", wrapper.BahdanauMonotonicAttentionV2), + ) + def test_layer_output(self, attention_cls): + attention = attention_cls(self.units, self.memory) + score = attention([self.query, self.state]) + self.evaluate(variables.variables_initializer(attention.variables)) + + score_val = self.evaluate(score) + self.assertLen(score_val, 2) + self.assertEqual(score_val[0].shape, (self.batch, self.timestep)) + self.assertEqual(score_val[1].shape, (self.batch, self.timestep)) + + @parameterized.named_parameters( + ("luong", wrapper.LuongAttentionV2), + ("luong_monotonic", wrapper.LuongMonotonicAttentionV2), + ("bahdanau", wrapper.BahdanauAttentionV2), + ("bahdanau_monotonic", wrapper.BahdanauMonotonicAttentionV2), + ) + def test_passing_memory_from_call(self, attention_cls): + attention = attention_cls(self.units, self.memory) + weights_before_query = attention.get_weights() + ref_score = attention([self.query, self.state]) + + self.evaluate(variables.global_variables_initializer()) + ref_score_val = self.evaluate(ref_score) + + all_weights = attention.get_weights() + config = attention.get_config() + # Simulate the twice invocation of calls here. + attention_from_config = attention_cls.from_config(config) + attention_from_config.build(self.memory.shape) + attention_from_config.set_weights(weights_before_query) + attention_from_config(self.memory, setup_memory=True) + attention_from_config.build([self.query.shape, self.state.shape]) + attention_from_config.set_weights(all_weights) + score = attention_from_config([self.query, self.state]) + + score_val = self.evaluate(score) + self.assertAllClose(ref_score_val, score_val) + + @parameterized.named_parameters( + ("luong", wrapper.LuongAttentionV2), + ("luong_monotonic", wrapper.LuongMonotonicAttentionV2), + ("bahdanau", wrapper.BahdanauAttentionV2), + ("bahdanau_monotonic", wrapper.BahdanauMonotonicAttentionV2), + ) + def test_save_load_layer(self, attention_cls): + vocab = 20 + embedding_dim = 6 + inputs = keras.layers.Input(shape=[self.timestep]) + encoder_input = keras.layers.Embedding( + vocab, embedding_dim, mask_zero=True)( + inputs) + encoder_output = keras.layers.UnifiedLSTM( + self.memory_size, return_sequences=True)( + encoder_input) + + attention = attention_cls(self.units, encoder_output) + query = keras.layers.Input(shape=[self.units]) + state = keras.layers.Input(shape=[self.timestep]) + + score = attention([query, state]) + + x = np.random.randint(vocab, size=(self.batch, self.timestep)) + x_test = np.random.randint(vocab, size=(self.batch, self.timestep)) + y = np.random.randn(self.batch, self.timestep) + model = keras.models.Model([inputs, query, state], score) + model.compile("rmsprop", "mse") + model.fit([x, self.query, self.state], (y, y)) + y_ref = model.predict_on_batch([x_test, self.query, self.state]) + + config = model.get_config() + weights = model.get_weights() + loaded_model = keras.models.Model.from_config( + config, custom_objects={attention_cls.__name__: attention_cls}) + loaded_model.set_weights(weights) + + y = loaded_model.predict_on_batch([x_test, self.query, self.state]) + + self.assertAllClose(y_ref, y) + + # TODO(scottzhu): Add tests for model.compile(run_eagerly=True) + + +class ResultSummary( + collections.namedtuple("ResultSummary", ("shape", "dtype", "mean"))): + pass + + +def get_result_summary(x): + if isinstance(x, np.ndarray): + return ResultSummary(x.shape, x.dtype, x.mean()) + return x + + +@test_util.run_all_in_graph_and_eager_modes +class AttentionWrapperV2Test(test.TestCase, parameterized.TestCase): + + def assertAllCloseOrEqual(self, x, y, **kwargs): + if isinstance(x, np.ndarray) or isinstance(x, float): + return super(AttentionWrapperV2Test, self).assertAllClose( + x, y, atol=1e-3, **kwargs) + else: + self.assertAllEqual(x, y, **kwargs) + + def setUp(self): + super(AttentionWrapperV2Test, self).setUp() + self.batch = 64 + self.units = 128 + self.encoder_timestep = 10 + self.encoder_dim = 256 + self.decoder_timestep = 12 + self.encoder_outputs = np.random.randn(self.batch, self.encoder_timestep, + self.encoder_dim) + self.encoder_sequence_length = np.random.randint( + self.encoder_timestep, size=(self.batch,)).astype(np.int32) + self.decoder_inputs = np.random.randn(self.batch, self.decoder_timestep, + self.units) + self.decoder_sequence_length = np.random.randint( + self.decoder_timestep, size=(self.batch,)).astype(np.int32) + + def _testWithAttention(self, + create_attention_mechanism, + expected_final_output, + expected_final_state, + attention_mechanism_depth=3, + alignment_history=False, + expected_final_alignment_history=None, + attention_layer_size=6, + attention_layer=None, + create_query_layer=False, + create_memory_layer=True, + create_attention_kwargs=None): + attention_layer_sizes = ([attention_layer_size] + if attention_layer_size is not None else None) + attention_layers = ([attention_layer] + if attention_layer is not None else None) + self._testWithMaybeMultiAttention( + is_multi=False, + create_attention_mechanisms=[create_attention_mechanism], + expected_final_output=expected_final_output, + expected_final_state=expected_final_state, + attention_mechanism_depths=[attention_mechanism_depth], + alignment_history=alignment_history, + expected_final_alignment_history=expected_final_alignment_history, + attention_layer_sizes=attention_layer_sizes, + attention_layers=attention_layers, + create_query_layer=create_query_layer, + create_memory_layer=create_memory_layer, + create_attention_kwargs=create_attention_kwargs) + + def _testWithMaybeMultiAttention(self, + is_multi, + create_attention_mechanisms, + expected_final_output, + expected_final_state, + attention_mechanism_depths, + alignment_history=False, + expected_final_alignment_history=None, + attention_layer_sizes=None, + attention_layers=None, + create_query_layer=False, + create_memory_layer=True, + create_attention_kwargs=None): + # Allow is_multi to be True with a single mechanism to enable test for + # passing in a single mechanism in a list. + assert len(create_attention_mechanisms) == 1 or is_multi + encoder_sequence_length = [3, 2, 3, 1, 1] + decoder_sequence_length = [2, 0, 1, 2, 3] + batch_size = 5 + encoder_max_time = 8 + decoder_max_time = 4 + input_depth = 7 + encoder_output_depth = 10 + cell_depth = 9 + create_attention_kwargs = create_attention_kwargs or {} + + if attention_layer_sizes is not None: + # Compute sum of attention_layer_sizes. Use encoder_output_depth if None. + attention_depth = sum(attention_layer_size or encoder_output_depth + for attention_layer_size in attention_layer_sizes) + elif attention_layers is not None: + # Compute sum of attention_layers output depth. + attention_depth = sum( + attention_layer.compute_output_shape( + [batch_size, cell_depth + encoder_output_depth]).dims[-1].value + for attention_layer in attention_layers) + else: + attention_depth = encoder_output_depth * len(create_attention_mechanisms) + + decoder_inputs = np.random.randn(batch_size, decoder_max_time, + input_depth).astype(np.float32) + encoder_outputs = np.random.randn(batch_size, encoder_max_time, + encoder_output_depth).astype(np.float32) + + attention_mechanisms = [] + for creator, depth in zip(create_attention_mechanisms, + attention_mechanism_depths): + # Create a memory layer with deterministic initializer to avoid randomness + # in the test between graph and eager. + if create_query_layer: + create_attention_kwargs["query_layer"] = keras.layers.Dense( + depth, kernel_initializer="ones", use_bias=False) + if create_memory_layer: + create_attention_kwargs["memory_layer"] = keras.layers.Dense( + depth, kernel_initializer="ones", use_bias=False) + + attention_mechanisms.append( + creator( + units=depth, + memory=encoder_outputs, + memory_sequence_length=encoder_sequence_length, + **create_attention_kwargs)) + + with self.cached_session(use_gpu=True): + attention_layer_size = attention_layer_sizes + attention_layer = attention_layers + if not is_multi: + if attention_layer_size is not None: + attention_layer_size = attention_layer_size[0] + if attention_layer is not None: + attention_layer = attention_layer[0] + cell = rnn_cell.LSTMCell(cell_depth, initializer="ones") + cell = wrapper.AttentionWrapper( + cell, + attention_mechanisms if is_multi else attention_mechanisms[0], + attention_layer_size=attention_layer_size, + alignment_history=alignment_history, + attention_layer=attention_layer) + # Set the attention_layer within AttentionWrapper to have deterministic + # kernel initializer, for testing purpose. + if cell._attention_layers is not None: + for layer in cell._attention_layers: + if getattr(layer, "kernel_initializer") is None: + layer.kernel_initializer = initializers.ones() + + sampler = sampler_py.TrainingSampler() + my_decoder = basic_decoder.BasicDecoderV2(cell=cell, sampler=sampler) + initial_state = cell.zero_state( + dtype=dtypes.float32, batch_size=batch_size) + final_outputs, final_state, _ = my_decoder( + decoder_inputs, + initial_state=initial_state, + sequence_length=decoder_sequence_length) + + self.assertIsInstance(final_outputs, basic_decoder.BasicDecoderOutput) + self.assertIsInstance(final_state, wrapper.AttentionWrapperState) + self.assertIsInstance(final_state.cell_state, rnn_cell.LSTMStateTuple) + + expected_time = ( + expected_final_state.time if context.executing_eagerly() else None) + self.assertEqual((batch_size, expected_time, attention_depth), + tuple(final_outputs.rnn_output.get_shape().as_list())) + self.assertEqual((batch_size, expected_time), + tuple(final_outputs.sample_id.get_shape().as_list())) + + self.assertEqual((batch_size, attention_depth), + tuple(final_state.attention.get_shape().as_list())) + self.assertEqual((batch_size, cell_depth), + tuple(final_state.cell_state.c.get_shape().as_list())) + self.assertEqual((batch_size, cell_depth), + tuple(final_state.cell_state.h.get_shape().as_list())) + + if alignment_history: + if is_multi: + state_alignment_history = [] + for history_array in final_state.alignment_history: + history = history_array.stack() + self.assertEqual((expected_time, batch_size, encoder_max_time), + tuple(history.get_shape().as_list())) + state_alignment_history.append(history) + state_alignment_history = tuple(state_alignment_history) + else: + state_alignment_history = final_state.alignment_history.stack() + self.assertEqual((expected_time, batch_size, encoder_max_time), + tuple(state_alignment_history.get_shape().as_list())) + nest.assert_same_structure(cell.state_size, + cell.zero_state(batch_size, dtypes.float32)) + # Remove the history from final_state for purposes of the + # remainder of the tests. + final_state = final_state._replace(alignment_history=()) # pylint: disable=protected-access + else: + state_alignment_history = () + + self.evaluate(variables.global_variables_initializer()) + eval_result = self.evaluate({ + "final_outputs": final_outputs, + "final_state": final_state, + "state_alignment_history": state_alignment_history, + }) + + final_output_info = nest.map_structure(get_result_summary, + eval_result["final_outputs"]) + final_state_info = nest.map_structure(get_result_summary, + eval_result["final_state"]) + print("final_output_info: ", final_output_info) + print("final_state_info: ", final_state_info) + + nest.map_structure(self.assertAllCloseOrEqual, expected_final_output, + final_output_info) + nest.map_structure(self.assertAllCloseOrEqual, expected_final_state, + final_state_info) + if alignment_history: # by default, the wrapper emits attention as output + final_alignment_history_info = nest.map_structure( + get_result_summary, eval_result["state_alignment_history"]) + print("final_alignment_history_info: ", final_alignment_history_info) + nest.map_structure( + self.assertAllCloseOrEqual, + # outputs are batch major but the stacked TensorArray is time major + expected_final_alignment_history, + final_alignment_history_info) + + @parameterized.parameters([np.float16, np.float32, np.float64]) + def _testBahdanauNormalizedDType(self, dtype): + encoder_outputs = self.encoder_outputs.astype(dtype) + decoder_inputs = self.decoder_inputs.astype(dtype) + attention_mechanism = wrapper.BahdanauAttentionV2( + units=self.units, + memory=encoder_outputs, + memory_sequence_length=self.encoder_sequence_length, + normalize=True, + dtype=dtype) + cell = rnn_cell.LSTMCell(self.units) + cell = wrapper.AttentionWrapper(cell, attention_mechanism) + + sampler = sampler_py.TrainingSampler() + my_decoder = basic_decoder.BasicDecoderV2(cell=cell, sampler=sampler) + + final_outputs, final_state, _ = my_decoder( + decoder_inputs, + initial_state=cell.zero_state(dtype=dtype, batch_size=self.batch), + sequence_length=self.decoder_sequence_length) + self.assertIsInstance(final_outputs, basic_decoder.BasicDecoderOutput) + self.assertEqual(final_outputs.rnn_output.dtype, dtype) + self.assertIsInstance(final_state, wrapper.AttentionWrapperState) + self.assertIsInstance(final_state.cell_state, rnn_cell.LSTMStateTuple) + + @parameterized.parameters([np.float16, np.float32, np.float64]) + def testLuongScaledDType(self, dtype): + # Test case for GitHub issue 18099 + encoder_outputs = self.encoder_outputs.astype(dtype) + decoder_inputs = self.decoder_inputs.astype(dtype) + attention_mechanism = wrapper.LuongAttentionV2( + units=self.units, + memory=encoder_outputs, + memory_sequence_length=self.encoder_sequence_length, + scale=True, + dtype=dtype, + ) + cell = rnn_cell.LSTMCell(self.units) + cell = wrapper.AttentionWrapper(cell, attention_mechanism) + + sampler = sampler_py.TrainingSampler() + my_decoder = basic_decoder.BasicDecoderV2(cell=cell, sampler=sampler) + + final_outputs, final_state, _ = my_decoder( + decoder_inputs, + initial_state=cell.zero_state(dtype=dtype, batch_size=self.batch), + sequence_length=self.decoder_sequence_length) + self.assertIsInstance(final_outputs, basic_decoder.BasicDecoderOutput) + self.assertEqual(final_outputs.rnn_output.dtype, dtype) + self.assertIsInstance(final_state, wrapper.AttentionWrapperState) + self.assertIsInstance(final_state.cell_state, rnn_cell.LSTMStateTuple) + + def testBahdanauNotNormalized(self): + create_attention_mechanism = wrapper.BahdanauAttentionV2 + create_attention_kwargs = {"kernel_initializer": "ones"} + expected_final_output = basic_decoder.BasicDecoderOutput( + rnn_output=ResultSummary( + shape=(5, 3, 6), dtype=np.dtype(np.float32), mean=4.8290324), + sample_id=ResultSummary(shape=(5, 3), dtype=np.dtype(np.int32), mean=0)) + expected_final_state = wrapper.AttentionWrapperState( + cell_state=rnn_cell.LSTMStateTuple( + c=ResultSummary( + shape=(5, 9), dtype=np.dtype(np.float32), mean=1.6432636), + h=ResultSummary( + shape=(5, 9), dtype=np.dtype(np.float32), mean=0.75866824)), + attention=ResultSummary( + shape=(5, 6), dtype=np.dtype(np.float32), mean=6.7445569), + time=3, + alignments=ResultSummary( + shape=(5, 8), dtype=np.dtype(np.float32), mean=0.125), + attention_state=ResultSummary( + shape=(5, 8), dtype=np.dtype(np.float32), mean=0.125), + alignment_history=()) + expected_final_alignment_history = ResultSummary( + shape=(3, 5, 8), dtype=np.dtype(np.float32), mean=0.125) + + self._testWithAttention( + create_attention_mechanism, + expected_final_output, + expected_final_state, + alignment_history=True, + create_query_layer=True, + expected_final_alignment_history=expected_final_alignment_history, + create_attention_kwargs=create_attention_kwargs) + + def testBahdanauNormalized(self): + create_attention_mechanism = wrapper.BahdanauAttentionV2 + create_attention_kwargs = {"kernel_initializer": "ones", "normalize": True} + + expected_final_output = basic_decoder.BasicDecoderOutput( + rnn_output=ResultSummary( + shape=(5, 3, 6), dtype=np.dtype("float32"), mean=3.9548259), + sample_id=ResultSummary( + shape=(5, 3), dtype=np.dtype("int32"), mean=0.0)) + expected_final_state = wrapper.AttentionWrapperState( + cell_state=rnn_cell.LSTMStateTuple( + c=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=1.4652209), + h=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=0.70997983)), + attention=ResultSummary( + shape=(5, 6), dtype=np.dtype("float32"), mean=6.3075728), + time=3, + alignments=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.125), + attention_state=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.125), + alignment_history=()) + + self._testWithAttention( + create_attention_mechanism, + expected_final_output, + expected_final_state, + create_query_layer=True, + create_attention_kwargs=create_attention_kwargs) + + def testLuongNotNormalized(self): + create_attention_mechanism = wrapper.LuongAttentionV2 + + expected_final_output = basic_decoder.BasicDecoderOutput( + rnn_output=ResultSummary( + shape=(5, 3, 6), dtype=np.dtype("float32"), mean=2.6605489), + sample_id=ResultSummary( + shape=(5, 3), dtype=np.dtype("int32"), mean=0.0)) + expected_final_state = wrapper.AttentionWrapperState( + cell_state=rnn_cell.LSTMStateTuple( + c=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=0.88403547), + h=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=0.37819088)), + attention=ResultSummary( + shape=(5, 6), dtype=np.dtype("float32"), mean=4.084631), + time=3, + alignments=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.125), + attention_state=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.125), + alignment_history=()) + + self._testWithAttention( + create_attention_mechanism, + expected_final_output, + expected_final_state, + attention_mechanism_depth=9) + + def testLuongScaled(self): + create_attention_mechanism = wrapper.LuongAttentionV2 + create_attention_kwargs = {"scale": True} + + expected_final_output = basic_decoder.BasicDecoderOutput( + rnn_output=ResultSummary( + shape=(5, 3, 6), dtype=np.dtype("float32"), mean=2.6605489), + sample_id=ResultSummary( + shape=(5, 3), dtype=np.dtype("int32"), mean=0.0)) + expected_final_state = wrapper.AttentionWrapperState( + cell_state=rnn_cell.LSTMStateTuple( + c=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=0.88403547), + h=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=0.37819088)), + attention=ResultSummary( + shape=(5, 6), dtype=np.dtype("float32"), mean=4.0846314), + time=3, + alignments=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.125), + attention_state=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.125), + alignment_history=()) + + self._testWithAttention( + create_attention_mechanism, + expected_final_output, + expected_final_state, + attention_mechanism_depth=9, + create_attention_kwargs=create_attention_kwargs) + + def testNotUseAttentionLayer(self): + create_attention_mechanism = wrapper.BahdanauAttentionV2 + create_attention_kwargs = {"kernel_initializer": "ones"} + + expected_final_output = basic_decoder.BasicDecoderOutput( + rnn_output=ResultSummary( + shape=(5, 3, 10), dtype=np.dtype("float32"), mean=0.072406612), + sample_id=ResultSummary( + shape=(5, 3), dtype=np.dtype("int32"), mean=3.86666666)) + expected_final_state = wrapper.AttentionWrapperState( + cell_state=rnn_cell.LSTMStateTuple( + c=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=1.032002), + h=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=0.61177742)), + attention=ResultSummary( + shape=(5, 10), dtype=np.dtype("float32"), mean=0.011346335), + time=3, + alignments=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.125), + attention_state=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.125), + alignment_history=()) + + self._testWithAttention( + create_attention_mechanism, + expected_final_output, + expected_final_state, + attention_layer_size=None, + create_query_layer=True, + create_attention_kwargs=create_attention_kwargs) + + def testBahdanauMonotonicNotNormalized(self): + create_attention_mechanism = wrapper.BahdanauMonotonicAttentionV2 + create_attention_kwargs = {"kernel_initializer": "ones"} + + expected_final_output = basic_decoder.BasicDecoderOutput( + rnn_output=ResultSummary( + shape=(5, 3, 6), dtype=np.dtype("float32"), mean=5.9850435), + sample_id=ResultSummary( + shape=(5, 3), dtype=np.dtype("int32"), mean=0.0)) + expected_final_state = wrapper.AttentionWrapperState( + cell_state=rnn_cell.LSTMStateTuple( + c=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=1.6752492), + h=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=0.76052248)), + attention=ResultSummary( + shape=(5, 6), dtype=np.dtype("float32"), mean=8.361186), + time=3, + alignments=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.10989678), + attention_state=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.10989678), + alignment_history=()) + expected_final_alignment_history = ResultSummary( + shape=(3, 5, 8), dtype=np.dtype("float32"), mean=0.117412611) + + self._testWithAttention( + create_attention_mechanism, + expected_final_output, + expected_final_state, + alignment_history=True, + expected_final_alignment_history=expected_final_alignment_history, + create_query_layer=True, + create_attention_kwargs=create_attention_kwargs) + + def testBahdanauMonotonicNormalized(self): + create_attention_mechanism = wrapper.BahdanauMonotonicAttentionV2 + create_attention_kwargs = {"kernel_initializer": "ones", + "normalize": True} + expected_final_output = basic_decoder.BasicDecoderOutput( + rnn_output=ResultSummary( + shape=(5, 3, 6), dtype=np.dtype("float32"), mean=4.5706983), + sample_id=ResultSummary( + shape=(5, 3), dtype=np.dtype("int32"), mean=0.0)) + expected_final_state = wrapper.AttentionWrapperState( + cell_state=rnn_cell.LSTMStateTuple( + c=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=1.6005473), + h=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=0.77863038)), + attention=ResultSummary( + shape=(5, 6), dtype=np.dtype("float32"), mean=7.3326721), + time=3, + alignments=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.12258384), + attention_state=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.12258384), + alignment_history=()) + expected_final_alignment_history = ResultSummary( + shape=(3, 5, 8), dtype=np.dtype("float32"), mean=0.12258384) + + self._testWithAttention( + create_attention_mechanism, + expected_final_output, + expected_final_state, + alignment_history=True, + expected_final_alignment_history=expected_final_alignment_history, + create_query_layer=True, + create_attention_kwargs=create_attention_kwargs) + + def testLuongMonotonicNotNormalized(self): + create_attention_mechanism = wrapper.LuongMonotonicAttentionV2 + + expected_final_output = basic_decoder.BasicDecoderOutput( + rnn_output=ResultSummary( + shape=(5, 3, 6), dtype=np.dtype("float32"), mean=3.159497), + sample_id=ResultSummary( + shape=(5, 3), dtype=np.dtype("int32"), mean=0.0)) + expected_final_state = wrapper.AttentionWrapperState( + cell_state=rnn_cell.LSTMStateTuple( + c=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=1.072384), + h=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=0.50331038)), + attention=ResultSummary( + shape=(5, 6), dtype=np.dtype("float32"), mean=5.3079605), + time=3, + alignments=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.11467695), + attention_state=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.11467695), + alignment_history=()) + expected_final_alignment_history = ResultSummary( + shape=(3, 5, 8), dtype=np.dtype("float32"), mean=0.11899644) + + self._testWithAttention( + create_attention_mechanism, + expected_final_output, + expected_final_state, + attention_mechanism_depth=9, + alignment_history=True, + expected_final_alignment_history=expected_final_alignment_history) + + def testLuongMonotonicScaled(self): + create_attention_mechanism = wrapper.LuongMonotonicAttentionV2 + create_attention_kwargs = {"scale": True} + + expected_final_output = basic_decoder.BasicDecoderOutput( + rnn_output=ResultSummary( + shape=(5, 3, 6), dtype=np.dtype("float32"), mean=3.159497), + sample_id=ResultSummary( + shape=(5, 3), dtype=np.dtype("int32"), mean=0.0)) + expected_final_state = wrapper.AttentionWrapperState( + cell_state=rnn_cell.LSTMStateTuple( + c=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=1.072384), + h=ResultSummary( + shape=(5, 9), dtype=np.dtype("float32"), mean=0.50331038)), + attention=ResultSummary( + shape=(5, 6), dtype=np.dtype("float32"), mean=5.3079605), + time=3, + alignments=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.11467695), + attention_state=ResultSummary( + shape=(5, 8), dtype=np.dtype("float32"), mean=0.11467695), + alignment_history=()) + expected_final_alignment_history = ResultSummary( + shape=(3, 5, 8), dtype=np.dtype("float32"), mean=0.11899644) + + self._testWithAttention( + create_attention_mechanism, + expected_final_output, + expected_final_state, + attention_mechanism_depth=9, + alignment_history=True, + expected_final_alignment_history=expected_final_alignment_history, + create_attention_kwargs=create_attention_kwargs) + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/seq2seq/python/kernel_tests/basic_decoder_test.py b/tensorflow/contrib/seq2seq/python/kernel_tests/basic_decoder_test.py index b7f9f3fb090356a1c8d2bfb5044712ff93e267ce..abcf71c61b6e6df9462bf06323b8b11d5cc0d9a8 100644 --- a/tensorflow/contrib/seq2seq/python/kernel_tests/basic_decoder_test.py +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/basic_decoder_test.py @@ -34,8 +34,6 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops import rnn_cell from tensorflow.python.ops import variables from tensorflow.python.ops import variable_scope -from tensorflow.python.ops.distributions import bernoulli -from tensorflow.python.ops.distributions import categorical from tensorflow.python.platform import test # pylint: enable=g-import-not-at-top @@ -517,7 +515,7 @@ class BasicDecoderTest(test.TestCase): vocabulary_size) # The sample function samples categorically from the logits. - sample_fn = lambda x: categorical.Categorical(logits=x).sample() + sample_fn = lambda x: helper_py.categorical_sample(logits=x) # The next inputs are a one-hot encoding of the sampled labels. next_inputs_fn = ( lambda x: array_ops.one_hot(x, vocabulary_size, dtype=dtypes.float32)) @@ -599,7 +597,7 @@ class BasicDecoderTest(test.TestCase): # The sample function samples independent bernoullis from the logits. sample_fn = ( - lambda x: bernoulli.Bernoulli(logits=x, dtype=dtypes.bool).sample()) + lambda x: helper_py.bernoulli_sample(logits=x, dtype=dtypes.bool)) # The next inputs are a one-hot encoding of the sampled labels. next_inputs_fn = math_ops.to_float end_fn = lambda sample_ids: sample_ids[:, end_token] diff --git a/tensorflow/contrib/seq2seq/python/kernel_tests/basic_decoder_v2_test.py b/tensorflow/contrib/seq2seq/python/kernel_tests/basic_decoder_v2_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2341ebb77ab6ecad1e979bc8bed0080128a804da --- /dev/null +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/basic_decoder_v2_test.py @@ -0,0 +1,670 @@ +# 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. +# ============================================================================== +"""Tests for contrib.seq2seq.python.seq2seq.basic_decoder_v2.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from absl.testing import parameterized +import numpy as np + +from tensorflow.contrib.seq2seq.python.ops import basic_decoder +from tensorflow.contrib.seq2seq.python.ops import sampler as sampler_py + +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import tensor_shape +from tensorflow.python.keras import keras_parameterized +from tensorflow.python.layers import core as layers_core +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import rnn_cell +from tensorflow.python.ops import variables +from tensorflow.python.platform import test + + +@keras_parameterized.run_all_keras_modes +class BasicDecoderTest(keras_parameterized.TestCase): + """Unit test for basic_decoder.BasicDecoderV2.""" + + @parameterized.named_parameters( + ("use_output_layer", True), + ("without_output_layer", False)) + def testStepWithTrainingHelperOutputLayer(self, use_output_layer): + sequence_length = [3, 4, 3, 1, 0] + batch_size = 5 + max_time = 8 + input_depth = 7 + cell_depth = 10 + output_layer_depth = 3 + + with self.cached_session(use_gpu=True): + inputs = np.random.randn(batch_size, max_time, + input_depth).astype(np.float32) + input_t = constant_op.constant(inputs) + cell = rnn_cell.LSTMCell(cell_depth) + sampler = sampler_py.TrainingSampler(time_major=False) + if use_output_layer: + output_layer = layers_core.Dense(output_layer_depth, use_bias=False) + expected_output_depth = output_layer_depth + else: + output_layer = None + expected_output_depth = cell_depth + initial_state = cell.zero_state(dtype=dtypes.float32, + batch_size=batch_size) + my_decoder = basic_decoder.BasicDecoderV2( + cell=cell, + sampler=sampler, + output_layer=output_layer) + + (first_finished, + first_inputs, + first_state) = my_decoder.initialize(input_t, + initial_state=initial_state, + sequence_length=sequence_length) + output_size = my_decoder.output_size + output_dtype = my_decoder.output_dtype + self.assertEqual( + basic_decoder.BasicDecoderOutput(expected_output_depth, + tensor_shape.TensorShape([])), + output_size) + self.assertEqual( + basic_decoder.BasicDecoderOutput(dtypes.float32, dtypes.int32), + output_dtype) + + (step_outputs, step_state, step_next_inputs, + step_finished) = my_decoder.step( + constant_op.constant(0), first_inputs, first_state) + batch_size_t = my_decoder.batch_size + + self.assertTrue(isinstance(first_state, rnn_cell.LSTMStateTuple)) + self.assertTrue(isinstance(step_state, rnn_cell.LSTMStateTuple)) + self.assertTrue( + isinstance(step_outputs, basic_decoder.BasicDecoderOutput)) + self.assertEqual((batch_size, expected_output_depth), + step_outputs[0].get_shape()) + self.assertEqual((batch_size,), step_outputs[1].get_shape()) + self.assertEqual((batch_size, cell_depth), first_state[0].get_shape()) + self.assertEqual((batch_size, cell_depth), first_state[1].get_shape()) + self.assertEqual((batch_size, cell_depth), step_state[0].get_shape()) + self.assertEqual((batch_size, cell_depth), step_state[1].get_shape()) + + if use_output_layer: + # The output layer was accessed + self.assertEqual(len(output_layer.variables), 1) + + self.evaluate(variables.global_variables_initializer()) + eval_result = self.evaluate({ + "batch_size": batch_size_t, + "first_finished": first_finished, + "first_inputs": first_inputs, + "first_state": first_state, + "step_outputs": step_outputs, + "step_state": step_state, + "step_next_inputs": step_next_inputs, + "step_finished": step_finished + }) + + self.assertAllEqual([False, False, False, False, True], + eval_result["first_finished"]) + self.assertAllEqual([False, False, False, True, True], + eval_result["step_finished"]) + self.assertEqual(output_dtype.sample_id, + eval_result["step_outputs"].sample_id.dtype) + self.assertAllEqual( + np.argmax(eval_result["step_outputs"].rnn_output, -1), + eval_result["step_outputs"].sample_id) + + def DISABLED_testStepWithGreedyEmbeddingHelper(self): + batch_size = 5 + vocabulary_size = 7 + cell_depth = vocabulary_size # cell's logits must match vocabulary size + input_depth = 10 + start_tokens = np.random.randint(0, vocabulary_size, size=batch_size) + end_token = 1 + + with self.cached_session(use_gpu=True): + embeddings = np.random.randn(vocabulary_size, + input_depth).astype(np.float32) + embeddings_t = constant_op.constant(embeddings) + cell = rnn_cell.LSTMCell(vocabulary_size) + sampler = sampler_py.GreedyEmbeddingSampler() + initial_state = cell.zero_state( + dtype=dtypes.float32, batch_size=batch_size) + my_decoder = basic_decoder.BasicDecoderV2( + cell=cell, + sampler=sampler) + (first_finished, first_inputs, first_state) = my_decoder.initialize( + embeddings_t, + start_tokens=start_tokens, + end_token=end_token, + initial_state=initial_state) + output_size = my_decoder.output_size + output_dtype = my_decoder.output_dtype + self.assertEqual( + basic_decoder.BasicDecoderOutput(cell_depth, + tensor_shape.TensorShape([])), + output_size) + self.assertEqual( + basic_decoder.BasicDecoderOutput(dtypes.float32, dtypes.int32), + output_dtype) + + (step_outputs, step_state, step_next_inputs, + step_finished) = my_decoder.step( + constant_op.constant(0), first_inputs, first_state) + batch_size_t = my_decoder.batch_size + + self.assertTrue(isinstance(first_state, rnn_cell.LSTMStateTuple)) + self.assertTrue(isinstance(step_state, rnn_cell.LSTMStateTuple)) + self.assertTrue( + isinstance(step_outputs, basic_decoder.BasicDecoderOutput)) + self.assertEqual((batch_size, cell_depth), step_outputs[0].get_shape()) + self.assertEqual((batch_size,), step_outputs[1].get_shape()) + self.assertEqual((batch_size, cell_depth), first_state[0].get_shape()) + self.assertEqual((batch_size, cell_depth), first_state[1].get_shape()) + self.assertEqual((batch_size, cell_depth), step_state[0].get_shape()) + self.assertEqual((batch_size, cell_depth), step_state[1].get_shape()) + + self.evaluate(variables.global_variables_initializer()) + eval_result = self.evaluate({ + "batch_size": batch_size_t, + "first_finished": first_finished, + "first_inputs": first_inputs, + "first_state": first_state, + "step_outputs": step_outputs, + "step_state": step_state, + "step_next_inputs": step_next_inputs, + "step_finished": step_finished + }) + + expected_sample_ids = np.argmax( + eval_result["step_outputs"].rnn_output, -1) + expected_step_finished = (expected_sample_ids == end_token) + expected_step_next_inputs = embeddings[expected_sample_ids] + self.assertAllEqual([False, False, False, False, False], + eval_result["first_finished"]) + self.assertAllEqual(expected_step_finished, eval_result["step_finished"]) + self.assertEqual(output_dtype.sample_id, + eval_result["step_outputs"].sample_id.dtype) + self.assertAllEqual(expected_sample_ids, + eval_result["step_outputs"].sample_id) + self.assertAllEqual(expected_step_next_inputs, + eval_result["step_next_inputs"]) + + def testStepWithSampleEmbeddingHelper(self): + batch_size = 5 + vocabulary_size = 7 + cell_depth = vocabulary_size # cell's logits must match vocabulary size + input_depth = 10 + np.random.seed(0) + start_tokens = np.random.randint(0, vocabulary_size, size=batch_size) + end_token = 1 + + with self.cached_session(use_gpu=True): + embeddings = np.random.randn(vocabulary_size, + input_depth).astype(np.float32) + embeddings_t = constant_op.constant(embeddings) + cell = rnn_cell.LSTMCell(vocabulary_size) + sampler = sampler_py.SampleEmbeddingSampler(seed=0) + initial_state = cell.zero_state( + dtype=dtypes.float32, batch_size=batch_size) + my_decoder = basic_decoder.BasicDecoderV2(cell=cell, sampler=sampler) + (first_finished, + first_inputs, + first_state) = my_decoder.initialize(embeddings_t, + start_tokens=start_tokens, + end_token=end_token, + initial_state=initial_state) + output_size = my_decoder.output_size + output_dtype = my_decoder.output_dtype + self.assertEqual( + basic_decoder.BasicDecoderOutput(cell_depth, + tensor_shape.TensorShape([])), + output_size) + self.assertEqual( + basic_decoder.BasicDecoderOutput(dtypes.float32, dtypes.int32), + output_dtype) + + (step_outputs, step_state, step_next_inputs, + step_finished) = my_decoder.step( + constant_op.constant(0), first_inputs, first_state) + batch_size_t = my_decoder.batch_size + + self.assertTrue(isinstance(first_state, rnn_cell.LSTMStateTuple)) + self.assertTrue(isinstance(step_state, rnn_cell.LSTMStateTuple)) + self.assertTrue( + isinstance(step_outputs, basic_decoder.BasicDecoderOutput)) + self.assertEqual((batch_size, cell_depth), step_outputs[0].get_shape()) + self.assertEqual((batch_size,), step_outputs[1].get_shape()) + self.assertEqual((batch_size, cell_depth), first_state[0].get_shape()) + self.assertEqual((batch_size, cell_depth), first_state[1].get_shape()) + self.assertEqual((batch_size, cell_depth), step_state[0].get_shape()) + self.assertEqual((batch_size, cell_depth), step_state[1].get_shape()) + + self.evaluate(variables.global_variables_initializer()) + eval_result = self.evaluate({ + "batch_size": batch_size_t, + "first_finished": first_finished, + "first_inputs": first_inputs, + "first_state": first_state, + "step_outputs": step_outputs, + "step_state": step_state, + "step_next_inputs": step_next_inputs, + "step_finished": step_finished + }) + + sample_ids = eval_result["step_outputs"].sample_id + self.assertEqual(output_dtype.sample_id, sample_ids.dtype) + expected_step_finished = (sample_ids == end_token) + expected_step_next_inputs = embeddings[sample_ids] + self.assertAllEqual(expected_step_finished, + eval_result["step_finished"]) + self.assertAllEqual(expected_step_next_inputs, + eval_result["step_next_inputs"]) + + def testStepWithScheduledEmbeddingTrainingHelper(self): + sequence_length = [3, 4, 3, 1, 0] + batch_size = 5 + max_time = 8 + input_depth = 7 + vocabulary_size = 10 + + with self.cached_session(use_gpu=True): + inputs = np.random.randn( + batch_size, max_time, input_depth).astype(np.float32) + input_t = constant_op.constant(inputs) + embeddings = np.random.randn( + vocabulary_size, input_depth).astype(np.float32) + half = constant_op.constant(0.5) + cell = rnn_cell.LSTMCell(vocabulary_size) + sampler = sampler_py.ScheduledEmbeddingTrainingSampler( + sampling_probability=half, + time_major=False) + initial_state = cell.zero_state( + dtype=dtypes.float32, batch_size=batch_size) + my_decoder = basic_decoder.BasicDecoderV2( + cell=cell, + sampler=sampler) + (first_finished, first_inputs, first_state) = my_decoder.initialize( + input_t, sequence_length=sequence_length, embedding=embeddings, + initial_state=initial_state) + output_size = my_decoder.output_size + output_dtype = my_decoder.output_dtype + self.assertEqual( + basic_decoder.BasicDecoderOutput(vocabulary_size, + tensor_shape.TensorShape([])), + output_size) + self.assertEqual( + basic_decoder.BasicDecoderOutput(dtypes.float32, dtypes.int32), + output_dtype) + + (step_outputs, step_state, step_next_inputs, + step_finished) = my_decoder.step( + constant_op.constant(0), first_inputs, first_state) + batch_size_t = my_decoder.batch_size + + self.assertTrue(isinstance(first_state, rnn_cell.LSTMStateTuple)) + self.assertTrue(isinstance(step_state, rnn_cell.LSTMStateTuple)) + self.assertTrue( + isinstance(step_outputs, basic_decoder.BasicDecoderOutput)) + self.assertEqual((batch_size, vocabulary_size), + step_outputs[0].get_shape()) + self.assertEqual((batch_size,), step_outputs[1].get_shape()) + self.assertEqual((batch_size, vocabulary_size), + first_state[0].get_shape()) + self.assertEqual((batch_size, vocabulary_size), + first_state[1].get_shape()) + self.assertEqual((batch_size, vocabulary_size), + step_state[0].get_shape()) + self.assertEqual((batch_size, vocabulary_size), + step_state[1].get_shape()) + self.assertEqual((batch_size, input_depth), + step_next_inputs.get_shape()) + + self.evaluate(variables.global_variables_initializer()) + eval_result = self.evaluate({ + "batch_size": batch_size_t, + "first_finished": first_finished, + "first_inputs": first_inputs, + "first_state": first_state, + "step_outputs": step_outputs, + "step_state": step_state, + "step_next_inputs": step_next_inputs, + "step_finished": step_finished + }) + + self.assertAllEqual([False, False, False, False, True], + eval_result["first_finished"]) + self.assertAllEqual([False, False, False, True, True], + eval_result["step_finished"]) + sample_ids = eval_result["step_outputs"].sample_id + self.assertEqual(output_dtype.sample_id, sample_ids.dtype) + batch_where_not_sampling = np.where(sample_ids == -1) + batch_where_sampling = np.where(sample_ids > -1) + self.assertAllClose( + eval_result["step_next_inputs"][batch_where_sampling], + embeddings[sample_ids[batch_where_sampling]]) + self.assertAllClose( + eval_result["step_next_inputs"][batch_where_not_sampling], + np.squeeze(inputs[batch_where_not_sampling, 1], axis=0)) + + def _testStepWithScheduledOutputTrainingHelper( + self, sampling_probability, use_next_inputs_fn, use_auxiliary_inputs): + sequence_length = [3, 4, 3, 1, 0] + batch_size = 5 + max_time = 8 + input_depth = 7 + cell_depth = input_depth + if use_auxiliary_inputs: + auxiliary_input_depth = 4 + auxiliary_inputs = np.random.randn( + batch_size, max_time, auxiliary_input_depth).astype(np.float32) + else: + auxiliary_inputs = None + + with self.cached_session(use_gpu=True): + inputs = np.random.randn(batch_size, max_time, + input_depth).astype(np.float32) + input_t = constant_op.constant(inputs) + cell = rnn_cell.LSTMCell(cell_depth) + sampling_probability = constant_op.constant(sampling_probability) + + if use_next_inputs_fn: + def next_inputs_fn(outputs): + # Use deterministic function for test. + samples = math_ops.argmax(outputs, axis=1) + return array_ops.one_hot(samples, cell_depth, dtype=dtypes.float32) + else: + next_inputs_fn = None + + sampler = sampler_py.ScheduledOutputTrainingSampler( + sampling_probability=sampling_probability, + time_major=False, + next_inputs_fn=next_inputs_fn) + initial_state = cell.zero_state( + dtype=dtypes.float32, batch_size=batch_size) + + my_decoder = basic_decoder.BasicDecoderV2( + cell=cell, + sampler=sampler) + + (first_finished, + first_inputs, + first_state) = my_decoder.initialize(input_t, + sequence_length=sequence_length, + initial_state=initial_state, + auxiliary_inputs=auxiliary_inputs) + output_size = my_decoder.output_size + output_dtype = my_decoder.output_dtype + self.assertEqual( + basic_decoder.BasicDecoderOutput(cell_depth, + tensor_shape.TensorShape([])), + output_size) + self.assertEqual( + basic_decoder.BasicDecoderOutput(dtypes.float32, dtypes.int32), + output_dtype) + + (step_outputs, step_state, step_next_inputs, + step_finished) = my_decoder.step( + constant_op.constant(0), first_inputs, first_state) + + if use_next_inputs_fn: + output_after_next_inputs_fn = next_inputs_fn(step_outputs.rnn_output) + + batch_size_t = my_decoder.batch_size + + self.assertTrue(isinstance(first_state, rnn_cell.LSTMStateTuple)) + self.assertTrue(isinstance(step_state, rnn_cell.LSTMStateTuple)) + self.assertTrue( + isinstance(step_outputs, basic_decoder.BasicDecoderOutput)) + self.assertEqual((batch_size, cell_depth), step_outputs[0].get_shape()) + self.assertEqual((batch_size,), step_outputs[1].get_shape()) + self.assertEqual((batch_size, cell_depth), first_state[0].get_shape()) + self.assertEqual((batch_size, cell_depth), first_state[1].get_shape()) + self.assertEqual((batch_size, cell_depth), step_state[0].get_shape()) + self.assertEqual((batch_size, cell_depth), step_state[1].get_shape()) + + self.evaluate(variables.global_variables_initializer()) + + fetches = { + "batch_size": batch_size_t, + "first_finished": first_finished, + "first_inputs": first_inputs, + "first_state": first_state, + "step_outputs": step_outputs, + "step_state": step_state, + "step_next_inputs": step_next_inputs, + "step_finished": step_finished + } + if use_next_inputs_fn: + fetches["output_after_next_inputs_fn"] = output_after_next_inputs_fn + + eval_result = self.evaluate(fetches) + + self.assertAllEqual([False, False, False, False, True], + eval_result["first_finished"]) + self.assertAllEqual([False, False, False, True, True], + eval_result["step_finished"]) + + sample_ids = eval_result["step_outputs"].sample_id + self.assertEqual(output_dtype.sample_id, sample_ids.dtype) + batch_where_not_sampling = np.where(np.logical_not(sample_ids)) + batch_where_sampling = np.where(sample_ids) + + auxiliary_inputs_to_concat = ( + auxiliary_inputs[:, 1] if use_auxiliary_inputs else + np.array([]).reshape(batch_size, 0).astype(np.float32)) + + expected_next_sampling_inputs = np.concatenate( + (eval_result["output_after_next_inputs_fn"][batch_where_sampling] + if use_next_inputs_fn else + eval_result["step_outputs"].rnn_output[batch_where_sampling], + auxiliary_inputs_to_concat[batch_where_sampling]), + axis=-1) + self.assertAllClose( + eval_result["step_next_inputs"][batch_where_sampling], + expected_next_sampling_inputs) + + self.assertAllClose( + eval_result["step_next_inputs"][batch_where_not_sampling], + np.concatenate( + (np.squeeze(inputs[batch_where_not_sampling, 1], axis=0), + auxiliary_inputs_to_concat[batch_where_not_sampling]), + axis=-1)) + + def testStepWithScheduledOutputTrainingHelperWithoutNextInputsFnOrAuxInputs( + self): + self._testStepWithScheduledOutputTrainingHelper( + sampling_probability=0.5, use_next_inputs_fn=False, + use_auxiliary_inputs=False) + + def testStepWithScheduledOutputTrainingHelperWithNextInputsFn(self): + self._testStepWithScheduledOutputTrainingHelper( + sampling_probability=0.5, use_next_inputs_fn=True, + use_auxiliary_inputs=False) + + def testStepWithScheduledOutputTrainingHelperWithAuxiliaryInputs(self): + self._testStepWithScheduledOutputTrainingHelper( + sampling_probability=0.5, use_next_inputs_fn=False, + use_auxiliary_inputs=True) + + def testStepWithScheduledOutputTrainingHelperWithNextInputsFnAndAuxInputs( + self): + self._testStepWithScheduledOutputTrainingHelper( + sampling_probability=0.5, use_next_inputs_fn=True, + use_auxiliary_inputs=True) + + def testStepWithScheduledOutputTrainingHelperWithNoSampling(self): + self._testStepWithScheduledOutputTrainingHelper( + sampling_probability=0.0, use_next_inputs_fn=True, + use_auxiliary_inputs=True) + + def testStepWithInferenceHelperCategorical(self): + batch_size = 5 + vocabulary_size = 7 + cell_depth = vocabulary_size + start_token = 0 + end_token = 6 + + start_inputs = array_ops.one_hot( + np.ones(batch_size, dtype=np.int32) * start_token, + vocabulary_size) + + # The sample function samples categorically from the logits. + sample_fn = lambda x: sampler_py.categorical_sample(logits=x) + # The next inputs are a one-hot encoding of the sampled labels. + next_inputs_fn = ( + lambda x: array_ops.one_hot(x, vocabulary_size, dtype=dtypes.float32)) + end_fn = lambda sample_ids: math_ops.equal(sample_ids, end_token) + + with self.cached_session(use_gpu=True): + cell = rnn_cell.LSTMCell(vocabulary_size) + sampler = sampler_py.InferenceSampler( + sample_fn, sample_shape=(), sample_dtype=dtypes.int32, end_fn=end_fn, + next_inputs_fn=next_inputs_fn) + initial_state = cell.zero_state( + dtype=dtypes.float32, batch_size=batch_size) + my_decoder = basic_decoder.BasicDecoderV2( + cell=cell, + sampler=sampler) + (first_finished, first_inputs, first_state) = my_decoder.initialize( + start_inputs, initial_state=initial_state) + + output_size = my_decoder.output_size + output_dtype = my_decoder.output_dtype + self.assertEqual( + basic_decoder.BasicDecoderOutput(cell_depth, + tensor_shape.TensorShape([])), + output_size) + self.assertEqual( + basic_decoder.BasicDecoderOutput(dtypes.float32, dtypes.int32), + output_dtype) + + (step_outputs, step_state, step_next_inputs, + step_finished) = my_decoder.step( + constant_op.constant(0), first_inputs, first_state) + batch_size_t = my_decoder.batch_size + + self.assertTrue(isinstance(first_state, rnn_cell.LSTMStateTuple)) + self.assertTrue(isinstance(step_state, rnn_cell.LSTMStateTuple)) + self.assertTrue( + isinstance(step_outputs, basic_decoder.BasicDecoderOutput)) + self.assertEqual((batch_size, cell_depth), step_outputs[0].get_shape()) + self.assertEqual((batch_size,), step_outputs[1].get_shape()) + self.assertEqual((batch_size, cell_depth), first_state[0].get_shape()) + self.assertEqual((batch_size, cell_depth), first_state[1].get_shape()) + self.assertEqual((batch_size, cell_depth), step_state[0].get_shape()) + self.assertEqual((batch_size, cell_depth), step_state[1].get_shape()) + + self.evaluate(variables.global_variables_initializer()) + eval_result = self.evaluate({ + "batch_size": batch_size_t, + "first_finished": first_finished, + "first_inputs": first_inputs, + "first_state": first_state, + "step_outputs": step_outputs, + "step_state": step_state, + "step_next_inputs": step_next_inputs, + "step_finished": step_finished + }) + + sample_ids = eval_result["step_outputs"].sample_id + self.assertEqual(output_dtype.sample_id, sample_ids.dtype) + expected_step_finished = (sample_ids == end_token) + expected_step_next_inputs = np.zeros((batch_size, vocabulary_size)) + expected_step_next_inputs[np.arange(batch_size), sample_ids] = 1.0 + self.assertAllEqual(expected_step_finished, + eval_result["step_finished"]) + self.assertAllEqual(expected_step_next_inputs, + eval_result["step_next_inputs"]) + + def testStepWithInferenceHelperMultilabel(self): + batch_size = 5 + vocabulary_size = 7 + cell_depth = vocabulary_size + start_token = 0 + end_token = 6 + + start_inputs = array_ops.one_hot( + np.ones(batch_size, dtype=np.int32) * start_token, + vocabulary_size) + + # The sample function samples independent bernoullis from the logits. + sample_fn = ( + lambda x: sampler_py.bernoulli_sample(logits=x, dtype=dtypes.bool)) + # The next inputs are a one-hot encoding of the sampled labels. + next_inputs_fn = math_ops.to_float + end_fn = lambda sample_ids: sample_ids[:, end_token] + + with self.cached_session(use_gpu=True): + cell = rnn_cell.LSTMCell(vocabulary_size) + sampler = sampler_py.InferenceSampler( + sample_fn, sample_shape=[cell_depth], sample_dtype=dtypes.bool, + end_fn=end_fn, next_inputs_fn=next_inputs_fn) + initial_state = cell.zero_state( + dtype=dtypes.float32, batch_size=batch_size) + my_decoder = basic_decoder.BasicDecoderV2( + cell=cell, + sampler=sampler) + (first_finished, first_inputs, first_state) = my_decoder.initialize( + start_inputs, initial_state=initial_state) + output_size = my_decoder.output_size + output_dtype = my_decoder.output_dtype + self.assertEqual( + basic_decoder.BasicDecoderOutput(cell_depth, cell_depth), + output_size) + self.assertEqual( + basic_decoder.BasicDecoderOutput(dtypes.float32, dtypes.bool), + output_dtype) + + (step_outputs, step_state, step_next_inputs, + step_finished) = my_decoder.step( + constant_op.constant(0), first_inputs, first_state) + batch_size_t = my_decoder.batch_size + + self.assertTrue(isinstance(first_state, rnn_cell.LSTMStateTuple)) + self.assertTrue(isinstance(step_state, rnn_cell.LSTMStateTuple)) + self.assertTrue( + isinstance(step_outputs, basic_decoder.BasicDecoderOutput)) + self.assertEqual((batch_size, cell_depth), step_outputs[0].get_shape()) + self.assertEqual((batch_size, cell_depth), step_outputs[1].get_shape()) + self.assertEqual((batch_size, cell_depth), first_state[0].get_shape()) + self.assertEqual((batch_size, cell_depth), first_state[1].get_shape()) + self.assertEqual((batch_size, cell_depth), step_state[0].get_shape()) + self.assertEqual((batch_size, cell_depth), step_state[1].get_shape()) + + self.evaluate(variables.global_variables_initializer()) + eval_result = self.evaluate({ + "batch_size": batch_size_t, + "first_finished": first_finished, + "first_inputs": first_inputs, + "first_state": first_state, + "step_outputs": step_outputs, + "step_state": step_state, + "step_next_inputs": step_next_inputs, + "step_finished": step_finished + }) + + sample_ids = eval_result["step_outputs"].sample_id + self.assertEqual(output_dtype.sample_id, sample_ids.dtype) + expected_step_finished = sample_ids[:, end_token] + expected_step_next_inputs = sample_ids.astype(np.float32) + self.assertAllEqual(expected_step_finished, + eval_result["step_finished"]) + self.assertAllEqual(expected_step_next_inputs, + eval_result["step_next_inputs"]) + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/seq2seq/python/kernel_tests/decoder_v2_test.py b/tensorflow/contrib/seq2seq/python/kernel_tests/decoder_v2_test.py new file mode 100644 index 0000000000000000000000000000000000000000..b5bba2b32e940aa4d5984821ebd3845d7f272549 --- /dev/null +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/decoder_v2_test.py @@ -0,0 +1,169 @@ +# 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. +# ============================================================================== +"""Tests for contrib.seq2seq.python.seq2seq.decoder.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.seq2seq.python.ops import basic_decoder +from tensorflow.contrib.seq2seq.python.ops import sampler as sampler_py +from tensorflow.python.eager import context +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.keras import keras_parameterized +from tensorflow.python.ops import rnn +from tensorflow.python.ops import rnn_cell +from tensorflow.python.ops import variables +from tensorflow.python.platform import test + + +@keras_parameterized.run_all_keras_modes +class DecodeV2RNNTest(keras_parameterized.TestCase, test.TestCase): + """Tests for DecoderV2.""" + + def _testDecodeRNN(self, time_major, maximum_iterations=None): + + sequence_length = [3, 4, 3, 1, 0] + batch_size = 5 + max_time = 8 + input_depth = 7 + cell_depth = 10 + max_out = max(sequence_length) + + with self.cached_session(use_gpu=True): + if time_major: + inputs = np.random.randn(max_time, batch_size, + input_depth).astype(np.float32) + else: + inputs = np.random.randn(batch_size, max_time, + input_depth).astype(np.float32) + input_t = constant_op.constant(inputs) + cell = rnn_cell.LSTMCell(cell_depth) + sampler = sampler_py.TrainingSampler(time_major=time_major) + my_decoder = basic_decoder.BasicDecoderV2( + cell=cell, + sampler=sampler, + output_time_major=time_major, + maximum_iterations=maximum_iterations) + + initial_state = cell.zero_state( + dtype=dtypes.float32, batch_size=batch_size) + (final_outputs, unused_final_state, final_sequence_length) = my_decoder( + input_t, initial_state=initial_state, sequence_length=sequence_length) + + def _t(shape): + if time_major: + return (shape[1], shape[0]) + shape[2:] + return shape + + if not context.executing_eagerly(): + self.assertEqual((batch_size,), + tuple(final_sequence_length.get_shape().as_list())) + self.assertEqual( + _t((batch_size, None, cell_depth)), + tuple(final_outputs.rnn_output.get_shape().as_list())) + self.assertEqual( + _t((batch_size, None)), + tuple(final_outputs.sample_id.get_shape().as_list())) + + self.evaluate(variables.global_variables_initializer()) + final_outputs = self.evaluate(final_outputs) + final_sequence_length = self.evaluate(final_sequence_length) + + # Mostly a smoke test + time_steps = max_out + expected_length = sequence_length + if maximum_iterations is not None: + time_steps = min(max_out, maximum_iterations) + expected_length = [min(x, maximum_iterations) for x in expected_length] + if context.executing_eagerly() and maximum_iterations != 0: + # Only check the shape of output when maximum_iterations > 0, see + # b/123431432 for more details. + self.assertEqual( + _t((batch_size, time_steps, cell_depth)), + final_outputs.rnn_output.shape) + self.assertEqual( + _t((batch_size, time_steps)), final_outputs.sample_id.shape) + self.assertItemsEqual(expected_length, final_sequence_length) + + def testDynamicDecodeRNNBatchMajor(self): + self._testDecodeRNN(time_major=False) + + def testDynamicDecodeRNNTimeMajor(self): + self._testDecodeRNN(time_major=True) + + def testDynamicDecodeRNNZeroMaxIters(self): + self._testDecodeRNN(time_major=True, maximum_iterations=0) + + def testDynamicDecodeRNNOneMaxIter(self): + self._testDecodeRNN(time_major=True, maximum_iterations=1) + + def _testDynamicDecodeRNNWithTrainingHelperMatchesDynamicRNN( + self, use_sequence_length): + sequence_length = [3, 4, 3, 1, 0] + batch_size = 5 + max_time = 8 + input_depth = 7 + cell_depth = 10 + max_out = max(sequence_length) + + with self.cached_session(use_gpu=True): + inputs = np.random.randn(batch_size, max_time, + input_depth).astype(np.float32) + inputs = constant_op.constant(inputs) + + cell = rnn_cell.LSTMCell(cell_depth) + zero_state = cell.zero_state(dtype=dtypes.float32, batch_size=batch_size) + sampler = sampler_py.TrainingSampler() + my_decoder = basic_decoder.BasicDecoderV2( + cell=cell, sampler=sampler, impute_finished=use_sequence_length) + + final_decoder_outputs, final_decoder_state, _ = my_decoder( + inputs, initial_state=zero_state, sequence_length=sequence_length) + + final_rnn_outputs, final_rnn_state = rnn.dynamic_rnn( + cell, + inputs, + sequence_length=sequence_length if use_sequence_length else None, + initial_state=zero_state) + + self.evaluate(variables.global_variables_initializer()) + eval_result = self.evaluate({ + "final_decoder_outputs": final_decoder_outputs, + "final_decoder_state": final_decoder_state, + "final_rnn_outputs": final_rnn_outputs, + "final_rnn_state": final_rnn_state + }) + + # Decoder only runs out to max_out; ensure values are identical + # to dynamic_rnn, which also zeros out outputs and passes along state. + self.assertAllClose(eval_result["final_decoder_outputs"].rnn_output, + eval_result["final_rnn_outputs"][:, 0:max_out, :]) + if use_sequence_length: + self.assertAllClose(eval_result["final_decoder_state"], + eval_result["final_rnn_state"]) + + def testDynamicDecodeRNNWithTrainingHelperMatchesDynamicRNNWithSeqLen(self): + self._testDynamicDecodeRNNWithTrainingHelperMatchesDynamicRNN( + use_sequence_length=True) + + def testDynamicDecodeRNNWithTrainingHelperMatchesDynamicRNNNoSeqLen(self): + self._testDynamicDecodeRNNWithTrainingHelperMatchesDynamicRNN( + use_sequence_length=False) + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/seq2seq/python/kernel_tests/loss_test.py b/tensorflow/contrib/seq2seq/python/kernel_tests/loss_test.py index 5aa32b532ffcf5772f6ace26662f5e5471cf6923..41b2a53ca5b178be9b04446c81d832575e5ed75b 100644 --- a/tensorflow/contrib/seq2seq/python/kernel_tests/loss_test.py +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/loss_test.py @@ -14,80 +14,254 @@ # ============================================================================== """Tests for contrib.seq2seq.python.seq2seq.loss_ops.""" -# pylint: disable=unused-import,g-bad-import-order from __future__ import absolute_import from __future__ import division from __future__ import print_function -# pylint: enable=unused-import import numpy as np from tensorflow.contrib.seq2seq.python.ops import loss from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes +from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops -from tensorflow.python.ops import init_ops -from tensorflow.python.ops import variable_scope from tensorflow.python.platform import test +@test_util.run_all_in_graph_and_eager_modes class LossTest(test.TestCase): + def setUp(self): + self.batch_size = 2 + self.sequence_length = 3 + self.number_of_classes = 5 + logits = [ + constant_op.constant(i + 0.5, shape=[self.batch_size, + self.number_of_classes]) + for i in range(self.sequence_length) + ] + self.logits = array_ops.stack(logits, axis=1) + targets = [ + constant_op.constant(i, dtypes.int32, shape=[self.batch_size]) + for i in range(self.sequence_length) + ] + self.targets = array_ops.stack(targets, axis=1) + weights = [ + constant_op.constant(1.0, shape=[self.batch_size]) + for _ in range(self.sequence_length) + ] + self.weights = array_ops.stack(weights, axis=1) + # expected_loss = sparse_softmax_cross_entropy_with_logits(targets, logits) + # where targets = [0, 1, 2], and logits = [[0.5] * 5, [1.5] * 5, [2.5] * 5] + self.expected_loss = 1.60944 + def testSequenceLoss(self): - with self.session(use_gpu=True) as sess: - with variable_scope.variable_scope( - 'root', initializer=init_ops.constant_initializer(0.5)): - batch_size = 2 - sequence_length = 3 - number_of_classes = 5 - logits = [ - constant_op.constant( - i + 0.5, shape=[batch_size, number_of_classes]) - for i in range(sequence_length) - ] - logits = array_ops.stack(logits, axis=1) - targets = [ - constant_op.constant( - i, dtypes.int32, shape=[batch_size]) - for i in range(sequence_length) - ] - targets = array_ops.stack(targets, axis=1) - weights = [ - constant_op.constant( - 1.0, shape=[batch_size]) for i in range(sequence_length) - ] - weights = array_ops.stack(weights, axis=1) - - average_loss_per_example = loss.sequence_loss( - logits, targets, weights, - average_across_timesteps=True, - average_across_batch=True) - res = sess.run(average_loss_per_example) - self.assertAllClose(1.60944, res) - - average_loss_per_sequence = loss.sequence_loss( - logits, targets, weights, - average_across_timesteps=False, - average_across_batch=True) - res = sess.run(average_loss_per_sequence) - compare_per_sequence = np.ones((sequence_length)) * 1.60944 - self.assertAllClose(compare_per_sequence, res) - - average_loss_per_batch = loss.sequence_loss( - logits, targets, weights, - average_across_timesteps=True, - average_across_batch=False) - res = sess.run(average_loss_per_batch) - compare_per_batch = np.ones((batch_size)) * 1.60944 - self.assertAllClose(compare_per_batch, res) - - total_loss = loss.sequence_loss( - logits, targets, weights, - average_across_timesteps=False, - average_across_batch=False) - res = sess.run(total_loss) - compare_total = np.ones((batch_size, sequence_length)) * 1.60944 - self.assertAllClose(compare_total, res) + with self.test_session(use_gpu=True): + average_loss_per_example = loss.sequence_loss( + self.logits, self.targets, self.weights, + average_across_timesteps=True, + average_across_batch=True) + res = self.evaluate(average_loss_per_example) + self.assertAllClose(self.expected_loss, res) + + average_loss_per_sequence = loss.sequence_loss( + self.logits, self.targets, self.weights, + average_across_timesteps=False, + average_across_batch=True) + res = self.evaluate(average_loss_per_sequence) + compare_per_sequence = np.full((self.sequence_length), self.expected_loss) + self.assertAllClose(compare_per_sequence, res) + + average_loss_per_batch = loss.sequence_loss( + self.logits, self.targets, self.weights, + average_across_timesteps=True, + average_across_batch=False) + res = self.evaluate(average_loss_per_batch) + compare_per_batch = np.full((self.batch_size), self.expected_loss) + self.assertAllClose(compare_per_batch, res) + + total_loss = loss.sequence_loss( + self.logits, self.targets, self.weights, + average_across_timesteps=False, + average_across_batch=False) + res = self.evaluate(total_loss) + compare_total = np.full((self.batch_size, self.sequence_length), + self.expected_loss) + self.assertAllClose(compare_total, res) + + def testSequenceLossClass(self): + with self.test_session(use_gpu=True): + seq_loss = loss.SequenceLoss(average_across_timesteps=True, + average_across_batch=True, + sum_over_timesteps=False, + sum_over_batch=False) + average_loss_per_example = seq_loss( + self.targets, self.logits, self.weights) + res = self.evaluate(average_loss_per_example) + self.assertAllClose(self.expected_loss, res) + + seq_loss = loss.SequenceLoss(average_across_timesteps=False, + average_across_batch=True, + sum_over_timesteps=False, + sum_over_batch=False) + average_loss_per_sequence = seq_loss( + self.targets, self.logits, self.weights) + res = self.evaluate(average_loss_per_sequence) + compare_per_sequence = np.full((self.sequence_length), self.expected_loss) + self.assertAllClose(compare_per_sequence, res) + + seq_loss = loss.SequenceLoss(average_across_timesteps=True, + average_across_batch=False, + sum_over_timesteps=False, + sum_over_batch=False) + average_loss_per_batch = seq_loss( + self.targets, self.logits, self.weights) + res = self.evaluate(average_loss_per_batch) + compare_per_batch = np.full((self.batch_size), self.expected_loss) + self.assertAllClose(compare_per_batch, res) + + seq_loss = loss.SequenceLoss(average_across_timesteps=False, + average_across_batch=False, + sum_over_timesteps=False, + sum_over_batch=False) + total_loss = seq_loss( + self.targets, self.logits, self.weights) + res = self.evaluate(total_loss) + compare_total = np.full((self.batch_size, self.sequence_length), + self.expected_loss) + self.assertAllClose(compare_total, res) + + def testSumReduction(self): + with self.test_session(use_gpu=True): + seq_loss = loss.SequenceLoss(average_across_timesteps=False, + average_across_batch=False, + sum_over_timesteps=True, + sum_over_batch=True) + average_loss_per_example = seq_loss( + self.targets, self.logits, self.weights) + res = self.evaluate(average_loss_per_example) + self.assertAllClose(self.expected_loss, res) + + seq_loss = loss.SequenceLoss(average_across_timesteps=False, + average_across_batch=False, + sum_over_timesteps=False, + sum_over_batch=True) + average_loss_per_sequence = seq_loss( + self.targets, self.logits, self.weights) + res = self.evaluate(average_loss_per_sequence) + compare_per_sequence = np.full((self.sequence_length), self.expected_loss) + self.assertAllClose(compare_per_sequence, res) + + seq_loss = loss.SequenceLoss(average_across_timesteps=False, + average_across_batch=False, + sum_over_timesteps=True, + sum_over_batch=False) + average_loss_per_batch = seq_loss( + self.targets, self.logits, self.weights) + res = self.evaluate(average_loss_per_batch) + compare_per_batch = np.full((self.batch_size), self.expected_loss) + self.assertAllClose(compare_per_batch, res) + + seq_loss = loss.SequenceLoss(average_across_timesteps=False, + average_across_batch=False, + sum_over_timesteps=False, + sum_over_batch=False) + total_loss = seq_loss( + self.targets, self.logits, self.weights) + res = self.evaluate(total_loss) + compare_total = np.full((self.batch_size, self.sequence_length), + self.expected_loss) + self.assertAllClose(compare_total, res) + + def testWeightedSumReduction(self): + weights = [ + constant_op.constant(1.0, shape=[self.batch_size]) + for _ in range(self.sequence_length) + ] + # Make the last element in the sequence to have zero weights. + weights[-1] = constant_op.constant(0.0, shape=[self.batch_size]) + self.weights = array_ops.stack(weights, axis=1) + with self.test_session(use_gpu=True): + seq_loss = loss.SequenceLoss(average_across_timesteps=False, + average_across_batch=False, + sum_over_timesteps=True, + sum_over_batch=True) + average_loss_per_example = seq_loss( + self.targets, self.logits, self.weights) + res = self.evaluate(average_loss_per_example) + self.assertAllClose(self.expected_loss, res) + + seq_loss = loss.SequenceLoss(average_across_timesteps=False, + average_across_batch=False, + sum_over_timesteps=False, + sum_over_batch=True) + average_loss_per_sequence = seq_loss( + self.targets, self.logits, self.weights) + res = self.evaluate(average_loss_per_sequence) + compare_per_sequence = np.full((self.sequence_length), self.expected_loss) + # The last element in every sequence are zeros, which will be filtered. + compare_per_sequence[-1] = 0. + self.assertAllClose(compare_per_sequence, res) + + seq_loss = loss.SequenceLoss(average_across_timesteps=False, + average_across_batch=False, + sum_over_timesteps=True, + sum_over_batch=False) + average_loss_per_batch = seq_loss(self.targets, self.logits, self.weights) + res = self.evaluate(average_loss_per_batch) + compare_per_batch = np.full((self.batch_size), self.expected_loss) + self.assertAllClose(compare_per_batch, res) + + seq_loss = loss.SequenceLoss(average_across_timesteps=False, + average_across_batch=False, + sum_over_timesteps=False, + sum_over_batch=False) + total_loss = seq_loss(self.targets, self.logits, self.weights) + res = self.evaluate(total_loss) + compare_total = np.full((self.batch_size, self.sequence_length), + self.expected_loss) + # The last element in every sequence are zeros, which will be filtered. + compare_total[:, -1] = 0 + self.assertAllClose(compare_total, res) + + def testZeroWeights(self): + weights = [ + constant_op.constant(0.0, shape=[self.batch_size]) + for _ in range(self.sequence_length) + ] + weights = array_ops.stack(weights, axis=1) + with self.test_session(use_gpu=True): + average_loss_per_example = loss.sequence_loss( + self.logits, self.targets, weights, + average_across_timesteps=True, + average_across_batch=True) + res = self.evaluate(average_loss_per_example) + self.assertAllClose(0.0, res) + + average_loss_per_sequence = loss.sequence_loss( + self.logits, self.targets, weights, + average_across_timesteps=False, + average_across_batch=True) + res = self.evaluate(average_loss_per_sequence) + compare_per_sequence = np.zeros((self.sequence_length)) + self.assertAllClose(compare_per_sequence, res) + + average_loss_per_batch = loss.sequence_loss( + self.logits, self.targets, weights, + average_across_timesteps=True, + average_across_batch=False) + res = self.evaluate(average_loss_per_batch) + compare_per_batch = np.zeros((self.batch_size)) + self.assertAllClose(compare_per_batch, res) + + total_loss = loss.sequence_loss( + self.logits, self.targets, weights, + average_across_timesteps=False, + average_across_batch=False) + res = self.evaluate(total_loss) + compare_total = np.zeros((self.batch_size, self.sequence_length)) + self.assertAllClose(compare_total, res) + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py b/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py index 60ec3efffe771a3a6d6f36ed4b51a34ef9509612..5bcf0af8897ba8bc868951d03a18081e24a00f35 100644 --- a/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py +++ b/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py @@ -28,6 +28,9 @@ from tensorflow.contrib.framework.python.framework import tensor_util from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape +from tensorflow.python.keras import initializers +from tensorflow.python.keras import layers +from tensorflow.python.keras.engine import base_layer_utils from tensorflow.python.layers import base as layers_base from tensorflow.python.layers import core as layers_core from tensorflow.python.ops import array_ops @@ -72,77 +75,6 @@ class AttentionMechanism(object): raise NotImplementedError -def _prepare_memory(memory, memory_sequence_length, check_inner_dims_defined): - """Convert to tensor and possibly mask `memory`. - - Args: - memory: `Tensor`, shaped `[batch_size, max_time, ...]`. - memory_sequence_length: `int32` `Tensor`, shaped `[batch_size]`. - check_inner_dims_defined: Python boolean. If `True`, the `memory` - argument's shape is checked to ensure all but the two outermost - dimensions are fully defined. - - Returns: - A (possibly masked), checked, new `memory`. - - Raises: - ValueError: If `check_inner_dims_defined` is `True` and not - `memory.shape[2:].is_fully_defined()`. - """ - memory = nest.map_structure( - lambda m: ops.convert_to_tensor(m, name="memory"), memory) - if memory_sequence_length is not None: - memory_sequence_length = ops.convert_to_tensor( - memory_sequence_length, name="memory_sequence_length") - if check_inner_dims_defined: - def _check_dims(m): - if not m.get_shape()[2:].is_fully_defined(): - raise ValueError("Expected memory %s to have fully defined inner dims, " - "but saw shape: %s" % (m.name, m.get_shape())) - nest.map_structure(_check_dims, memory) - if memory_sequence_length is None: - seq_len_mask = None - else: - seq_len_mask = array_ops.sequence_mask( - memory_sequence_length, - maxlen=array_ops.shape(nest.flatten(memory)[0])[1], - dtype=nest.flatten(memory)[0].dtype) - seq_len_batch_size = ( - tensor_shape.dimension_value(memory_sequence_length.shape[0]) - or array_ops.shape(memory_sequence_length)[0]) - def _maybe_mask(m, seq_len_mask): - rank = m.get_shape().ndims - rank = rank if rank is not None else array_ops.rank(m) - extra_ones = array_ops.ones(rank - 2, dtype=dtypes.int32) - m_batch_size = tensor_shape.dimension_value( - m.shape[0]) or array_ops.shape(m)[0] - if memory_sequence_length is not None: - message = ("memory_sequence_length and memory tensor batch sizes do not " - "match.") - with ops.control_dependencies([ - check_ops.assert_equal( - seq_len_batch_size, m_batch_size, message=message)]): - seq_len_mask = array_ops.reshape( - seq_len_mask, - array_ops.concat((array_ops.shape(seq_len_mask), extra_ones), 0)) - return m * seq_len_mask - else: - return m - return nest.map_structure(lambda m: _maybe_mask(m, seq_len_mask), memory) - - -def _maybe_mask_score(score, memory_sequence_length, score_mask_value): - if memory_sequence_length is None: - return score - message = ("All values in memory_sequence_length must greater than zero.") - with ops.control_dependencies( - [check_ops.assert_positive(memory_sequence_length, message=message)]): - score_mask = array_ops.sequence_mask( - memory_sequence_length, maxlen=array_ops.shape(score)[1]) - score_mask_values = score_mask_value * array_ops.ones_like(score) - return array_ops.where(score_mask, score, score_mask_values) - - class _BaseAttentionMechanism(AttentionMechanism): """A base AttentionMechanism class providing common functionality. @@ -205,12 +137,14 @@ class _BaseAttentionMechanism(AttentionMechanism): self._memory_layer.dtype).as_numpy_dtype(-np.inf) self._probability_fn = lambda score, prev: ( # pylint:disable=g-long-lambda probability_fn( - _maybe_mask_score(score, memory_sequence_length, score_mask_value), + _maybe_mask_score(score, + memory_sequence_length=memory_sequence_length, + score_mask_value=score_mask_value), prev)) with ops.name_scope( name, "BaseAttentionMechanismInit", nest.flatten(memory)): self._values = _prepare_memory( - memory, memory_sequence_length, + memory, memory_sequence_length=memory_sequence_length, check_inner_dims_defined=check_inner_dims_defined) self._keys = ( self.memory_layer(self._values) if self.memory_layer # pylint: disable=not-callable @@ -286,6 +220,376 @@ class _BaseAttentionMechanism(AttentionMechanism): return self.initial_alignments(batch_size, dtype) +class _BaseAttentionMechanismV2(AttentionMechanism, layers.Layer): + """A base AttentionMechanism class providing common functionality. + + Common functionality includes: + 1. Storing the query and memory layers. + 2. Preprocessing and storing the memory. + + Note that this layer takes memory as its init parameter, which is an + anti-pattern of Keras API, we have to keep the memory as init parameter for + performance and dependency reason. Under the hood, during `__init__()`, it + will invoke `base_layer.__call__(memory, setup_memory=True)`. This will let + keras to keep track of the memory tensor as the input of this layer. Once + the `__init__()` is done, then user can query the attention by + `score = att_obj([query, state])`, and use it as a normal keras layer. + + Special attention is needed when adding using this class as the base layer for + new attention: + 1. Build() could be invoked at least twice. So please make sure weights are + not duplicated. + 2. Layer.get_weights() might return different set of weights if the instance + has `query_layer`. The query_layer weights is not initialized until the + memory is configured. + + Also note that this layer does not work with Keras model when + `model.compile(run_eagerly=True)` due to the fact that this layer is stateful. + The support for that will be added in a future version. + """ + + def __init__(self, + memory, + probability_fn, + query_layer=None, + memory_layer=None, + memory_sequence_length=None, + **kwargs): + """Construct base AttentionMechanism class. + + Args: + memory: The memory to query; usually the output of an RNN encoder. This + tensor should be shaped `[batch_size, max_time, ...]`. + probability_fn: A `callable`. Converts the score and previous alignments + to probabilities. Its signature should be: + `probabilities = probability_fn(score, state)`. + query_layer: (optional): Instance of `tf.keras.Layer`. The layer's depth + must match the depth of `memory_layer`. If `query_layer` is not + provided, the shape of `query` must match that of `memory_layer`. + memory_layer: (optional): Instance of `tf.keras.Layer`. The layer's + depth must match the depth of `query_layer`. + If `memory_layer` is not provided, the shape of `memory` must match + that of `query_layer`. + memory_sequence_length (optional): Sequence lengths for the batch entries + in memory. If provided, the memory tensor rows are masked with zeros + for values past the respective sequence lengths. + **kwargs: Dictionary that contains other common arguments for layer + creation. + """ + if (query_layer is not None + and not isinstance(query_layer, layers.Layer)): + raise TypeError( + "query_layer is not a Layer: %s" % type(query_layer).__name__) + if (memory_layer is not None + and not isinstance(memory_layer, layers.Layer)): + raise TypeError( + "memory_layer is not a Layer: %s" % type(memory_layer).__name__) + self.query_layer = query_layer + self.memory_layer = memory_layer + if self.memory_layer is not None and "dtype" not in kwargs: + kwargs["dtype"] = self.memory_layer.dtype + super(_BaseAttentionMechanismV2, self).__init__(**kwargs) + if not callable(probability_fn): + raise TypeError("probability_fn must be callable, saw type: %s" % + type(probability_fn).__name__) + self.probability_fn = probability_fn + + self.keys = None + self.values = None + self.batch_size = None + self._memory_initialized = False + self._check_inner_dims_defined = True + self.supports_masking = True + self.score_mask_value = dtypes.as_dtype(self.dtype).as_numpy_dtype(-np.inf) + + if memory is not None: + # Setup the memory by self.__call__() with memory and memory_seq_length. + # This will make the attention follow the keras convention which takes + # all the tensor inputs via __call__(). + if memory_sequence_length is None: + inputs = memory + else: + inputs = [memory, memory_sequence_length] + + self.values = super(_BaseAttentionMechanismV2, self).__call__( + inputs, setup_memory=True) + + def build(self, input_shape): + if not self._memory_initialized: + # This is for setting up the memory, which contains memory and optional + # memory_sequence_length. Build the memory_layer with memory shape. + if self.memory_layer is not None and not self.memory_layer.built: + if isinstance(input_shape, list): + self.memory_layer.build(input_shape[0]) + else: + self.memory_layer.build(input_shape) + else: + # The input_shape should be query.shape and state.shape. Use the query + # to init the query layer. + if self.query_layer is not None and not self.query_layer.built: + self.query_layer.build(input_shape[0]) + + def __call__(self, inputs, **kwargs): + """Preprocess the inputs before calling `base_layer.__call__()`. + + Note that there are situation here, one for setup memory, and one with + actual query and state. + 1. When the memory has not been configured, we just pass all the param to + base_layer.__call__(), which will then invoke self.call() with proper + inputs, which allows this class to setup memory. + 2. When the memory has already been setup, the input should contain query + and state, and optionally processed memory. If the processed memory is + not included in the input, we will have to append it to the inputs and + give it to the base_layer.__call__(). The processed memory is the output + of first invocation of self.__call__(). If we don't add it here, then from + keras perspective, the graph is disconnected since the output from + previous call is never used. + + Args: + inputs: the inputs tensors. + **kwargs: dict, other keyeword arguments for the `__call__()` + """ + if self._memory_initialized: + if len(inputs) not in (2, 3): + raise ValueError("Expect the inputs to have 2 or 3 tensors, got %d" % + len(inputs)) + if len(inputs) == 2: + # We append the calculated memory here so that the graph will be + # connected. + inputs.append(self.values) + return super(_BaseAttentionMechanismV2, self).__call__(inputs, **kwargs) + + def call(self, inputs, mask=None, setup_memory=False, **kwargs): + """Setup the memory or query the attention. + + There are two case here, one for setup memory, and the second is query the + attention score. `setup_memory` is the flag to indicate which mode it is. + The input list will be treated differently based on that flag. + + Args: + inputs: a list of tensor that could either be `query` and `state`, or + `memory` and `memory_sequence_length`. + `query` is the tensor of dtype matching `memory` and shape + `[batch_size, query_depth]`. + `state` is the tensor of dtype matching `memory` and shape + `[batch_size, alignments_size]`. (`alignments_size` is memory's + `max_time`). + `memory` is the memory to query; usually the output of an RNN encoder. + The tensor should be shaped `[batch_size, max_time, ...]`. + `memory_sequence_length` (optional) is the sequence lengths for the + batch entries in memory. If provided, the memory tensor rows are masked + with zeros for values past the respective sequence lengths. + mask: optional bool tensor with shape `[batch, max_time]` for the mask of + memory. If it is not None, the corresponding item of the memory should + be filtered out during calculation. + setup_memory: boolean, whether the input is for setting up memory, or + query attention. + **kwargs: Dict, other keyword arguments for the call method. + Returns: + Either processed memory or attention score, based on `setup_memory`. + """ + if setup_memory: + if isinstance(inputs, list): + if len(inputs) not in (1, 2): + raise ValueError("Expect inputs to have 1 or 2 tensors, got %d" % + len(inputs)) + memory = inputs[0] + memory_sequence_length = inputs[1] if len(inputs) == 2 else None + memory_mask = mask + else: + memory, memory_sequence_length = inputs, None + memory_mask = mask + self._setup_memory(memory, memory_sequence_length, memory_mask) + # We force the self.built to false here since only memory is initialized, + # but the real query/state has not been call() yet. The layer should be + # build and call again. + self.built = False + # Return the processed memory in order to create the Keras connectivity + # data for it. + return self.values + else: + if not self._memory_initialized: + raise ValueError("Cannot query the attention before the setup of " + "memory") + if len(inputs) not in (2, 3): + raise ValueError("Expect the inputs to have query, state, and optional " + "processed memory, got %d items" % len(inputs)) + # Ignore the rest of the inputs and only care about the query and state + query, state = inputs[0], inputs[1] + return self._calculate_attention(query, state) + + def _setup_memory(self, memory, memory_sequence_length=None, + memory_mask=None): + """Pre-process the memory before actually query the memory. + + This should only be called once at the first invocation of call(). + + Args: + memory: The memory to query; usually the output of an RNN encoder. This + tensor should be shaped `[batch_size, max_time, ...]`. + memory_sequence_length (optional): Sequence lengths for the batch entries + in memory. If provided, the memory tensor rows are masked with zeros for + values past the respective sequence lengths. + memory_mask: (Optional) The boolean tensor with shape `[batch_size, + max_time]`. For any value equal to False, the corresponding value in + memory should be ignored. + """ + if self._memory_initialized: + raise ValueError("The memory for the attention has already been setup.") + if memory_sequence_length is not None and memory_mask is not None: + raise ValueError("memory_sequence_length and memory_mask cannot be " + "used at same time for attention.") + with ops.name_scope( + self.name, "BaseAttentionMechanismInit", nest.flatten(memory)): + self.values = _prepare_memory( + memory, + memory_sequence_length=memory_sequence_length, + memory_mask=memory_mask, + check_inner_dims_defined=self._check_inner_dims_defined) + # Mark the value as check since the memory and memory mask might not + # passed from __call__(), which does not have proper keras metadata. + # TODO(omalleyt): Remove this hack once the mask the has proper keras + # history. + base_layer_utils.mark_checked(self.values) + if self.memory_layer is not None: + self.keys = self.memory_layer(self.values) + else: + self.keys = self.values + self.batch_size = ( + tensor_shape.dimension_value(self.keys.shape[0]) or + array_ops.shape(self.keys)[0]) + self._alignments_size = (tensor_shape.dimension_value(self.keys.shape[1]) + or array_ops.shape(self.keys)[1]) + if memory_mask is not None: + unwrapped_probability_fn = self.probability_fn + def _mask_probability_fn(score, prev): + return unwrapped_probability_fn( + _maybe_mask_score( + score, + memory_mask=memory_mask, + memory_sequence_length=memory_sequence_length, + score_mask_value=self.score_mask_value), prev) + self.probability_fn = _mask_probability_fn + self._memory_initialized = True + + def _calculate_attention(self, query, state): + raise NotImplementedError( + "_calculate_attention need to be implemented by subclasses.") + + def compute_mask(self, inputs, mask=None): + # There real input of the attention is query and state, and the memory layer + # mask shouldn't be pass down. Returning None for all output mask here. + return None, None + + def get_config(self): + config = {} + # Since the probability_fn is likely to be a wrapped function, the child + # class should preserve the original function and how its wrapped. + + if self.query_layer is not None: + config["query_layer"] = { + "class_name": self.query_layer.__class__.__name__, + "config": self.query_layer.get_config(), + } + if self.memory_layer is not None: + config["memory_layer"] = { + "class_name": self.memory_layer.__class__.__name__, + "config": self.memory_layer.get_config(), + } + # memory is a required init parameter and its a tensor. It cannot be + # serialized to config, so we put a placeholder for it. + config["memory"] = None + base_config = super(_BaseAttentionMechanismV2, self).get_config() + return dict(list(base_config.items()) + list(config.items())) + + def _process_probability_fn(self, func_name): + """Helper method to retrieve the probably function by string input.""" + valid_probability_fns = { + "softmax": nn_ops.softmax, + "hardmax": hardmax, + } + if func_name not in valid_probability_fns.keys(): + raise ValueError("Invalid probability function: %s, options are %s" % + (func_name, valid_probability_fns.keys())) + return valid_probability_fns[func_name] + + @classmethod + def deserialize_inner_layer_from_config(cls, config, custom_objects): + """Helper method that reconstruct the query and memory from the config. + + In the get_config() method, the query and memory layer configs are + serialized into dict for persistence, this method perform the reverse action + to reconstruct the layer from the config. + + Args: + config: dict, the configs that will be used to reconstruct the object. + custom_objects: dict mapping class names (or function names) of custom + (non-Keras) objects to class/functions. + Returns: + config: dict, the config with layer instance created, which is ready to be + used as init parameters. + """ + # Reconstruct the query and memory layer for parent class. + from tensorflow.python.keras.layers import deserialize as deserialize_layer # pylint: disable=g-import-not-at-top + # Instead of updating the input, create a copy and use that. + config = config.copy() + query_layer_config = config.pop("query_layer", None) + if query_layer_config: + query_layer = deserialize_layer(query_layer_config, + custom_objects=custom_objects) + config["query_layer"] = query_layer + memory_layer_config = config.pop("memory_layer", None) + if memory_layer_config: + memory_layer = deserialize_layer(memory_layer_config, + custom_objects=custom_objects) + config["memory_layer"] = memory_layer + return config + + @property + def alignments_size(self): + return self._alignments_size + + @property + def state_size(self): + return self._alignments_size + + def initial_alignments(self, batch_size, dtype): + """Creates the initial alignment values for the `AttentionWrapper` class. + + This is important for AttentionMechanisms that use the previous alignment + to calculate the alignment at the next time step (e.g. monotonic attention). + + The default behavior is to return a tensor of all zeros. + + Args: + batch_size: `int32` scalar, the batch_size. + dtype: The `dtype`. + + Returns: + A `dtype` tensor shaped `[batch_size, alignments_size]` + (`alignments_size` is the values' `max_time`). + """ + max_time = self._alignments_size + return _zero_state_tensors(max_time, batch_size, dtype) + + def initial_state(self, batch_size, dtype): + """Creates the initial state values for the `AttentionWrapper` class. + + This is important for AttentionMechanisms that use the previous alignment + to calculate the alignment at the next time step (e.g. monotonic attention). + + The default behavior is to return the same output as initial_alignments. + + Args: + batch_size: `int32` scalar, the batch_size. + dtype: The `dtype`. + + Returns: + A structure of all-zero tensors with shapes as described by `state_size`. + """ + return self.initial_alignments(batch_size, dtype) + + def _luong_score(query, keys, scale): """Implements Luong-style (multiplicative) scoring function. @@ -304,7 +608,7 @@ def _luong_score(query, keys, scale): Args: query: Tensor, shape `[batch_size, num_units]` to compare to keys. keys: Processed memory, shape `[batch_size, max_time, num_units]`. - scale: Whether to apply a scale to the score function. + scale: the optional tensor to scale the attention score. Returns: A `[batch_size, max_time]` tensor of unnormalized score values. @@ -320,7 +624,6 @@ def _luong_score(query, keys, scale): "Query (%s) has units: %s. Keys (%s) have units: %s. " "Perhaps you need to set num_units to the keys' dimension (%s)?" % (query, depth, keys, key_units, key_units)) - dtype = query.dtype # Reshape from [batch_size, depth] to [batch_size, 1, depth] # for matmul. @@ -338,12 +641,8 @@ def _luong_score(query, keys, scale): score = math_ops.matmul(query, keys, transpose_b=True) score = array_ops.squeeze(score, [1]) - if scale: - # Scalar used in weight scaling - g = variable_scope.get_variable( - "attention_g", dtype=dtype, - initializer=init_ops.ones_initializer, shape=()) - score = g * score + if scale is not None: + score = scale * score return score @@ -354,8 +653,8 @@ class LuongAttention(_BaseAttentionMechanism): as described in: Minh-Thang Luong, Hieu Pham, Christopher D. Manning. - "Effective Approaches to Attention-based Neural Machine Translation." - EMNLP 2015. https://arxiv.org/abs/1508.04025 + [Effective Approaches to Attention-based Neural Machine Translation. + EMNLP 2015.](https://arxiv.org/abs/1508.04025) The second is the scaled form inspired partly by the normalized form of Bahdanau attention. @@ -429,13 +728,133 @@ class LuongAttention(_BaseAttentionMechanism): `max_time`). """ with variable_scope.variable_scope(None, "luong_attention", [query]): - score = _luong_score(query, self._keys, self._scale) + attention_g = None + if self._scale: + attention_g = variable_scope.get_variable( + "attention_g", dtype=query.dtype, + initializer=init_ops.ones_initializer, shape=()) + score = _luong_score(query, self._keys, attention_g) alignments = self._probability_fn(score, state) next_state = alignments return alignments, next_state -def _bahdanau_score(processed_query, keys, normalize): +class LuongAttentionV2(_BaseAttentionMechanismV2): + """Implements Luong-style (multiplicative) attention scoring. + + This attention has two forms. The first is standard Luong attention, + as described in: + + Minh-Thang Luong, Hieu Pham, Christopher D. Manning. + [Effective Approaches to Attention-based Neural Machine Translation. + EMNLP 2015.](https://arxiv.org/abs/1508.04025) + + The second is the scaled form inspired partly by the normalized form of + Bahdanau attention. + + To enable the second form, construct the object with parameter + `scale=True`. + """ + + def __init__(self, + units, + memory, + memory_sequence_length=None, + scale=False, + probability_fn="softmax", + dtype=None, + name="LuongAttention", + **kwargs): + """Construct the AttentionMechanism mechanism. + + Args: + units: The depth of the attention mechanism. + memory: The memory to query; usually the output of an RNN encoder. This + tensor should be shaped `[batch_size, max_time, ...]`. + memory_sequence_length: (optional): Sequence lengths for the batch entries + in memory. If provided, the memory tensor rows are masked with zeros + for values past the respective sequence lengths. + scale: Python boolean. Whether to scale the energy term. + probability_fn: (optional) string, the name of function to convert the + attention score to probabilities. The default is `softmax` which is + `tf.nn.softmax`. Other options is `hardmax`, which is hardmax() within + this module. Any other value will result intovalidation error. Default + to use `softmax`. + dtype: The data type for the memory layer of the attention mechanism. + name: Name to use when creating ops. + **kwargs: Dictionary that contains other common arguments for layer + creation. + """ + # For LuongAttention, we only transform the memory layer; thus + # num_units **must** match expected the query depth. + self.probability_fn_name = probability_fn + probability_fn = self._process_probability_fn(self.probability_fn_name) + wrapped_probability_fn = lambda score, _: probability_fn(score) + if dtype is None: + dtype = dtypes.float32 + memory_layer = kwargs.pop("memory_layer", None) + if not memory_layer: + memory_layer = layers.Dense( + units, name="memory_layer", use_bias=False, dtype=dtype) + self.units = units + self.scale = scale + self.scale_weight = None + super(LuongAttentionV2, self).__init__( + memory=memory, + memory_sequence_length=memory_sequence_length, + query_layer=None, + memory_layer=memory_layer, + probability_fn=wrapped_probability_fn, + name=name, + dtype=dtype, + **kwargs) + + def build(self, input_shape): + super(LuongAttentionV2, self).build(input_shape) + if self.scale and self.scale_weight is None: + self.scale_weight = self.add_weight( + "attention_g", initializer=init_ops.ones_initializer, shape=()) + self.built = True + + def _calculate_attention(self, query, state): + """Score the query based on the keys and values. + + Args: + query: Tensor of dtype matching `self.values` and shape + `[batch_size, query_depth]`. + state: Tensor of dtype matching `self.values` and shape + `[batch_size, alignments_size]` + (`alignments_size` is memory's `max_time`). + + Returns: + alignments: Tensor of dtype matching `self.values` and shape + `[batch_size, alignments_size]` (`alignments_size` is memory's + `max_time`). + next_state: Same as the alignments. + """ + score = _luong_score(query, self.keys, self.scale_weight) + alignments = self.probability_fn(score, state) + next_state = alignments + return alignments, next_state + + def get_config(self): + config = { + "units": self.units, + "scale": self.scale, + "probability_fn": self.probability_fn_name, + } + base_config = super(LuongAttentionV2, self).get_config() + return dict(list(base_config.items()) + list(config.items())) + + @classmethod + def from_config(cls, config, custom_objects=None): + config = _BaseAttentionMechanismV2.deserialize_inner_layer_from_config( + config, custom_objects=custom_objects) + return cls(**config) + + +def _bahdanau_score(processed_query, keys, attention_v, + attention_g=None, attention_b=None): """Implements Bahdanau-style (additive) scoring function. This attention has two forms. The first is Bhandanau attention, @@ -453,41 +872,28 @@ def _bahdanau_score(processed_query, keys, normalize): Training of Deep Neural Networks." https://arxiv.org/abs/1602.07868 - To enable the second form, set `normalize=True`. + To enable the second form, set please pass in attention_g and attention_b. Args: processed_query: Tensor, shape `[batch_size, num_units]` to compare to keys. keys: Processed memory, shape `[batch_size, max_time, num_units]`. - normalize: Whether to normalize the score function. + attention_v: Tensor, shape `[num_units]`. + attention_g: Optional scalar tensor for normalization. + attention_b: Optional tensor with shape `[num_units]` for normalization. Returns: A `[batch_size, max_time]` tensor of unnormalized score values. """ - dtype = processed_query.dtype - # Get the number of hidden units from the trailing dimension of keys - num_units = tensor_shape.dimension_value( - keys.shape[2]) or array_ops.shape(keys)[2] # Reshape from [batch_size, ...] to [batch_size, 1, ...] for broadcasting. processed_query = array_ops.expand_dims(processed_query, 1) - v = variable_scope.get_variable( - "attention_v", [num_units], dtype=dtype) - if normalize: - # Scalar used in weight normalization - g = variable_scope.get_variable( - "attention_g", dtype=dtype, - initializer=init_ops.constant_initializer(math.sqrt((1. / num_units))), - shape=()) - # Bias added prior to the nonlinearity - b = variable_scope.get_variable( - "attention_b", [num_units], dtype=dtype, - initializer=init_ops.zeros_initializer()) - # normed_v = g * v / ||v|| - normed_v = g * v * math_ops.rsqrt( - math_ops.reduce_sum(math_ops.square(v))) + if attention_g is not None and attention_b is not None: + normed_v = attention_g * attention_v * math_ops.rsqrt( + math_ops.reduce_sum(math_ops.square(attention_v))) return math_ops.reduce_sum( - normed_v * math_ops.tanh(keys + processed_query + b), [2]) + normed_v * math_ops.tanh(keys + processed_query + attention_b), [2]) else: - return math_ops.reduce_sum(v * math_ops.tanh(keys + processed_query), [2]) + return math_ops.reduce_sum( + attention_v * math_ops.tanh(keys + processed_query), [2]) class BahdanauAttention(_BaseAttentionMechanism): @@ -578,12 +984,169 @@ class BahdanauAttention(_BaseAttentionMechanism): """ with variable_scope.variable_scope(None, "bahdanau_attention", [query]): processed_query = self.query_layer(query) if self.query_layer else query - score = _bahdanau_score(processed_query, self._keys, self._normalize) + attention_v = variable_scope.get_variable( + "attention_v", [self._num_units], dtype=query.dtype) + if not self._normalize: + attention_g = None + attention_b = None + else: + attention_g = variable_scope.get_variable( + "attention_g", dtype=query.dtype, + initializer=init_ops.constant_initializer( + math.sqrt((1. / self._num_units))), + shape=()) + attention_b = variable_scope.get_variable( + "attention_b", [self._num_units], dtype=query.dtype, + initializer=init_ops.zeros_initializer()) + + score = _bahdanau_score(processed_query, self._keys, attention_v, + attention_g=attention_g, attention_b=attention_b) alignments = self._probability_fn(score, state) next_state = alignments return alignments, next_state +class BahdanauAttentionV2(_BaseAttentionMechanismV2): + """Implements Bahdanau-style (additive) attention. + + This attention has two forms. The first is Bahdanau attention, + as described in: + + Dzmitry Bahdanau, Kyunghyun Cho, Yoshua Bengio. + "Neural Machine Translation by Jointly Learning to Align and Translate." + ICLR 2015. https://arxiv.org/abs/1409.0473 + + The second is the normalized form. This form is inspired by the + weight normalization article: + + Tim Salimans, Diederik P. Kingma. + "Weight Normalization: A Simple Reparameterization to Accelerate + Training of Deep Neural Networks." + https://arxiv.org/abs/1602.07868 + + To enable the second form, construct the object with parameter + `normalize=True`. + """ + + def __init__(self, + units, + memory, + memory_sequence_length=None, + normalize=False, + probability_fn="softmax", + kernel_initializer="glorot_uniform", + dtype=None, + name="BahdanauAttention", + **kwargs): + """Construct the Attention mechanism. + + Args: + units: The depth of the query mechanism. + memory: The memory to query; usually the output of an RNN encoder. This + tensor should be shaped `[batch_size, max_time, ...]`. + memory_sequence_length: (optional): Sequence lengths for the batch entries + in memory. If provided, the memory tensor rows are masked with zeros + for values past the respective sequence lengths. + normalize: Python boolean. Whether to normalize the energy term. + probability_fn: (optional) string, the name of function to convert the + attention score to probabilities. The default is `softmax` which is + `tf.nn.softmax`. Other options is `hardmax`, which is hardmax() within + this module. Any other value will result into validation error. Default + to use `softmax`. + kernel_initializer: (optional), the name of the initializer for the + attention kernel. + dtype: The data type for the query and memory layers of the attention + mechanism. + name: Name to use when creating ops. + **kwargs: Dictionary that contains other common arguments for layer + creation. + """ + self.probability_fn_name = probability_fn + probability_fn = self._process_probability_fn(self.probability_fn_name) + wrapped_probability_fn = lambda score, _: probability_fn(score) + if dtype is None: + dtype = dtypes.float32 + query_layer = kwargs.pop("query_layer", None) + if not query_layer: + query_layer = layers.Dense( + units, name="query_layer", use_bias=False, dtype=dtype) + memory_layer = kwargs.pop("memory_layer", None) + if not memory_layer: + memory_layer = layers.Dense( + units, name="memory_layer", use_bias=False, dtype=dtype) + self.units = units + self.normalize = normalize + self.kernel_initializer = initializers.get(kernel_initializer) + self.attention_v = None + self.attention_g = None + self.attention_b = None + super(BahdanauAttentionV2, self).__init__( + memory=memory, + memory_sequence_length=memory_sequence_length, + query_layer=query_layer, + memory_layer=memory_layer, + probability_fn=wrapped_probability_fn, + name=name, + dtype=dtype, + **kwargs) + + def build(self, input_shape): + super(BahdanauAttentionV2, self).build(input_shape) + if self.attention_v is None: + self.attention_v = self.add_weight( + "attention_v", [self.units], + dtype=self.dtype, + initializer=self.kernel_initializer) + if self.normalize and self.attention_g is None and self.attention_b is None: + self.attention_g = self.add_weight( + "attention_g", initializer=init_ops.constant_initializer( + math.sqrt((1. / self.units))), shape=()) + self.attention_b = self.add_weight( + "attention_b", shape=[self.units], + initializer=init_ops.zeros_initializer()) + self.built = True + + def _calculate_attention(self, query, state): + """Score the query based on the keys and values. + + Args: + query: Tensor of dtype matching `self.values` and shape + `[batch_size, query_depth]`. + state: Tensor of dtype matching `self.values` and shape + `[batch_size, alignments_size]` + (`alignments_size` is memory's `max_time`). + + Returns: + alignments: Tensor of dtype matching `self.values` and shape + `[batch_size, alignments_size]` (`alignments_size` is memory's + `max_time`). + next_state: same as alignments. + """ + processed_query = self.query_layer(query) if self.query_layer else query + score = _bahdanau_score(processed_query, self.keys, self.attention_v, + attention_g=self.attention_g, + attention_b=self.attention_b) + alignments = self.probability_fn(score, state) + next_state = alignments + return alignments, next_state + + def get_config(self): + config = { + "units": self.units, + "normalize": self.normalize, + "probability_fn": self.probability_fn_name, + "kernel_initializer": initializers.serialize(self.kernel_initializer) + } + base_config = super(BahdanauAttentionV2, self).get_config() + return dict(list(base_config.items()) + list(config.items())) + + @classmethod + def from_config(cls, config, custom_objects=None): + config = _BaseAttentionMechanismV2.deserialize_inner_layer_from_config( + config, custom_objects=custom_objects) + return cls(**config) + + def safe_cumprod(x, *args, **kwargs): """Computes cumprod of x in logspace using cumsum to avoid underflow. @@ -766,6 +1329,34 @@ class _BaseMonotonicAttentionMechanism(_BaseAttentionMechanism): dtype=dtype) +class _BaseMonotonicAttentionMechanismV2(_BaseAttentionMechanismV2): + """Base attention mechanism for monotonic attention. + + Simply overrides the initial_alignments function to provide a dirac + distribution, which is needed in order for the monotonic attention + distributions to have the correct behavior. + """ + + def initial_alignments(self, batch_size, dtype): + """Creates the initial alignment values for the monotonic attentions. + + Initializes to dirac distributions, i.e. [1, 0, 0, ...memory length..., 0] + for all entries in the batch. + + Args: + batch_size: `int32` scalar, the batch_size. + dtype: The `dtype`. + + Returns: + A `dtype` tensor shaped `[batch_size, alignments_size]` + (`alignments_size` is the values' `max_time`). + """ + max_time = self._alignments_size + return array_ops.one_hot( + array_ops.zeros((batch_size,), dtype=dtypes.int32), max_time, + dtype=dtype) + + class BahdanauMonotonicAttention(_BaseMonotonicAttentionMechanism): """Monotonic attention mechanism with Bahadanau-style energy function. @@ -860,7 +1451,22 @@ class BahdanauMonotonicAttention(_BaseMonotonicAttentionMechanism): with variable_scope.variable_scope( None, "bahdanau_monotonic_attention", [query]): processed_query = self.query_layer(query) if self.query_layer else query - score = _bahdanau_score(processed_query, self._keys, self._normalize) + attention_v = variable_scope.get_variable( + "attention_v", [self._num_units], dtype=query.dtype) + if not self._normalize: + attention_g = None + attention_b = None + else: + attention_g = variable_scope.get_variable( + "attention_g", dtype=query.dtype, + initializer=init_ops.constant_initializer( + math.sqrt((1. / self._num_units))), + shape=()) + attention_b = variable_scope.get_variable( + "attention_b", [self._num_units], dtype=query.dtype, + initializer=init_ops.zeros_initializer()) + score = _bahdanau_score(processed_query, self._keys, attention_v, + attention_g=attention_g, attention_b=attention_b) score_bias = variable_scope.get_variable( "attention_score_bias", dtype=processed_query.dtype, initializer=self._score_bias_init) @@ -870,6 +1476,164 @@ class BahdanauMonotonicAttention(_BaseMonotonicAttentionMechanism): return alignments, next_state +class BahdanauMonotonicAttentionV2(_BaseMonotonicAttentionMechanismV2): + """Monotonic attention mechanism with Bahadanau-style energy function. + + This type of attention enforces a monotonic constraint on the attention + distributions; that is once the model attends to a given point in the memory + it can't attend to any prior points at subsequence output timesteps. It + achieves this by using the _monotonic_probability_fn instead of softmax to + construct its attention distributions. Since the attention scores are passed + through a sigmoid, a learnable scalar bias parameter is applied after the + score function and before the sigmoid. Otherwise, it is equivalent to + BahdanauAttention. This approach is proposed in + + Colin Raffel, Minh-Thang Luong, Peter J. Liu, Ron J. Weiss, Douglas Eck, + "Online and Linear-Time Attention by Enforcing Monotonic Alignments." + ICML 2017. https://arxiv.org/abs/1704.00784 + """ + + def __init__(self, + units, + memory, + memory_sequence_length=None, + normalize=False, + sigmoid_noise=0., + sigmoid_noise_seed=None, + score_bias_init=0., + mode="parallel", + kernel_initializer="glorot_uniform", + dtype=None, + name="BahdanauMonotonicAttention", + **kwargs): + """Construct the Attention mechanism. + + Args: + units: The depth of the query mechanism. + memory: The memory to query; usually the output of an RNN encoder. This + tensor should be shaped `[batch_size, max_time, ...]`. + memory_sequence_length: (optional): Sequence lengths for the batch entries + in memory. If provided, the memory tensor rows are masked with zeros + for values past the respective sequence lengths. + normalize: Python boolean. Whether to normalize the energy term. + sigmoid_noise: Standard deviation of pre-sigmoid noise. See the docstring + for `_monotonic_probability_fn` for more information. + sigmoid_noise_seed: (optional) Random seed for pre-sigmoid noise. + score_bias_init: Initial value for score bias scalar. It's recommended to + initialize this to a negative value when the length of the memory is + large. + mode: How to compute the attention distribution. Must be one of + 'recursive', 'parallel', or 'hard'. See the docstring for + `tf.contrib.seq2seq.monotonic_attention` for more information. + kernel_initializer: (optional), the name of the initializer for the + attention kernel. + dtype: The data type for the query and memory layers of the attention + mechanism. + name: Name to use when creating ops. + **kwargs: Dictionary that contains other common arguments for layer + creation. + """ + # Set up the monotonic probability fn with supplied parameters + if dtype is None: + dtype = dtypes.float32 + wrapped_probability_fn = functools.partial( + _monotonic_probability_fn, sigmoid_noise=sigmoid_noise, mode=mode, + seed=sigmoid_noise_seed) + query_layer = kwargs.pop("query_layer", None) + if not query_layer: + query_layer = layers.Dense( + units, name="query_layer", use_bias=False, dtype=dtype) + memory_layer = kwargs.pop("memory_layer", None) + if not memory_layer: + memory_layer = layers.Dense( + units, name="memory_layer", use_bias=False, dtype=dtype) + self.units = units + self.normalize = normalize + self.sigmoid_noise = sigmoid_noise + self.sigmoid_noise_seed = sigmoid_noise_seed + self.score_bias_init = score_bias_init + self.mode = mode + self.kernel_initializer = initializers.get(kernel_initializer) + self.attention_v = None + self.attention_score_bias = None + self.attention_g = None + self.attention_b = None + super(BahdanauMonotonicAttentionV2, self).__init__( + memory=memory, + memory_sequence_length=memory_sequence_length, + query_layer=query_layer, + memory_layer=memory_layer, + probability_fn=wrapped_probability_fn, + name=name, + dtype=dtype, + **kwargs) + + def build(self, input_shape): + super(BahdanauMonotonicAttentionV2, self).build(input_shape) + if self.attention_v is None: + self.attention_v = self.add_weight( + "attention_v", [self.units], dtype=self.dtype, + initializer=self.kernel_initializer) + if self.attention_score_bias is None: + self.attention_score_bias = self.add_weight( + "attention_score_bias", shape=(), dtype=self.dtype, + initializer=init_ops.constant_initializer( + self.score_bias_init, dtype=self.dtype)) + if self.normalize and self.attention_g is None and self.attention_b is None: + self.attention_g = self.add_weight( + "attention_g", dtype=self.dtype, + initializer=init_ops.constant_initializer( + math.sqrt((1. / self.units))), + shape=()) + self.attention_b = self.add_weight( + "attention_b", [self.units], dtype=self.dtype, + initializer=init_ops.zeros_initializer()) + self.built = True + + def _calculate_attention(self, query, state): + """Score the query based on the keys and values. + + Args: + query: Tensor of dtype matching `self.values` and shape + `[batch_size, query_depth]`. + state: Tensor of dtype matching `self.values` and shape + `[batch_size, alignments_size]` + (`alignments_size` is memory's `max_time`). + + Returns: + alignments: Tensor of dtype matching `self.values` and shape + `[batch_size, alignments_size]` (`alignments_size` is memory's + `max_time`). + """ + processed_query = self.query_layer(query) if self.query_layer else query + score = _bahdanau_score(processed_query, self.keys, self.attention_v, + attention_g=self.attention_g, + attention_b=self.attention_b) + score += self.attention_score_bias + alignments = self.probability_fn(score, state) + next_state = alignments + return alignments, next_state + + def get_config(self): + config = { + "units": self.units, + "normalize": self.normalize, + "sigmoid_noise": self.sigmoid_noise, + "sigmoid_noise_seed": self.sigmoid_noise_seed, + "score_bias_init": self.score_bias_init, + "mode": self.mode, + "kernel_initializer": initializers.serialize(self.kernel_initializer), + } + base_config = super(BahdanauMonotonicAttentionV2, self).get_config() + return dict(list(base_config.items()) + list(config.items())) + + @classmethod + def from_config(cls, config, custom_objects=None): + config = _BaseAttentionMechanismV2.deserialize_inner_layer_from_config( + config, custom_objects=custom_objects) + return cls(**config) + + class LuongMonotonicAttention(_BaseMonotonicAttentionMechanism): """Monotonic attention mechanism with Luong-style energy function. @@ -960,7 +1724,12 @@ class LuongMonotonicAttention(_BaseMonotonicAttentionMechanism): """ with variable_scope.variable_scope(None, "luong_monotonic_attention", [query]): - score = _luong_score(query, self._keys, self._scale) + attention_g = None + if self._scale: + attention_g = variable_scope.get_variable( + "attention_g", dtype=query.dtype, + initializer=init_ops.ones_initializer, shape=()) + score = _luong_score(query, self._keys, attention_g) score_bias = variable_scope.get_variable( "attention_score_bias", dtype=query.dtype, initializer=self._score_bias_init) @@ -970,6 +1739,139 @@ class LuongMonotonicAttention(_BaseMonotonicAttentionMechanism): return alignments, next_state +class LuongMonotonicAttentionV2(_BaseMonotonicAttentionMechanismV2): + """Monotonic attention mechanism with Luong-style energy function. + + This type of attention enforces a monotonic constraint on the attention + distributions; that is once the model attends to a given point in the memory + it can't attend to any prior points at subsequence output timesteps. It + achieves this by using the _monotonic_probability_fn instead of softmax to + construct its attention distributions. Otherwise, it is equivalent to + LuongAttention. This approach is proposed in + + [Colin Raffel, Minh-Thang Luong, Peter J. Liu, Ron J. Weiss, Douglas Eck, + "Online and Linear-Time Attention by Enforcing Monotonic Alignments." + ICML 2017.](https://arxiv.org/abs/1704.00784) + """ + + def __init__(self, + units, + memory, + memory_sequence_length=None, + scale=False, + sigmoid_noise=0., + sigmoid_noise_seed=None, + score_bias_init=0., + mode="parallel", + dtype=None, + name="LuongMonotonicAttention", + **kwargs): + """Construct the Attention mechanism. + + Args: + units: The depth of the query mechanism. + memory: The memory to query; usually the output of an RNN encoder. This + tensor should be shaped `[batch_size, max_time, ...]`. + memory_sequence_length: (optional): Sequence lengths for the batch entries + in memory. If provided, the memory tensor rows are masked with zeros + for values past the respective sequence lengths. + scale: Python boolean. Whether to scale the energy term. + sigmoid_noise: Standard deviation of pre-sigmoid noise. See the docstring + for `_monotonic_probability_fn` for more information. + sigmoid_noise_seed: (optional) Random seed for pre-sigmoid noise. + score_bias_init: Initial value for score bias scalar. It's recommended to + initialize this to a negative value when the length of the memory is + large. + mode: How to compute the attention distribution. Must be one of + 'recursive', 'parallel', or 'hard'. See the docstring for + `tf.contrib.seq2seq.monotonic_attention` for more information. + dtype: The data type for the query and memory layers of the attention + mechanism. + name: Name to use when creating ops. + **kwargs: Dictionary that contains other common arguments for layer + creation. + """ + # Set up the monotonic probability fn with supplied parameters + if dtype is None: + dtype = dtypes.float32 + wrapped_probability_fn = functools.partial( + _monotonic_probability_fn, sigmoid_noise=sigmoid_noise, mode=mode, + seed=sigmoid_noise_seed) + memory_layer = kwargs.pop("memory_layer", None) + if not memory_layer: + memory_layer = layers.Dense( + units, name="memory_layer", use_bias=False, dtype=dtype) + self.units = units + self.scale = scale + self.sigmoid_noise = sigmoid_noise + self.sigmoid_noise_seed = sigmoid_noise_seed + self.score_bias_init = score_bias_init + self.mode = mode + self.attention_g = None + self.attention_score_bias = None + super(LuongMonotonicAttentionV2, self).__init__( + memory=memory, + memory_sequence_length=memory_sequence_length, + query_layer=None, + memory_layer=memory_layer, + probability_fn=wrapped_probability_fn, + name=name, + dtype=dtype, + **kwargs) + + def build(self, input_shape): + super(LuongMonotonicAttentionV2, self).build(input_shape) + if self.scale and self.attention_g is None: + self.attention_g = self.add_weight( + "attention_g", initializer=init_ops.ones_initializer, shape=()) + if self.attention_score_bias is None: + self.attention_score_bias = self.add_weight( + "attention_score_bias", shape=(), + initializer=init_ops.constant_initializer( + self.score_bias_init, dtype=self.dtype)) + self.built = True + + def _calculate_attention(self, query, state): + """Score the query based on the keys and values. + + Args: + query: Tensor of dtype matching `self.values` and shape + `[batch_size, query_depth]`. + state: Tensor of dtype matching `self.values` and shape + `[batch_size, alignments_size]` + (`alignments_size` is memory's `max_time`). + + Returns: + alignments: Tensor of dtype matching `self.values` and shape + `[batch_size, alignments_size]` (`alignments_size` is memory's + `max_time`). + next_state: Same as alignments + """ + score = _luong_score(query, self.keys, self.attention_g) + score += self.attention_score_bias + alignments = self.probability_fn(score, state) + next_state = alignments + return alignments, next_state + + def get_config(self): + config = { + "units": self.units, + "scale": self.scale, + "sigmoid_noise": self.sigmoid_noise, + "sigmoid_noise_seed": self.sigmoid_noise_seed, + "score_bias_init": self.score_bias_init, + "mode": self.mode, + } + base_config = super(LuongMonotonicAttentionV2, self).get_config() + return dict(list(base_config.items()) + list(config.items())) + + @classmethod + def from_config(cls, config, custom_objects=None): + config = _BaseAttentionMechanismV2.deserialize_inner_layer_from_config( + config, custom_objects=custom_objects) + return cls(**config) + + class AttentionWrapperState( collections.namedtuple("AttentionWrapperState", ("cell_state", "attention", "time", "alignments", @@ -1026,6 +1928,82 @@ class AttentionWrapperState( super(AttentionWrapperState, self)._replace(**kwargs)) +def _prepare_memory(memory, memory_sequence_length=None, memory_mask=None, + check_inner_dims_defined=True): + """Convert to tensor and possibly mask `memory`. + + Args: + memory: `Tensor`, shaped `[batch_size, max_time, ...]`. + memory_sequence_length: `int32` `Tensor`, shaped `[batch_size]`. + memory_mask: `boolean` tensor with shape [batch_size, max_time]. The memory + should be skipped when the corresponding mask is False. + check_inner_dims_defined: Python boolean. If `True`, the `memory` + argument's shape is checked to ensure all but the two outermost + dimensions are fully defined. + + Returns: + A (possibly masked), checked, new `memory`. + + Raises: + ValueError: If `check_inner_dims_defined` is `True` and not + `memory.shape[2:].is_fully_defined()`. + """ + memory = nest.map_structure( + lambda m: ops.convert_to_tensor(m, name="memory"), memory) + if memory_sequence_length is not None and memory_mask is not None: + raise ValueError("memory_sequence_length and memory_mask can't be provided " + "at same time.") + if memory_sequence_length is not None: + memory_sequence_length = ops.convert_to_tensor( + memory_sequence_length, name="memory_sequence_length") + if check_inner_dims_defined: + def _check_dims(m): + if not m.get_shape()[2:].is_fully_defined(): + raise ValueError("Expected memory %s to have fully defined inner dims, " + "but saw shape: %s" % (m.name, m.get_shape())) + nest.map_structure(_check_dims, memory) + if memory_sequence_length is None and memory_mask is None: + return memory + elif memory_sequence_length is not None: + seq_len_mask = array_ops.sequence_mask( + memory_sequence_length, + maxlen=array_ops.shape(nest.flatten(memory)[0])[1], + dtype=nest.flatten(memory)[0].dtype) + else: + # For memory_mask is not None + seq_len_mask = math_ops.cast( + memory_mask, dtype=nest.flatten(memory)[0].dtype) + def _maybe_mask(m, seq_len_mask): + """Mask the memory based on the memory mask.""" + rank = m.get_shape().ndims + rank = rank if rank is not None else array_ops.rank(m) + extra_ones = array_ops.ones(rank - 2, dtype=dtypes.int32) + seq_len_mask = array_ops.reshape( + seq_len_mask, + array_ops.concat((array_ops.shape(seq_len_mask), extra_ones), 0)) + return m * seq_len_mask + + return nest.map_structure(lambda m: _maybe_mask(m, seq_len_mask), memory) + + +def _maybe_mask_score(score, memory_sequence_length=None, memory_mask=None, + score_mask_value=None): + """Mask the attention score based on the masks.""" + if memory_sequence_length is None and memory_mask is None: + return score + if memory_sequence_length is not None and memory_mask is not None: + raise ValueError("memory_sequence_length and memory_mask can't be provided " + "at same time.") + if memory_sequence_length is not None: + message = "All values in memory_sequence_length must greater than zero." + with ops.control_dependencies( + [check_ops.assert_positive(memory_sequence_length, message=message)]): + memory_mask = array_ops.sequence_mask( + memory_sequence_length, maxlen=array_ops.shape(score)[1]) + score_mask_values = score_mask_value * array_ops.ones_like(score) + return array_ops.where(memory_mask, score, score_mask_values) + + def hardmax(logits, name=None): """Returns batched one-hot vectors. @@ -1050,8 +2028,14 @@ def hardmax(logits, name=None): def _compute_attention(attention_mechanism, cell_output, attention_state, attention_layer): """Computes the attention and alignments for a given attention_mechanism.""" - alignments, next_attention_state = attention_mechanism( - cell_output, state=attention_state) + if isinstance(attention_mechanism, _BaseAttentionMechanismV2): + alignments, next_attention_state = attention_mechanism( + [cell_output, attention_state]) + else: + # For other class, assume they are following _BaseAttentionMechanism, which + # takes query and state as separate parameter. + alignments, next_attention_state = attention_mechanism( + cell_output, state=attention_state) # Reshape from [batch_size, memory_time] to [batch_size, 1, memory_time] expanded_alignments = array_ops.expand_dims(alignments, 1) diff --git a/tensorflow/contrib/seq2seq/python/ops/basic_decoder.py b/tensorflow/contrib/seq2seq/python/ops/basic_decoder.py index 7eb95e5a70de985dca0d4b565ba03bdf454b6161..16dfa7ed8268d761dee49ec0146efabcaaef1393 100644 --- a/tensorflow/contrib/seq2seq/python/ops/basic_decoder.py +++ b/tensorflow/contrib/seq2seq/python/ops/basic_decoder.py @@ -23,8 +23,10 @@ import collections from tensorflow.contrib.seq2seq.python.ops import decoder from tensorflow.contrib.seq2seq.python.ops import helper as helper_py +from tensorflow.contrib.seq2seq.python.ops import sampler as sampler_py from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape +from tensorflow.python.keras import layers from tensorflow.python.layers import base as layers_base from tensorflow.python.ops import rnn_cell_impl from tensorflow.python.util import nest @@ -146,3 +148,102 @@ class BasicDecoder(decoder.Decoder): sample_ids=sample_ids) outputs = BasicDecoderOutput(cell_outputs, sample_ids) return (outputs, next_state, next_inputs, finished) + + +class BasicDecoderV2(decoder.BaseDecoder): + """Basic sampling decoder.""" + + def __init__(self, cell, sampler, output_layer=None, **kwargs): + """Initialize BasicDecoder. + + Args: + cell: An `RNNCell` instance. + sampler: A `Sampler` instance. + output_layer: (Optional) An instance of `tf.layers.Layer`, i.e., + `tf.layers.Dense`. Optional layer to apply to the RNN output prior to + storing the result or sampling. + **kwargs: Other keyward arguments for layer creation. + + Raises: + TypeError: if `cell`, `helper` or `output_layer` have an incorrect type. + """ + rnn_cell_impl.assert_like_rnncell("cell", cell) + if not isinstance(sampler, sampler_py.Sampler): + raise TypeError("sampler must be a Sampler, received: %s" % (sampler,)) + if (output_layer is not None and + not isinstance(output_layer, layers.Layer)): + raise TypeError( + "output_layer must be a Layer, received: %s" % (output_layer,)) + self.cell = cell + self.sampler = sampler + self.output_layer = output_layer + super(BasicDecoderV2, self).__init__(**kwargs) + + def initialize(self, inputs, initial_state=None, **kwargs): + """Initialize the decoder.""" + # Assume the dtype of the cell is the output_size structure + # containing the input_state's first component's dtype. + self._cell_dtype = nest.flatten(initial_state)[0].dtype + return self.sampler.initialize(inputs, **kwargs) + (initial_state,) + + @property + def batch_size(self): + return self.sampler.batch_size + + def _rnn_output_size(self): + size = tensor_shape.TensorShape(self.cell.output_size) + if self.output_layer is None: + return size + else: + # To use layer's compute_output_shape, we need to convert the + # RNNCell's output_size entries into shapes with an unknown + # batch size. We then pass this through the layer's + # compute_output_shape and read off all but the first (batch) + # dimensions to get the output size of the rnn with the layer + # applied to the top. + output_shape_with_unknown_batch = nest.map_structure( + lambda s: tensor_shape.TensorShape([None]).concatenate(s), size) + layer_output_shape = self.output_layer.compute_output_shape( + output_shape_with_unknown_batch) + return nest.map_structure(lambda s: s[1:], layer_output_shape) + + @property + def output_size(self): + # Return the cell output and the id + return BasicDecoderOutput( + rnn_output=self._rnn_output_size(), + sample_id=self.sampler.sample_ids_shape) + + @property + def output_dtype(self): + # Assume the dtype of the cell is the output_size structure + # containing the input_state's first component's dtype. + # Return that structure and the sample_ids_dtype from the helper. + dtype = self._cell_dtype + return BasicDecoderOutput( + nest.map_structure(lambda _: dtype, self._rnn_output_size()), + self.sampler.sample_ids_dtype) + + def step(self, time, inputs, state): + """Perform a decoding step. + + Args: + time: scalar `int32` tensor. + inputs: A (structure of) input tensors. + state: A (structure of) state tensors and TensorArrays. + + Returns: + `(outputs, next_state, next_inputs, finished)`. + """ + cell_outputs, cell_state = self.cell(inputs, state) + if self.output_layer is not None: + cell_outputs = self.output_layer(cell_outputs) + sample_ids = self.sampler.sample( + time=time, outputs=cell_outputs, state=cell_state) + (finished, next_inputs, next_state) = self.sampler.next_inputs( + time=time, + outputs=cell_outputs, + state=cell_state, + sample_ids=sample_ids) + outputs = BasicDecoderOutput(cell_outputs, sample_ids) + return (outputs, next_state, next_inputs, finished) diff --git a/tensorflow/contrib/seq2seq/python/ops/decoder.py b/tensorflow/contrib/seq2seq/python/ops/decoder.py index f58268eff525a4b592c79acb32207e1a3f62bdc7..33f7bac8159401175ce57c0463fff1398c1dd9bb 100644 --- a/tensorflow/contrib/seq2seq/python/ops/decoder.py +++ b/tensorflow/contrib/seq2seq/python/ops/decoder.py @@ -27,6 +27,7 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.framework import tensor_util +from tensorflow.python.keras import layers from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import control_flow_util @@ -135,6 +136,127 @@ class Decoder(object): return False +class BaseDecoder(layers.Layer): + """An RNN Decoder that is based on a Keras layer. + + Concepts used by this interface: + - `inputs`: (structure of) tensors and TensorArrays that is passed as input to + the RNNCell composing the decoder, at each time step. + - `state`: (structure of) tensors and TensorArrays that is passed to the + RNNCell instance as the state. + - `memory`: (sturecute of) tensors that is usually the full output of the + encoder, which will be used for the attention wrapper for the RNNCell. + - `finished`: boolean tensor telling whether each sequence in the batch is + finished. + - `outputs`: Instance of BasicDecoderOutput. Result of the decoding, at each + time step. + """ + + def __init__(self, + output_time_major=False, + impute_finished=False, + maximum_iterations=None, + parallel_iterations=32, + swap_memory=False, + **kwargs): + self.output_time_major = output_time_major + self.impute_finished = impute_finished + self.maximum_iterations = maximum_iterations + self.parallel_iterations = parallel_iterations + self.swap_memory = swap_memory + super(BaseDecoder, self).__init__(**kwargs) + + def call(self, inputs, initial_state=None, **kwargs): + init_kwargs = kwargs + init_kwargs["initial_state"] = initial_state + return dynamic_decode(self, + output_time_major=self.output_time_major, + impute_finished=self.impute_finished, + maximum_iterations=self.maximum_iterations, + parallel_iterations=self.parallel_iterations, + swap_memory=self.swap_memory, + decoder_init_input=inputs, + decoder_init_kwargs=init_kwargs) + + @property + def batch_size(self): + """The batch size of input values.""" + raise NotImplementedError + + @property + def output_size(self): + """A (possibly nested tuple of...) integer[s] or `TensorShape` object[s].""" + raise NotImplementedError + + @property + def output_dtype(self): + """A (possibly nested tuple of...) dtype[s].""" + raise NotImplementedError + + def initialize(self, inputs, initial_state=None, **kwargs): + """Called before any decoding iterations. + + This methods must compute initial input values and initial state. + + Args: + inputs: (structure of) tensors that contains the input for the decoder. In + the normal case, its a tensor with shape [batch, timestep, embedding]. + initial_state: (structure of) tensors that contains the initial state for + the RNNCell. + **kwargs: Other arguments that are passed in from layer.call() method. It + could contains item like input sequence_length, or masking for input. + + Returns: + `(finished, initial_inputs, initial_state)`: initial values of + 'finished' flags, inputs and state. + """ + raise NotImplementedError + + def step(self, time, inputs, state): + """Called per step of decoding (but only once for dynamic decoding). + + Args: + time: Scalar `int32` tensor. Current step number. + inputs: RNNCell input (possibly nested tuple of) tensor[s] for this time + step. + state: RNNCell state (possibly nested tuple of) tensor[s] from previous + time step. + + Returns: + `(outputs, next_state, next_inputs, finished)`: `outputs` is an object + containing the decoder output, `next_state` is a (structure of) state + tensors and TensorArrays, `next_inputs` is the tensor that should be used + as input for the next step, `finished` is a boolean tensor telling whether + the sequence is complete, for each sequence in the batch. + """ + raise NotImplementedError + + def finalize(self, outputs, final_state, sequence_lengths): + raise NotImplementedError + + @property + def tracks_own_finished(self): + """Describes whether the Decoder keeps track of finished states. + + Most decoders will emit a true/false `finished` value independently + at each time step. In this case, the `dynamic_decode` function keeps track + of which batch entries are already finished, and performs a logical OR to + insert new batches to the finished set. + + Some decoders, however, shuffle batches / beams between time steps and + `dynamic_decode` will mix up the finished state across these entries because + it does not track the reshuffle across time steps. In this case, it is + up to the decoder to declare that it will keep track of its own finished + state by setting this property to `True`. + + Returns: + Python bool. + """ + return False + + # TODO(scottzhu): Add build/get_config/from_config and other layer methods. + + def _create_zero_outputs(size, dtype, batch_size): """Create a zero outputs Tensor structure.""" def _create(s, d): @@ -149,7 +271,8 @@ def dynamic_decode(decoder, maximum_iterations=None, parallel_iterations=32, swap_memory=False, - scope=None): + scope=None, + **kwargs): """Perform dynamic decoding with `decoder`. Calls initialize() once and step() repeatedly on the Decoder object. @@ -171,6 +294,9 @@ def dynamic_decode(decoder, parallel_iterations: Argument passed to `tf.while_loop`. swap_memory: Argument passed to `tf.while_loop`. scope: Optional variable scope to use. + **kwargs: dict, other keyword arguments for dynamic_decode. It might contain + arguments for `BaseDecoder` to initialize, which takes all tensor inputs + during call(). Returns: `(final_outputs, final_state, final_sequence_lengths)`. @@ -179,7 +305,7 @@ def dynamic_decode(decoder, TypeError: if `decoder` is not an instance of `Decoder`. ValueError: if `maximum_iterations` is provided but is not a scalar. """ - if not isinstance(decoder, Decoder): + if not isinstance(decoder, (Decoder, BaseDecoder)): raise TypeError("Expected decoder to be type Decoder, but saw: %s" % type(decoder)) @@ -204,7 +330,14 @@ def dynamic_decode(decoder, if maximum_iterations.get_shape().ndims != 0: raise ValueError("maximum_iterations must be a scalar") - initial_finished, initial_inputs, initial_state = decoder.initialize() + if isinstance(decoder, Decoder): + initial_finished, initial_inputs, initial_state = decoder.initialize() + else: + # For BaseDecoder that takes tensor inputs during call. + decoder_init_input = kwargs.pop("decoder_init_input", None) + decoder_init_kwargs = kwargs.pop("decoder_init_kwargs", {}) + initial_finished, initial_inputs, initial_state = decoder.initialize( + decoder_init_input, **decoder_init_kwargs) zero_outputs = _create_zero_outputs(decoder.output_size, decoder.output_dtype, @@ -222,7 +355,7 @@ def dynamic_decode(decoder, def _shape(batch_size, from_shape): if (not isinstance(from_shape, tensor_shape.TensorShape) or from_shape.ndims == 0): - return tensor_shape.TensorShape(None) + return None else: batch_size = tensor_util.constant_value( ops.convert_to_tensor( diff --git a/tensorflow/contrib/seq2seq/python/ops/helper.py b/tensorflow/contrib/seq2seq/python/ops/helper.py index 3245cc5e72154289ea3ba000b9a30586a7ad03a9..033c2eb0801d5a51ee937f5e960faa91a6f1ae54 100644 --- a/tensorflow/contrib/seq2seq/python/ops/helper.py +++ b/tensorflow/contrib/seq2seq/python/ops/helper.py @@ -32,9 +32,8 @@ from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import embedding_ops from tensorflow.python.ops import gen_array_ops from tensorflow.python.ops import math_ops +from tensorflow.python.ops import random_ops from tensorflow.python.ops import tensor_array_ops -from tensorflow.python.ops.distributions import bernoulli -from tensorflow.python.ops.distributions import categorical from tensorflow.python.util import nest __all__ = [ @@ -51,6 +50,68 @@ __all__ = [ _transpose_batch_time = decoder._transpose_batch_time # pylint: disable=protected-access +# The following sample functions (_call_sampler, bernoulli_sample, +# categorical_sample) mimic TensorFlow Probability distribution semantics. + + +def _call_sampler(sample_n_fn, sample_shape, name=None): + """Reshapes vector of samples.""" + with ops.name_scope(name, "call_sampler", values=[sample_shape]): + sample_shape = ops.convert_to_tensor( + sample_shape, dtype=dtypes.int32, name="sample_shape") + # Ensure sample_shape is a vector (vs just a scalar). + pad = math_ops.cast(math_ops.equal(array_ops.rank(sample_shape), 0), + dtypes.int32) + sample_shape = array_ops.reshape( + sample_shape, + array_ops.pad(array_ops.shape(sample_shape), + paddings=[[pad, 0]], + constant_values=1)) + samples = sample_n_fn(math_ops.reduce_prod(sample_shape)) + batch_event_shape = array_ops.shape(samples)[1:] + final_shape = array_ops.concat([sample_shape, batch_event_shape], 0) + return array_ops.reshape(samples, final_shape) + + +def bernoulli_sample(probs=None, logits=None, dtype=dtypes.int32, + sample_shape=(), seed=None): + """Samples from Bernoulli distribution.""" + if probs is None: + probs = math_ops.sigmoid(logits, name="probs") + else: + probs = ops.convert_to_tensor(probs, name="probs") + batch_shape_tensor = array_ops.shape(probs) + def _sample_n(n): + """Sample vector of Bernoullis.""" + new_shape = array_ops.concat([[n], batch_shape_tensor], 0) + uniform = random_ops.random_uniform( + new_shape, seed=seed, dtype=probs.dtype) + return math_ops.cast(math_ops.less(uniform, probs), dtype) + return _call_sampler(_sample_n, sample_shape) + + +def categorical_sample(logits, dtype=dtypes.int32, + sample_shape=(), seed=None): + """Samples from categorical distribution.""" + logits = ops.convert_to_tensor(logits, name="logits") + event_size = array_ops.shape(logits)[-1] + batch_shape_tensor = array_ops.shape(logits)[:-1] + def _sample_n(n): + """Sample vector of categoricals.""" + if logits.shape.ndims == 2: + logits_2d = logits + else: + logits_2d = array_ops.reshape(logits, [-1, event_size]) + sample_dtype = dtypes.int64 if logits.dtype.size > 4 else dtypes.int32 + draws = random_ops.multinomial( + logits_2d, n, seed=seed, output_dtype=sample_dtype) + draws = array_ops.reshape( + array_ops.transpose(draws), + array_ops.concat([[n], batch_shape_tensor], 0)) + return math_ops.cast(draws, dtype) + return _call_sampler(_sample_n, sample_shape) + + def _unstack_ta(inp): return tensor_array_ops.TensorArray( dtype=inp.dtype, size=array_ops.shape(inp)[0], @@ -307,14 +368,14 @@ class ScheduledEmbeddingTrainingHelper(TrainingHelper): with ops.name_scope(name, "ScheduledEmbeddingTrainingHelperSample", [time, outputs, state]): # Return -1s where we did not sample, and sample_ids elsewhere - select_sampler = bernoulli.Bernoulli( - probs=self._sampling_probability, dtype=dtypes.bool) - select_sample = select_sampler.sample( - sample_shape=self.batch_size, seed=self._scheduling_seed) - sample_id_sampler = categorical.Categorical(logits=outputs) + select_sample = bernoulli_sample( + probs=self._sampling_probability, + dtype=dtypes.bool, + sample_shape=self.batch_size, + seed=self._scheduling_seed) return array_ops.where( select_sample, - sample_id_sampler.sample(seed=self._seed), + categorical_sample(logits=outputs, seed=self._seed), gen_array_ops.fill([self.batch_size], -1)) def next_inputs(self, time, outputs, state, sample_ids, name=None): @@ -425,8 +486,10 @@ class ScheduledOutputTrainingHelper(TrainingHelper): def sample(self, time, outputs, state, name=None): with ops.name_scope(name, "ScheduledOutputTrainingHelperSample", [time, outputs, state]): - sampler = bernoulli.Bernoulli(probs=self._sampling_probability) - return sampler.sample(sample_shape=self.batch_size, seed=self._seed) + return bernoulli_sample( + probs=self._sampling_probability, + sample_shape=self.batch_size, + seed=self._seed) def next_inputs(self, time, outputs, state, sample_ids, name=None): with ops.name_scope(name, "ScheduledOutputTrainingHelperNextInputs", @@ -610,8 +673,7 @@ class SampleEmbeddingHelper(GreedyEmbeddingHelper): else: logits = outputs / self._softmax_temperature - sample_id_sampler = categorical.Categorical(logits=logits) - sample_ids = sample_id_sampler.sample(seed=self._seed) + sample_ids = categorical_sample(logits=logits, seed=self._seed) return sample_ids diff --git a/tensorflow/contrib/seq2seq/python/ops/loss.py b/tensorflow/contrib/seq2seq/python/ops/loss.py index 39a6d2f58b140706a94d83273d3327edd1891368..0fbfd6187030f14ac105a18b3e09b7a42d4de32a 100644 --- a/tensorflow/contrib/seq2seq/python/ops/loss.py +++ b/tensorflow/contrib/seq2seq/python/ops/loss.py @@ -20,11 +20,12 @@ from __future__ import division from __future__ import print_function from tensorflow.python.framework import ops +from tensorflow.python.keras.losses import Loss from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_ops -__all__ = ["sequence_loss"] +__all__ = ["sequence_loss", "SequenceLoss"] def sequence_loss(logits, @@ -32,16 +33,26 @@ def sequence_loss(logits, weights, average_across_timesteps=True, average_across_batch=True, + sum_over_timesteps=False, + sum_over_batch=False, softmax_loss_function=None, name=None): """Weighted cross-entropy loss for a sequence of logits. - Depending on the values of `average_across_timesteps` and - `average_across_batch`, the return Tensor will have rank 0, 1, or 2 as these - arguments reduce the cross-entropy at each target, which has shape - `[batch_size, sequence_length]`, over their respective dimensions. For - example, if `average_across_timesteps` is `True` and `average_across_batch` - is `False`, then the return Tensor will have shape `[batch_size]`. + Depending on the values of `average_across_timesteps` / `sum_over_timesteps` + and `average_across_batch` / `sum_over_batch`, the return Tensor will have + rank 0, 1, or 2 as these arguments reduce the cross-entropy at each target, + which has shape `[batch_size, sequence_length]`, over their respective + dimensions. For example, if `average_across_timesteps` is `True` and + `average_across_batch` is `False`, then the return Tensor will have shape + `[batch_size]`. + + Note that `average_across_timesteps` and `sum_over_timesteps` cannot be True + at same time. Same for `average_across_batch` and `sum_over_batch`. + + The recommended loss reduction in tf 2.0 has been changed to sum_over, instead + of weighted average. User are recommend to use `sum_over_timesteps` and + `sum_over_batch` for reduction. Args: logits: A Tensor of shape @@ -58,6 +69,12 @@ def sequence_loss(logits, dimension and divide the cost by the total label weight across timesteps. average_across_batch: If set, sum the cost across the batch dimension and divide the returned cost by the batch size. + sum_over_timesteps: If set, sum the cost across the sequence dimension and + divide the size of the sequence. Note that any element with 0 weights will + be excluded from size calculation. + sum_over_batch: if set, sum the cost across the batch dimension and divide + the total cost by the batch size. Not that any element with 0 weights will + be excluded from size calculation. softmax_loss_function: Function (labels, logits) -> loss-batch to be used instead of the standard softmax (the default if this is None). **Note that to avoid confusion, it is required for the function to accept @@ -78,11 +95,15 @@ def sequence_loss(logits, raise ValueError("Logits must be a " "[batch_size x sequence_length x logits] tensor") if len(targets.get_shape()) != 2: - raise ValueError("Targets must be a [batch_size x sequence_length] " - "tensor") + raise ValueError("Targets must be a [batch_size x sequence_length] tensor") if len(weights.get_shape()) != 2: - raise ValueError("Weights must be a [batch_size x sequence_length] " - "tensor") + raise ValueError("Weights must be a [batch_size x sequence_length] tensor") + if average_across_timesteps and sum_over_timesteps: + raise ValueError("average_across_timesteps and sum_over_timesteps cannot " + "be set to True at same time.") + if average_across_batch and sum_over_batch: + raise ValueError("average_across_batch and sum_over_batch cannot be set " + "to True at same time.") with ops.name_scope(name, "sequence_loss", [logits, targets, weights]): num_classes = array_ops.shape(logits)[2] logits_flat = array_ops.reshape(logits, [-1, num_classes]) @@ -96,20 +117,56 @@ def sequence_loss(logits, if average_across_timesteps and average_across_batch: crossent = math_ops.reduce_sum(crossent) total_size = math_ops.reduce_sum(weights) - total_size += 1e-12 # to avoid division by 0 for all-0 weights - crossent /= total_size + crossent = math_ops.div_no_nan(crossent, total_size) + elif sum_over_timesteps and sum_over_batch: + crossent = math_ops.reduce_sum(crossent) + total_count = math_ops.cast(math_ops.count_nonzero(weights), + crossent.dtype) + crossent = math_ops.div_no_nan(crossent, total_count) else: - batch_size = array_ops.shape(logits)[0] - sequence_length = array_ops.shape(logits)[1] - crossent = array_ops.reshape(crossent, [batch_size, sequence_length]) - if average_across_timesteps and not average_across_batch: - crossent = math_ops.reduce_sum(crossent, axis=[1]) - total_size = math_ops.reduce_sum(weights, axis=[1]) - total_size += 1e-12 # to avoid division by 0 for all-0 weights - crossent /= total_size - if not average_across_timesteps and average_across_batch: - crossent = math_ops.reduce_sum(crossent, axis=[0]) - total_size = math_ops.reduce_sum(weights, axis=[0]) - total_size += 1e-12 # to avoid division by 0 for all-0 weights - crossent /= total_size + crossent = array_ops.reshape(crossent, array_ops.shape(logits)[0:2]) + if average_across_timesteps or average_across_batch: + reduce_axis = [0] if average_across_batch else [1] + crossent = math_ops.reduce_sum(crossent, axis=reduce_axis) + total_size = math_ops.reduce_sum(weights, axis=reduce_axis) + crossent = math_ops.div_no_nan(crossent, total_size) + elif sum_over_timesteps or sum_over_batch: + reduce_axis = [0] if sum_over_batch else [1] + crossent = math_ops.reduce_sum(crossent, axis=reduce_axis) + total_count = math_ops.cast( + math_ops.count_nonzero(weights, axis=reduce_axis), + dtype=crossent.dtype) + crossent = math_ops.div_no_nan(crossent, total_count) return crossent + + +class SequenceLoss(Loss): + """Weighted cross-entropy loss for a sequence of logits.""" + + def __init__(self, + average_across_timesteps=False, + average_across_batch=False, + sum_over_timesteps=True, + sum_over_batch=True, + softmax_loss_function=None, + name=None): + super(SequenceLoss, self).__init__(name=name) + self.average_across_timesteps = average_across_timesteps + self.average_across_batch = average_across_batch + self.sum_over_timesteps = sum_over_timesteps + self.sum_over_batch = sum_over_batch + self.softmax_loss_function = softmax_loss_function + + def __call__(self, y_true, y_pred, sample_weight=None): + """Override the parent __call__ to have a customized reduce behavior.""" + return sequence_loss(y_pred, y_true, sample_weight, + average_across_timesteps=self.average_across_timesteps, + average_across_batch=self.average_across_batch, + sum_over_timesteps=self.sum_over_timesteps, + sum_over_batch=self.sum_over_batch, + softmax_loss_function=self.softmax_loss_function, + name=self.name) + + def call(self, y_true, y_pred): + # Skip this method since the __call__ contains real implementation. + pass diff --git a/tensorflow/contrib/seq2seq/python/ops/sampler.py b/tensorflow/contrib/seq2seq/python/ops/sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..9e3e48b3bc61c0ff94ae0a1794767c7ff6914969 --- /dev/null +++ b/tensorflow/contrib/seq2seq/python/ops/sampler.py @@ -0,0 +1,765 @@ +# 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. +# ============================================================================== +"""A library of sampler for use with SamplingDecoders.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import abc + +import six + +from tensorflow.contrib.seq2seq.python.ops import decoder +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.ops import control_flow_ops +from tensorflow.python.ops import embedding_ops +from tensorflow.python.ops import gen_array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import random_ops +from tensorflow.python.ops import tensor_array_ops +from tensorflow.python.util import nest + +__all__ = [ + "Sampler", + "TrainingSampler", + "GreedyEmbeddingSampler", + "SampleEmbeddingSampler", + "CustomSampler", + "ScheduledEmbeddingTrainingSampler", + "ScheduledOutputTrainingSampler", + "InferenceSampler", +] + +_transpose_batch_time = decoder._transpose_batch_time # pylint: disable=protected-access + + +@six.add_metaclass(abc.ABCMeta) +class Sampler(object): + """Interface for implementing sampling in seq2seq decoders. + + Sampler instances are used by `BasicDecoder`. The normal usage of a sampler is + like below: + sampler = Sampler(init_args) + (initial_finished, initial_inputs) = sampler.initialize(input_tensors) + for time_step in range(time): + cell_output, cell_state = cell.call(cell_input, previous_state) + sample_ids = sampler.sample(time_step, cell_output, cell_state) + (finished, next_inputs, next_state) = sampler.next_inputs( + time_step,cell_output, cell_state) + + Note that all the tensor input should not be feed to Sampler as __init__() + parameters, instead, they should be feed by decoders via initialize(). + """ + + @abc.abstractmethod + def initialize(self, inputs, **kwargs): + """initialize the sampler with the input tensors. + + This method suppose to be only invoke once before the calling other methods + of the Sampler. + + Args: + inputs: A (structure of) input tensors, it could be a nested tuple or a + single tensor. + **kwargs: Other kwargs for initialization. It could contain tensors like + mask for inputs, or non tensor parameter. + + Returns: + `(initial_finished, initial_inputs)`. + """ + pass + + @abc.abstractmethod + def sample(self, time, outputs, state): + """Returns `sample_ids`.""" + pass + + @abc.abstractmethod + def next_inputs(self, time, outputs, state, sample_ids): + """Returns `(finished, next_inputs, next_state)`.""" + pass + + @abc.abstractproperty + def batch_size(self): + """Batch size of tensor returned by `sample`. + + Returns a scalar int32 tensor. The return value might not available before + the invocation of initialize(), in this case, ValueError is raised. + """ + raise NotImplementedError("batch_size has not been implemented") + + @abc.abstractproperty + def sample_ids_shape(self): + """Shape of tensor returned by `sample`, excluding the batch dimension. + + Returns a `TensorShape`. The return value might not available before the + invocation of initialize(). + """ + raise NotImplementedError("sample_ids_shape has not been implemented") + + @abc.abstractproperty + def sample_ids_dtype(self): + """DType of tensor returned by `sample`. + + Returns a DType. The return value might not available before the + invocation of initialize(). + """ + raise NotImplementedError("sample_ids_dtype has not been implemented") + + +class CustomSampler(Sampler): + """Base abstract class that allows the user to customize sampling.""" + + def __init__(self, + initialize_fn, + sample_fn, + next_inputs_fn, + sample_ids_shape=None, + sample_ids_dtype=None): + """Initializer. + + Args: + initialize_fn: callable that returns `(finished, next_inputs)` for the + first iteration. + sample_fn: callable that takes `(time, outputs, state)` and emits tensor + `sample_ids`. + next_inputs_fn: callable that takes `(time, outputs, state, sample_ids)` + and emits `(finished, next_inputs, next_state)`. + sample_ids_shape: Either a list of integers, or a 1-D Tensor of type + `int32`, the shape of each value in the `sample_ids` batch. Defaults to + a scalar. + sample_ids_dtype: The dtype of the `sample_ids` tensor. Defaults to int32. + """ + self._initialize_fn = initialize_fn + self._sample_fn = sample_fn + self._next_inputs_fn = next_inputs_fn + self._batch_size = None + self._sample_ids_shape = tensor_shape.TensorShape(sample_ids_shape or []) + self._sample_ids_dtype = sample_ids_dtype or dtypes.int32 + + @property + def batch_size(self): + if self._batch_size is None: + raise ValueError("batch_size accessed before initialize was called") + return self._batch_size + + @property + def sample_ids_shape(self): + return self._sample_ids_shape + + @property + def sample_ids_dtype(self): + return self._sample_ids_dtype + + def initialize(self, inputs, **kwargs): + (finished, next_inputs) = self._initialize_fn(inputs, **kwargs) + if self._batch_size is None: + self._batch_size = array_ops.size(finished) + return (finished, next_inputs) + + def sample(self, time, outputs, state): + return self._sample_fn(time=time, outputs=outputs, state=state) + + def next_inputs(self, time, outputs, state, sample_ids): + return self._next_inputs_fn( + time=time, outputs=outputs, state=state, sample_ids=sample_ids) + + +class TrainingSampler(Sampler): + """A Sampler for use during training. + + Only reads inputs. + + Returned sample_ids are the argmax of the RNN output logits. + """ + + def __init__(self, time_major=False): + """Initializer. + + Args: + time_major: Python bool. Whether the tensors in `inputs` are time major. + If `False` (default), they are assumed to be batch major. + + Raises: + ValueError: if `sequence_length` is not a 1D tensor. + """ + self.time_major = time_major + self._batch_size = None + + @property + def batch_size(self): + if self._batch_size is None: + raise ValueError("batch_size accessed before initialize was called") + return self._batch_size + + @property + def sample_ids_shape(self): + return tensor_shape.TensorShape([]) + + @property + def sample_ids_dtype(self): + return dtypes.int32 + + def initialize(self, inputs, sequence_length=None): + """Initialize the TrainSampler. + + Args: + inputs: A (structure of) input tensors. + sequence_length: An int32 vector tensor. + + Returns: + (finished, next_inputs), a tuple of two items. The first item is a boolean + vector to indicate whether the item in the batch has finished. The + second item is the first slide of input data based on the timestep + dimension (usually the second dim of the input). + """ + self.inputs = ops.convert_to_tensor(inputs, name="inputs") + if not self.time_major: + inputs = nest.map_structure(_transpose_batch_time, inputs) + + self.input_tas = nest.map_structure(_unstack_ta, inputs) + if sequence_length is None: + raise ValueError("sequence_length is required for TrainingSampler") + self.sequence_length = ops.convert_to_tensor( + sequence_length, name="sequence_length") + if self.sequence_length.get_shape().ndims != 1: + raise ValueError( + "Expected sequence_length to be a vector, but received shape: %s" % + self._sequence_length.get_shape()) + + self.zero_inputs = nest.map_structure( + lambda inp: array_ops.zeros_like(inp[0, :]), inputs) + + self._batch_size = array_ops.size(self.sequence_length) + + finished = math_ops.equal(0, self.sequence_length) + all_finished = math_ops.reduce_all(finished) + next_inputs = control_flow_ops.cond( + all_finished, + lambda: self.zero_inputs, + lambda: nest.map_structure(lambda inp: inp.read(0), self.input_tas)) + return (finished, next_inputs) + + def sample(self, time, outputs, state): + del state + sample_ids = math_ops.cast(math_ops.argmax(outputs, axis=-1), dtypes.int32) + return sample_ids + + def next_inputs(self, time, outputs, state, sample_ids): + del sample_ids + next_time = time + 1 + finished = (next_time >= self.sequence_length) + all_finished = math_ops.reduce_all(finished) + + def read_from_ta(inp): + return inp.read(next_time) + + next_inputs = control_flow_ops.cond( + all_finished, + lambda: self.zero_inputs, + lambda: nest.map_structure(read_from_ta, self.input_tas)) + return (finished, next_inputs, state) + + +class ScheduledEmbeddingTrainingSampler(TrainingSampler): + """A training sampler that adds scheduled sampling. + + Returns -1s for sample_ids where no sampling took place; valid sample id + values elsewhere. + """ + + def __init__(self, + sampling_probability, + embedding_fn=None, + time_major=False, + seed=None, + scheduling_seed=None): + """Initializer. + + Args: + sampling_probability: A `float32` 0-D or 1-D tensor: the probability of + sampling categorically from the output ids instead of reading directly + from the inputs. + embedding_fn: A callable that takes a vector tensor of `ids` (argmax ids), + or the `params` argument for `embedding_lookup`. + time_major: Python bool. Whether the tensors in `inputs` are time major. + If `False` (default), they are assumed to be batch major. + seed: The sampling seed. + scheduling_seed: The schedule decision rule sampling seed. + + Raises: + ValueError: if `sampling_probability` is not a scalar or vector. + """ + if callable(embedding_fn) or embedding_fn is None: + self.embedding_fn = embedding_fn + else: + raise ValueError("embedding_fn is expected to be callable, got %s" + % type(embedding_fn)) + self.sampling_probability = ops.convert_to_tensor( + sampling_probability, name="sampling_probability") + if self.sampling_probability.get_shape().ndims not in (0, 1): + raise ValueError( + "sampling_probability must be either a scalar or a vector. " + "saw shape: %s" % (self.sampling_probability.get_shape())) + self.seed = seed + self.scheduling_seed = scheduling_seed + super(ScheduledEmbeddingTrainingSampler, + self).__init__(time_major=time_major) + + def initialize(self, inputs, sequence_length=None, embedding=None): + if self.embedding_fn is None: + if embedding is None: + raise ValueError("embedding is required as a keyword argument for " + "ScheduledEmbeddingTrainingSampler") + self.embedding_fn = ( + lambda ids: embedding_ops.embedding_lookup(embedding, ids)) + return super(ScheduledEmbeddingTrainingSampler, self).initialize( + inputs, sequence_length=sequence_length) + + def sample(self, time, outputs, state): + del state + # Return -1s where we did not sample, and sample_ids elsewhere + select_sample = bernoulli_sample( + probs=self.sampling_probability, + dtype=dtypes.bool, + sample_shape=self.batch_size, + seed=self.scheduling_seed) + return array_ops.where(select_sample, + categorical_sample(logits=outputs, seed=self.seed), + gen_array_ops.fill([self.batch_size], -1)) + + def next_inputs(self, time, outputs, state, sample_ids): + (finished, base_next_inputs, state) = ( + super(ScheduledEmbeddingTrainingSampler, self).next_inputs( + time=time, outputs=outputs, state=state, sample_ids=sample_ids)) + + def maybe_sample(): + """Perform scheduled sampling.""" + where_sampling = math_ops.cast( + array_ops.where(sample_ids > -1), dtypes.int32) + where_not_sampling = math_ops.cast( + array_ops.where(sample_ids <= -1), dtypes.int32) + sample_ids_sampling = array_ops.gather_nd(sample_ids, where_sampling) + inputs_not_sampling = array_ops.gather_nd(base_next_inputs, + where_not_sampling) + sampled_next_inputs = self.embedding_fn(sample_ids_sampling) + base_shape = array_ops.shape(base_next_inputs) + return (array_ops.scatter_nd( + indices=where_sampling, updates=sampled_next_inputs, shape=base_shape) + + array_ops.scatter_nd( + indices=where_not_sampling, + updates=inputs_not_sampling, + shape=base_shape)) + + all_finished = math_ops.reduce_all(finished) + next_inputs = control_flow_ops.cond(all_finished, lambda: base_next_inputs, + maybe_sample) + return (finished, next_inputs, state) + + +class ScheduledOutputTrainingSampler(TrainingSampler): + """A training sampler that adds scheduled sampling directly to outputs. + + Returns False for sample_ids where no sampling took place; True elsewhere. + """ + + def __init__(self, + sampling_probability, + time_major=False, + seed=None, + next_inputs_fn=None): + """Initializer. + + Args: + sampling_probability: A `float32` scalar tensor: the probability of + sampling from the outputs instead of reading directly from the inputs. + time_major: Python bool. Whether the tensors in `inputs` are time major. + If `False` (default), they are assumed to be batch major. + seed: The sampling seed. + next_inputs_fn: (Optional) callable to apply to the RNN outputs to create + the next input when sampling. If `None` (default), the RNN outputs will + be used as the next inputs. + + Raises: + ValueError: if `sampling_probability` is not a scalar or vector. + """ + self.sampling_probability = ops.convert_to_tensor( + sampling_probability, name="sampling_probability") + if self.sampling_probability.get_shape().ndims not in (0, 1): + raise ValueError( + "sampling_probability must be either a scalar or a vector. " + "saw shape: %s" % (self._sampling_probability.get_shape())) + + self.seed = seed + self.next_inputs_fn = next_inputs_fn + + super(ScheduledOutputTrainingSampler, self).__init__(time_major=time_major) + + def initialize(self, inputs, sequence_length=None, auxiliary_inputs=None): + if auxiliary_inputs is None: + maybe_concatenated_inputs = inputs + else: + inputs = ops.convert_to_tensor(inputs) + auxiliary_inputs = ops.convert_to_tensor(auxiliary_inputs) + maybe_concatenated_inputs = nest.map_structure( + lambda x, y: array_ops.concat((x, y), -1), inputs, auxiliary_inputs) + if not self.time_major: + auxiliary_inputs = nest.map_structure(_transpose_batch_time, + auxiliary_inputs) + if auxiliary_inputs is not None: + self._auxiliary_input_tas = nest.map_structure(_unstack_ta, + auxiliary_inputs) + else: + self._auxiliary_input_tas = None + + return super(ScheduledOutputTrainingSampler, self).initialize( + maybe_concatenated_inputs, sequence_length=sequence_length) + + def sample(self, time, outputs, state): + del state + return bernoulli_sample( + probs=self.sampling_probability, + sample_shape=self.batch_size, + seed=self.seed) + + def next_inputs(self, time, outputs, state, sample_ids): + (finished, base_next_inputs, state) = ( + super(ScheduledOutputTrainingSampler, self).next_inputs( + time=time, outputs=outputs, state=state, sample_ids=sample_ids)) + sample_ids = math_ops.cast(sample_ids, dtypes.bool) + + def maybe_sample(): + """Perform scheduled sampling.""" + + def maybe_concatenate_auxiliary_inputs(outputs_, indices=None): + """Concatenate outputs with auxiliary inputs, if they exist.""" + if self._auxiliary_input_tas is None: + return outputs_ + + next_time = time + 1 + auxiliary_inputs = nest.map_structure(lambda ta: ta.read(next_time), + self._auxiliary_input_tas) + if indices is not None: + auxiliary_inputs = array_ops.gather_nd(auxiliary_inputs, indices) + return nest.map_structure(lambda x, y: array_ops.concat((x, y), -1), + outputs_, auxiliary_inputs) + + if self.next_inputs_fn is None: + return array_ops.where(sample_ids, + maybe_concatenate_auxiliary_inputs(outputs), + base_next_inputs) + + where_sampling = math_ops.cast(array_ops.where(sample_ids), dtypes.int32) + where_not_sampling = math_ops.cast( + array_ops.where(math_ops.logical_not(sample_ids)), dtypes.int32) + outputs_sampling = array_ops.gather_nd(outputs, where_sampling) + inputs_not_sampling = array_ops.gather_nd(base_next_inputs, + where_not_sampling) + sampled_next_inputs = maybe_concatenate_auxiliary_inputs( + self.next_inputs_fn(outputs_sampling), where_sampling) + + base_shape = array_ops.shape(base_next_inputs) + return (array_ops.scatter_nd( + indices=where_sampling, updates=sampled_next_inputs, shape=base_shape) + + array_ops.scatter_nd( + indices=where_not_sampling, + updates=inputs_not_sampling, + shape=base_shape)) + + all_finished = math_ops.reduce_all(finished) + no_samples = math_ops.logical_not(math_ops.reduce_any(sample_ids)) + next_inputs = control_flow_ops.cond( + math_ops.logical_or(all_finished, no_samples), lambda: base_next_inputs, + maybe_sample) + return (finished, next_inputs, state) + + +class GreedyEmbeddingSampler(Sampler): + """A sampler for use during inference. + + Uses the argmax of the output (treated as logits) and passes the + result through an embedding layer to get the next input. + """ + + def __init__(self, embedding_fn=None): + """Initializer. + + Args: + embedding_fn: A optional callable that takes a vector tensor of `ids` + (argmax ids), or the `params` argument for `embedding_lookup`. The + returned tensor will be passed to the decoder input. Default to use + `embedding_ops.embedding_lookup`. + """ + if embedding_fn is None or callable(embedding_fn): + self.embedding_fn = embedding_fn + else: + raise ValueError("embedding_fn is expected to be a callable, got %s" % + type(embedding_fn)) + self._batch_size = None + + @property + def batch_size(self): + if self._batch_size is None: + raise ValueError("batch_size accessed before initialize was called") + return self._batch_size + + @property + def sample_ids_shape(self): + return tensor_shape.TensorShape([]) + + @property + def sample_ids_dtype(self): + return dtypes.int32 + + def initialize(self, embedding, start_tokens=None, end_token=None): + """Initialize the GreedyEmbeddingSampler. + + Args: + embedding: tensor that contains embedding states matrix. It will be used + to generate generate outputs with start_tokens and end_tokens. The + embedding will be ignored if the embedding_fn has been provided at + __init__(). + start_tokens: `int32` vector shaped `[batch_size]`, the start tokens. + end_token: `int32` scalar, the token that marks end of decoding. + + Returns: + Tuple of two items: `(finished, self.start_inputs)`. + Raises: + ValueError: if `start_tokens` is not a 1D tensor or `end_token` is not a + scalar. + """ + if self.embedding_fn is None: + self.embedding_fn = ( + lambda ids: embedding_ops.embedding_lookup(embedding, ids)) + + self.start_tokens = ops.convert_to_tensor( + start_tokens, dtype=dtypes.int32, name="start_tokens") + self.end_token = ops.convert_to_tensor( + end_token, dtype=dtypes.int32, name="end_token") + if self.start_tokens.get_shape().ndims != 1: + raise ValueError("start_tokens must be a vector") + self._batch_size = array_ops.size(start_tokens) + if self.end_token.get_shape().ndims != 0: + raise ValueError("end_token must be a scalar") + self.start_inputs = self.embedding_fn(self.start_tokens) + + finished = array_ops.tile([False], [self._batch_size]) + return (finished, self.start_inputs) + + def sample(self, time, outputs, state): + """sample for GreedyEmbeddingHelper.""" + del time, state # unused by sample_fn + # Outputs are logits, use argmax to get the most probable id + if not isinstance(outputs, ops.Tensor): + raise TypeError( + "Expected outputs to be a single Tensor, got: %s" % type(outputs)) + sample_ids = math_ops.argmax(outputs, axis=-1, output_type=dtypes.int32) + return sample_ids + + def next_inputs(self, time, outputs, state, sample_ids): + """next_inputs_fn for GreedyEmbeddingHelper.""" + del time, outputs # unused by next_inputs_fn + finished = math_ops.equal(sample_ids, self.end_token) + all_finished = math_ops.reduce_all(finished) + next_inputs = control_flow_ops.cond( + all_finished, + # If we're finished, the next_inputs value doesn't matter + lambda: self.start_inputs, + lambda: self.embedding_fn(sample_ids)) + return (finished, next_inputs, state) + + +class SampleEmbeddingSampler(GreedyEmbeddingSampler): + """A sampler for use during inference. + + Uses sampling (from a distribution) instead of argmax and passes the + result through an embedding layer to get the next input. + """ + + def __init__(self, embedding_fn=None, softmax_temperature=None, seed=None): + """Initializer. + + Args: + embedding_fn: (Optional) A callable that takes a vector tensor of `ids` + (argmax ids), or the `params` argument for `embedding_lookup`. The + returned tensor will be passed to the decoder input. + softmax_temperature: (Optional) `float32` scalar, value to divide the + logits by before computing the softmax. Larger values (above 1.0) result + in more random samples, while smaller values push the sampling + distribution towards the argmax. Must be strictly greater than 0. + Defaults to 1.0. + seed: (Optional) The sampling seed. + + Raises: + ValueError: if `start_tokens` is not a 1D tensor or `end_token` is not a + scalar. + """ + super(SampleEmbeddingSampler, self).__init__(embedding_fn) + self.softmax_temperature = softmax_temperature + self.seed = seed + + def sample(self, time, outputs, state): + """sample for SampleEmbeddingHelper.""" + del time, state # unused by sample_fn + # Outputs are logits, we sample instead of argmax (greedy). + if not isinstance(outputs, ops.Tensor): + raise TypeError( + "Expected outputs to be a single Tensor, got: %s" % type(outputs)) + if self.softmax_temperature is None: + logits = outputs + else: + logits = outputs / self.softmax_temperature + + return categorical_sample(logits=logits, seed=self.seed) + + +class InferenceSampler(Sampler): + """A helper to use during inference with a custom sampling function.""" + + def __init__(self, + sample_fn, + sample_shape, + sample_dtype, + end_fn, + next_inputs_fn=None): + """Initializer. + + Args: + sample_fn: A callable that takes `outputs` and emits tensor `sample_ids`. + sample_shape: Either a list of integers, or a 1-D Tensor of type `int32`, + the shape of the each sample in the batch returned by `sample_fn`. + sample_dtype: the dtype of the sample returned by `sample_fn`. + end_fn: A callable that takes `sample_ids` and emits a `bool` vector + shaped `[batch_size]` indicating whether each sample is an end token. + next_inputs_fn: (Optional) A callable that takes `sample_ids` and returns + the next batch of inputs. If not provided, `sample_ids` is used as the + next batch of inputs. + """ + self.sample_fn = sample_fn + self.sample_shape = tensor_shape.TensorShape(sample_shape) + self.sample_dtype = sample_dtype + self.end_fn = end_fn + self.next_inputs_fn = next_inputs_fn + self._batch_size = None + + @property + def batch_size(self): + if self._batch_size is None: + raise ValueError("batch_size accessed before initialize was called") + return self._batch_size + + @property + def sample_ids_shape(self): + return self.sample_shape + + @property + def sample_ids_dtype(self): + return self.sample_dtype + + def initialize(self, start_inputs): + self.start_inputs = ops.convert_to_tensor(start_inputs, name="start_inputs") + self._batch_size = array_ops.shape(start_inputs)[0] + finished = array_ops.tile([False], [self._batch_size]) + return (finished, self.start_inputs) + + def sample(self, time, outputs, state): + del time, state # unused by sample + return self.sample_fn(outputs) + + def next_inputs(self, time, outputs, state, sample_ids): + del time, outputs # unused by next_inputs + if self.next_inputs_fn is None: + next_inputs = sample_ids + else: + next_inputs = self.next_inputs_fn(sample_ids) + finished = self.end_fn(sample_ids) + return (finished, next_inputs, state) + + +# The following sample functions (_call_sampler, bernoulli_sample, +# categorical_sample) mimic TensorFlow Probability distribution semantics. +def _call_sampler(sample_n_fn, sample_shape, name=None): + """Reshapes vector of samples.""" + with ops.name_scope(name, "call_sampler", values=[sample_shape]): + sample_shape = ops.convert_to_tensor( + sample_shape, dtype=dtypes.int32, name="sample_shape") + # Ensure sample_shape is a vector (vs just a scalar). + pad = math_ops.cast( + math_ops.equal(array_ops.rank(sample_shape), 0), dtypes.int32) + sample_shape = array_ops.reshape( + sample_shape, + array_ops.pad( + array_ops.shape(sample_shape), + paddings=[[pad, 0]], + constant_values=1)) + samples = sample_n_fn(math_ops.reduce_prod(sample_shape)) + batch_event_shape = array_ops.shape(samples)[1:] + final_shape = array_ops.concat([sample_shape, batch_event_shape], 0) + return array_ops.reshape(samples, final_shape) + + +def bernoulli_sample(probs=None, + logits=None, + dtype=dtypes.int32, + sample_shape=(), + seed=None): + """Samples from Bernoulli distribution.""" + if probs is None: + probs = math_ops.sigmoid(logits, name="probs") + else: + probs = ops.convert_to_tensor(probs, name="probs") + batch_shape_tensor = array_ops.shape(probs) + + def _sample_n(n): + """Sample vector of Bernoullis.""" + new_shape = array_ops.concat([[n], batch_shape_tensor], 0) + uniform = random_ops.random_uniform(new_shape, seed=seed, dtype=probs.dtype) + return math_ops.cast(math_ops.less(uniform, probs), dtype) + + return _call_sampler(_sample_n, sample_shape) + + +def categorical_sample(logits, dtype=dtypes.int32, sample_shape=(), seed=None): + """Samples from categorical distribution.""" + logits = ops.convert_to_tensor(logits, name="logits") + event_size = array_ops.shape(logits)[-1] + batch_shape_tensor = array_ops.shape(logits)[:-1] + + def _sample_n(n): + """Sample vector of categoricals.""" + if logits.shape.ndims == 2: + logits_2d = logits + else: + logits_2d = array_ops.reshape(logits, [-1, event_size]) + sample_dtype = dtypes.int64 if logits.dtype.size > 4 else dtypes.int32 + draws = random_ops.multinomial( + logits_2d, n, seed=seed, output_dtype=sample_dtype) + draws = array_ops.reshape( + array_ops.transpose(draws), + array_ops.concat([[n], batch_shape_tensor], 0)) + return math_ops.cast(draws, dtype) + + return _call_sampler(_sample_n, sample_shape) + + +def _unstack_ta(inp): + return tensor_array_ops.TensorArray( + dtype=inp.dtype, + size=array_ops.shape(inp)[0], + element_shape=inp.get_shape()[1:]).unstack(inp) diff --git a/tensorflow/contrib/session_bundle/exporter.py b/tensorflow/contrib/session_bundle/exporter.py index 08983337fccc138d40eb959cecc5bf9e47cf6cac..f3efd292cf5acba4319c8a5545a7f70fae4b5ce1 100644 --- a/tensorflow/contrib/session_bundle/exporter.py +++ b/tensorflow/contrib/session_bundle/exporter.py @@ -304,10 +304,10 @@ class Exporter(object): def parser(path): if os.name == "nt": match = re.match( - "^" + export_dir_base.replace("\\", "/") + "/(\\d{8})$", + r"^" + export_dir_base.replace("\\", "/") + r"/(\d{8})$", path.path.replace("\\", "/")) else: - match = re.match("^" + export_dir_base + "/(\\d{8})$", path.path) + match = re.match(r"^" + export_dir_base + r"/(\d{8})$", path.path) if not match: return None return path._replace(export_version=int(match.group(1))) diff --git a/tensorflow/contrib/session_bundle/gc_test.py b/tensorflow/contrib/session_bundle/gc_test.py index 8faf3ef3d4cd7ee0096265283070e25d06782254..02725bb1cbb4ef9ace29dcc58f6d23fb241d96b2 100644 --- a/tensorflow/contrib/session_bundle/gc_test.py +++ b/tensorflow/contrib/session_bundle/gc_test.py @@ -104,7 +104,7 @@ class GcTest(test_util.TensorFlowTestCase): # create a simple parser that pulls the export_version from the directory. def parser(path): - match = re.match("^" + base_dir + "/(\\d+)$", path.path) + match = re.match(r"^" + base_dir + r"/(\d+)$", path.path) if not match: return None return path._replace(export_version=int(match.group(1))) diff --git a/tensorflow/contrib/slim/python/slim/data/tfexample_decoder.py b/tensorflow/contrib/slim/python/slim/data/tfexample_decoder.py index 1b2b6acacca838f95cb758ae88f79263993ca69e..c63a3ca19b6a70cf7776c7fce4e0291ee94b775c 100644 --- a/tensorflow/contrib/slim/python/slim/data/tfexample_decoder.py +++ b/tensorflow/contrib/slim/python/slim/data/tfexample_decoder.py @@ -32,7 +32,7 @@ from tensorflow.python.framework import dtypes 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 functional_ops +from tensorflow.python.ops import map_fn from tensorflow.python.ops import image_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import parsing_ops @@ -396,8 +396,8 @@ class Image(ItemHandler): image_format = keys_to_tensors[self._format_key] if self._repeated: - return functional_ops.map_fn(lambda x: self._decode(x, image_format), - image_buffer, dtype=self._dtype) + return map_fn.map_fn(lambda x: self._decode(x, image_format), + image_buffer, dtype=self._dtype) else: return self._decode(image_buffer, image_format) diff --git a/tensorflow/contrib/slim/python/slim/nets/BUILD b/tensorflow/contrib/slim/python/slim/nets/BUILD index 8bbdf96384683c68648367c6433eeb89c64c22bf..e9595d1b324dbd3d570d2407a6620c5295b15548 100644 --- a/tensorflow/contrib/slim/python/slim/nets/BUILD +++ b/tensorflow/contrib/slim/python/slim/nets/BUILD @@ -115,9 +115,9 @@ py_library( py_test( name = "inception_v1_test", - size = "large", + size = "medium", srcs = ["inception_v1_test.py"], - shard_count = 3, + shard_count = 8, srcs_version = "PY2AND3", deps = [ ":inception_v1", @@ -135,9 +135,9 @@ py_test( py_test( name = "inception_v2_test", - size = "large", + size = "medium", srcs = ["inception_v2_test.py"], - shard_count = 3, + shard_count = 8, srcs_version = "PY2AND3", deps = [ ":inception_v2", @@ -155,9 +155,9 @@ py_test( py_test( name = "inception_v3_test", - size = "large", + size = "medium", srcs = ["inception_v3_test.py"], - shard_count = 3, + shard_count = 8, srcs_version = "PY2AND3", deps = [ ":inception_v3", @@ -233,8 +233,9 @@ py_library( py_test( name = "resnet_v1_test", - size = "large", + size = "medium", srcs = ["resnet_v1_test.py"], + shard_count = 4, srcs_version = "PY2AND3", deps = [ ":resnet_utils", @@ -268,8 +269,9 @@ py_library( py_test( name = "resnet_v2_test", - size = "large", + size = "medium", srcs = ["resnet_v2_test.py"], + shard_count = 4, srcs_version = "PY2AND3", deps = [ ":resnet_utils", diff --git a/tensorflow/contrib/sparsemax/BUILD b/tensorflow/contrib/sparsemax/BUILD index d7ba754f701d4b433e35ad8396eae7ee6132b97f..ed4eca1a60a6f0ccf629d8aa7906c02092e25ba0 100644 --- a/tensorflow/contrib/sparsemax/BUILD +++ b/tensorflow/contrib/sparsemax/BUILD @@ -49,6 +49,9 @@ cuda_py_tests( "//tensorflow/python:math_ops", "//tensorflow/python:platform_test", ], + tags = [ + "oss_serial", + ], ) cuda_py_tests( @@ -64,4 +67,7 @@ cuda_py_tests( "//tensorflow/python:math_ops", "//tensorflow/python:platform_test", ], + tags = [ + "oss_serial", + ], ) diff --git a/tensorflow/contrib/tensor_forest/BUILD b/tensorflow/contrib/tensor_forest/BUILD index 398ac314f4b520610ec100273b37c33bc4b5b43a..583bbf97c57cf263f65bc3b0a56b32cc2dce5482 100644 --- a/tensorflow/contrib/tensor_forest/BUILD +++ b/tensorflow/contrib/tensor_forest/BUILD @@ -537,8 +537,9 @@ py_library( py_test( name = "random_forest_test", - size = "large", + size = "medium", srcs = ["client/random_forest_test.py"], + shard_count = 6, srcs_version = "PY2AND3", tags = [ "noasan", diff --git a/tensorflow/contrib/tensor_forest/client/eval_metrics.py b/tensorflow/contrib/tensor_forest/client/eval_metrics.py index d8236a0a6fa6d0d0e383e454eb0146bb10b6f49d..0d87cea9fbaa8fe28b55ec996414a568d39efee3 100644 --- a/tensorflow/contrib/tensor_forest/client/eval_metrics.py +++ b/tensorflow/contrib/tensor_forest/client/eval_metrics.py @@ -50,9 +50,10 @@ def _accuracy(predictions, targets, weights=None): def _r2(probabilities, targets, weights=None): targets = math_ops.to_float(targets) y_mean = math_ops.reduce_mean(targets, 0) - squares_total = math_ops.reduce_sum(math_ops.square(targets - y_mean), 0) + squares_total = math_ops.reduce_sum( + math_ops.squared_difference(targets, y_mean), 0) squares_residuals = math_ops.reduce_sum( - math_ops.square(targets - probabilities), 0) + math_ops.squared_difference(targets, probabilities), 0) score = 1 - math_ops.reduce_sum(squares_residuals / squares_total) return metrics.mean(score, weights=weights) diff --git a/tensorflow/contrib/tensor_forest/kernels/model_ops.cc b/tensorflow/contrib/tensor_forest/kernels/model_ops.cc index b9aad36f3d25b9fb7b8b525be54fb7a39394b373..76b1d2b4da269cda71f5b49878f2933d7d9b5776 100644 --- a/tensorflow/contrib/tensor_forest/kernels/model_ops.cc +++ b/tensorflow/contrib/tensor_forest/kernels/model_ops.cc @@ -304,7 +304,7 @@ class TraverseTreeV4Op : public OpKernel { auto worker_threads = context->device()->tensorflow_cpu_worker_threads(); int num_threads = worker_threads->num_threads; const int64 costPerTraverse = 500; - auto traverse = [this, &set_leaf_ids, &data_set, decision_tree_resource, + auto traverse = [&set_leaf_ids, &data_set, decision_tree_resource, num_data](int64 start, int64 end) { CHECK(start <= end); CHECK(end <= num_data); diff --git a/tensorflow/contrib/tensor_forest/kernels/stats_ops.cc b/tensorflow/contrib/tensor_forest/kernels/stats_ops.cc index fe2c91c1047fe56710b1a86b2fa3206caf6ff3bc..0243f106814511c1b53a5aacb830b845214a00a3 100644 --- a/tensorflow/contrib/tensor_forest/kernels/stats_ops.cc +++ b/tensorflow/contrib/tensor_forest/kernels/stats_ops.cc @@ -307,7 +307,7 @@ class ProcessInputOp : public OpKernel { // from a digits run on local desktop. Heuristics might be necessary // if it really matters that much. const int64 costPerUpdate = 1000; - auto update = [this, &target, &leaf_ids_tensor, &num_targets, &data_set, + auto update = [&target, &leaf_ids_tensor, &num_targets, &data_set, fertile_stats_resource, &locks, &set_lock, &ready_to_split, num_data](int64 start, int64 end) { CHECK(start <= end); @@ -317,7 +317,7 @@ class ProcessInputOp : public OpKernel { static_cast(end), &ready_to_split); }; - auto update_collated = [this, &target, &num_targets, fertile_stats_resource, + auto update_collated = [&target, &num_targets, fertile_stats_resource, tree_resource, &leaf_examples, &set_lock, &ready_to_split, &data_set, num_leaves](int64 start, int64 end) { diff --git a/tensorflow/contrib/tensor_forest/kernels/tree_utils.h b/tensorflow/contrib/tensor_forest/kernels/tree_utils.h index e04eb60f9b27cfd8b6b4e1502594d4d310ae55cc..774da472f1543f938d1b607ebdef008f7b540211 100644 --- a/tensorflow/contrib/tensor_forest/kernels/tree_utils.h +++ b/tensorflow/contrib/tensor_forest/kernels/tree_utils.h @@ -18,10 +18,10 @@ #include #include "tensorflow/contrib/tensor_forest/kernels/data_spec.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_types.h" -#include "tensorflow/core/kernels/bounds_check.h" #include "tensorflow/core/lib/random/distribution_sampler.h" #include "tensorflow/core/lib/random/simple_philox.h" #include "tensorflow/core/lib/strings/strcat.h" diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.h b/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.h index d3edb43733761a906c6e5bf8b65f76e3e1ae56fc..3100a5a0e5da1103b61bd089cd433721686b9e72 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/decision-tree-resource.h @@ -32,7 +32,7 @@ class DecisionTreeResource : public ResourceBase { // Constructor. explicit DecisionTreeResource(const TensorForestParams& params); - string DebugString() override { + string DebugString() const override { return strings::StrCat("DecisionTree[size=", decision_tree_->decision_tree().nodes_size(), "]"); } diff --git a/tensorflow/contrib/tensor_forest/kernels/v4/fertile-stats-resource.h b/tensorflow/contrib/tensor_forest/kernels/v4/fertile-stats-resource.h index eea0be27caf0a022ba7acaacd359c75a2df4eedb..44f2b3f473b9eced06bd800b9cf0a5a0825ec3eb 100644 --- a/tensorflow/contrib/tensor_forest/kernels/v4/fertile-stats-resource.h +++ b/tensorflow/contrib/tensor_forest/kernels/v4/fertile-stats-resource.h @@ -40,7 +40,7 @@ class FertileStatsResource : public ResourceBase { model_op_ = LeafModelOperatorFactory::CreateLeafModelOperator(params_); } - string DebugString() override { return "FertileStats"; } + string DebugString() const override { return "FertileStats"; } void ExtractFromProto(const FertileStats& stats); diff --git a/tensorflow/contrib/tensorrt/BUILD b/tensorflow/contrib/tensorrt/BUILD index da123e1623f9abc12dd7a44f3d7e4127740a62df..91b6d2614a8963c21e35c385411dc4c9956e3146 100644 --- a/tensorflow/contrib/tensorrt/BUILD +++ b/tensorflow/contrib/tensorrt/BUILD @@ -11,568 +11,54 @@ exports_files(["LICENSE"]) load( "//tensorflow:tensorflow.bzl", - "tf_cc_test", - "tf_copts", "tf_cuda_library", - "tf_custom_op_library", "tf_custom_op_library_additional_deps", - "tf_gen_op_libs", - "tf_gen_op_wrapper_py", ) -load("//tensorflow:tensorflow.bzl", "cuda_py_test") -load("//tensorflow:tensorflow.bzl", "cuda_py_tests") -load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test") -load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") -load("//tensorflow:tensorflow.bzl", "tf_py_wrap_cc") load( "@local_config_tensorrt//:build_defs.bzl", "if_tensorrt", ) -exports_files(glob([ - "test/testdata/*", -])) - -tf_cuda_cc_test( - name = "tensorrt_test_cc", - size = "small", - srcs = ["tensorrt_test.cc"], - tags = [ - "no_windows", - "nomac", - ], - deps = [ - "//tensorflow/core:gpu_init", - "//tensorflow/core:lib", - "//tensorflow/core:stream_executor", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - ] + if_tensorrt([ - "@local_config_cuda//cuda:cuda_headers", - "@local_config_tensorrt//:nv_infer", - ]), -) - -tf_custom_op_library( - name = "python/ops/_trt_engine_op.so", - srcs = [ - "ops/trt_engine_op.cc", - ], - deps = [ - ":trt_shape_function", - "//tensorflow/core:lib_proto_parsing", - ] + if_tensorrt([ - "@local_config_tensorrt//:nv_infer", - ]), -) - tf_cuda_library( name = "trt_shape_function", srcs = ["shape_fn/trt_shfn.cc"], hdrs = ["shape_fn/trt_shfn.h"], visibility = ["//visibility:public"], deps = [ - ":trt_logging", - ":trt_plugins", + "//tensorflow/compiler/tf2tensorrt:trt_logging", + "//tensorflow/compiler/tf2tensorrt:trt_plugins", ] + if_tensorrt([ - "@local_config_tensorrt//:nv_infer", + "@local_config_tensorrt//:tensorrt", ]) + tf_custom_op_library_additional_deps(), ) -cc_library( - name = "trt_engine_op_kernel", - srcs = [ - "kernels/trt_engine_op.cc", - ], - hdrs = [ - "kernels/trt_engine_op.h", - ], - copts = tf_copts(), - visibility = ["//visibility:public"], - deps = [ - ":test_utils", - ":trt_allocator", - ":trt_conversion", - ":trt_logging", - ":trt_plugins", - ":trt_resources", - ":utils", - "//tensorflow/core:gpu_headers_lib", - "//tensorflow/core:lib_proto_parsing", - "//tensorflow/core:stream_executor_headers_lib", - "//tensorflow/core/grappler/costs:graph_properties", - ] + if_tensorrt([ - "@local_config_tensorrt//:nv_infer", - ]) + tf_custom_op_library_additional_deps(), - # TODO(laigd): fix this by merging header file in cc file. - alwayslink = 1, # buildozer: disable=alwayslink-with-hdrs -) - -tf_gen_op_libs( - op_lib_names = [ - "trt_engine_op", - ], -) - -tf_cuda_library( - name = "trt_logging", - srcs = ["log/trt_logger.cc"], - hdrs = ["log/trt_logger.h"], - visibility = ["//visibility:public"], - deps = [ - "//tensorflow/core:lib_proto_parsing", - ] + if_tensorrt([ - "@local_config_tensorrt//:nv_infer", - ]), -) - -tf_gen_op_wrapper_py( - name = "trt_engine_op", - deps = [ - ":trt_engine_op_op_lib", - ":trt_logging", - ":trt_shape_function", - ], -) - -tf_custom_op_py_library( - name = "trt_engine_op_loader", - srcs = ["python/ops/trt_engine_op.py"], - dso = [ - ":python/ops/_trt_engine_op.so", - ] + if_tensorrt([ - "@local_config_tensorrt//:nv_infer", - ]), - kernels = [ - ":trt_engine_op_kernel", - ":trt_engine_op_op_lib", - ":trt_shape_function", - ], - srcs_version = "PY2AND3", - deps = [ - "//tensorflow/contrib/util:util_py", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:resources", - ], -) - py_library( name = "init_py", srcs = [ "__init__.py", "python/__init__.py", + "python/trt_convert.py", ], srcs_version = "PY2AND3", deps = [ - ":tf_trt_integration_test_base", - ":trt_convert_py", - ":trt_ops_py", - "//tensorflow/python:errors", - ], -) - -py_library( - name = "trt_ops_py", - srcs_version = "PY2AND3", - deps = [ - ":trt_engine_op", - ":trt_engine_op_loader", - ], -) - -py_library( - name = "trt_convert_py", - srcs = ["python/trt_convert.py"], - srcs_version = "PY2AND3", - deps = [ - ":wrap_conversion", - "//tensorflow/python:graph_util", - "//tensorflow/python:session", - "//tensorflow/python:tf_optimizer", - "//tensorflow/python/saved_model:builder", - "//tensorflow/python/saved_model:loader", - "//tensorflow/python/saved_model:tag_constants", - ], -) - -# TODO(aaroey): this wrapper has been causing troubles of double linking, so -# either get rid of it, or split to make it contain minimum dependencies. -tf_py_wrap_cc( - name = "wrap_conversion", - srcs = ["trt_conversion.i"], - copts = tf_copts(), - swig_includes = [ - "//tensorflow/python:platform/base.i", - ], - deps = [ - ":test_utils", - ":trt_conversion", - ":trt_engine_op_kernel", - "//third_party/python_runtime:headers", - ], -) - -tf_cuda_library( - name = "trt_resources", - srcs = [ - "resources/trt_int8_calibrator.cc", - "resources/trt_resource_manager.cc", - ], - hdrs = [ - "resources/trt_int8_calibrator.h", - "resources/trt_resource_manager.h", - "resources/trt_resources.h", + "//tensorflow/python/compiler/tensorrt:init_py", ], - deps = [ - ":trt_allocator", - ":trt_logging", - ":utils", - "//tensorflow/core:framework_headers_lib", - "//tensorflow/core:framework_lite", - "//tensorflow/core:lib_proto_parsing", - ] + if_tensorrt([ - "@local_config_tensorrt//:nv_infer", - ]), ) -tf_cuda_library( - name = "trt_allocator", - srcs = ["resources/trt_allocator.cc"], - hdrs = ["resources/trt_allocator.h"], - deps = [ - "//tensorflow/core:framework_headers_lib", - "//tensorflow/core:framework_lite", - "//tensorflow/core:lib_proto_parsing", - ] + if_tensorrt([ - "@local_config_tensorrt//:nv_infer", - ]), -) +# The following rules forward the libraries that were moved in order to not +# break other internal targets. -tf_cc_test( - name = "trt_allocator_test", - size = "small", - srcs = ["resources/trt_allocator_test.cc"], - tags = [ - "no_windows", - "nomac", - ], - deps = [ - ":trt_allocator", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - ], -) - -# Library for the node-level conversion portion of TensorRT operation creation -tf_cuda_library( +alias( name = "trt_conversion", - srcs = [ - "convert/convert_graph.cc", - "convert/convert_nodes.cc", - "convert/trt_optimization_pass.cc", - ], - hdrs = [ - "convert/convert_graph.h", - "convert/convert_nodes.h", - "convert/trt_optimization_pass.h", - ], - deps = [ - ":segment", - ":test_utils", - ":trt_allocator", - ":trt_plugins", - ":trt_logging", - ":trt_resources", - ":utils", - "//tensorflow/core/grappler/clusters:cluster", - "//tensorflow/core/grappler/optimizers:custom_graph_optimizer", - "//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry", - "//tensorflow/core/grappler:grappler_item", - "//tensorflow/core/grappler:utils", - "//tensorflow/core:framework", - "//tensorflow/core:framework_lite", - "//tensorflow/core:gpu_runtime", - "//tensorflow/core:graph", - "//tensorflow/core:lib", - "//tensorflow/core:lib_internal", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core/grappler:devices", - "//tensorflow/core/grappler/clusters:virtual_cluster", - "//tensorflow/core/grappler/costs:graph_properties", - "//tensorflow/core/grappler/optimizers:meta_optimizer", - ] + if_tensorrt([ - "@local_config_tensorrt//:nv_infer", - ]) + tf_custom_op_library_additional_deps(), -) - -tf_cuda_cc_test( - name = "convert_graph_test", - size = "medium", - srcs = ["convert/convert_graph_test.cc"], - tags = [ - "no_cuda_on_cpu_tap", - "no_windows", - "nomac", - ], - deps = [ - ":trt_conversion", - "@com_google_googletest//:gtest", - "//tensorflow/cc:cc_ops", - "//tensorflow/cc:scope", - "//tensorflow/core/grappler:grappler_item", - "//tensorflow/core/grappler/clusters:cluster", - "//tensorflow/core:core_cpu", - "//tensorflow/core:core_cpu_base", - "//tensorflow/core:direct_session", - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - "//tensorflow/core:testlib", - ] + if_tensorrt([ - "@local_config_tensorrt//:nv_infer", - ]), -) - -tf_cuda_cc_test( - name = "convert_nodes_test", - size = "medium", - srcs = ["convert/convert_nodes_test.cc"], - tags = [ - "no_cuda_on_cpu_tap", - "no_windows", - "nomac", - ], - deps = [ - ":trt_logging", - ":trt_conversion", - ":trt_plugins", - "@com_google_googletest//:gtest", - "//tensorflow/cc:cc_ops", - "//tensorflow/cc:ops", - "//tensorflow/cc:scope", - "//tensorflow/core/grappler/costs:graph_properties", - "//tensorflow/core:core_cpu", - "//tensorflow/core:core_cpu_base", - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - "//tensorflow/core:testlib", - ] + if_tensorrt([ - "@local_config_cuda//cuda:cuda_headers", - "@local_config_tensorrt//:nv_infer", - ]), -) - -# Library for the segmenting portion of TensorRT operation creation -cc_library( - name = "segment", - srcs = ["segment/segment.cc"], - hdrs = [ - "segment/segment.h", - "segment/union_find.h", - ], - deps = [ - "//tensorflow/core:graph", - "//tensorflow/core:lib_proto_parsing", - "//tensorflow/core:protos_all_cc", - "@protobuf_archive//:protobuf_headers", - ], -) - -tf_cc_test( - name = "segment_test", - size = "small", - srcs = ["segment/segment_test.cc"], - tags = [ - "no_windows", - "nomac", - ], - deps = [ - ":segment", - "//tensorflow/cc:cc_ops", - "//tensorflow/cc:scope", - "//tensorflow/core:core_cpu", - "//tensorflow/core:lib", - "//tensorflow/core:ops", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - "//tensorflow/core:testlib", - ], -) - -# Library for the plugin factory -tf_cuda_library( - name = "trt_plugins", - srcs = [ - "plugin/trt_plugin.cc", - "plugin/trt_plugin_factory.cc", - "plugin/trt_plugin_utils.cc", - ], - hdrs = [ - "plugin/trt_plugin.h", - "plugin/trt_plugin_factory.h", - "plugin/trt_plugin_utils.h", - ], - deps = [ - "//tensorflow/core:framework_lite", - "//tensorflow/core:lib_proto_parsing", - ] + if_tensorrt([ - "@local_config_tensorrt//:nv_infer", - ]), -) - -tf_cuda_cc_test( - name = "trt_plugin_factory_test", - size = "small", - srcs = ["plugin/trt_plugin_factory_test.cc"], - tags = [ - "no_cuda_on_cpu_tap", - "no_windows", - "nomac", - ], - deps = [ - ":trt_plugins", - "//tensorflow/core:lib", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - ] + if_tensorrt([ - "@local_config_cuda//cuda:cuda_headers", - "@local_config_tensorrt//:nv_infer", - ]), -) - -py_library( - name = "tf_trt_integration_test_base", - srcs = ["test/tf_trt_integration_test_base.py"], - deps = [ - ":trt_convert_py", - ":trt_ops_py", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_test_lib", - ], -) - -cuda_py_test( - name = "trt_convert_test", - srcs = ["python/trt_convert_test.py"], - additional_deps = [ - ":trt_convert_py", - ":trt_ops_py", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:graph_util", - "//tensorflow/python/saved_model:builder", - "//tensorflow/python/saved_model:loader", - "//tensorflow/python/saved_model:signature_constants", - "//tensorflow/python/saved_model:signature_def_utils", - "//tensorflow/python/saved_model:tag_constants", - "//tensorflow/python/saved_model:utils", - "//tensorflow/python/tools:freeze_graph_lib", - "//tensorflow/python/tools:saved_model_utils", - ], - tags = [ - "no_cuda_on_cpu_tap", - "no_windows", - "nomac", - ], -) - -cuda_py_tests( - name = "tf_trt_integration_test", - srcs = [ - "test/base_test.py", - "test/batch_matmul_test.py", - "test/biasadd_matmul_test.py", - "test/binary_tensor_weight_broadcast_test.py", - "test/concatenation_test.py", - "test/const_broadcast_test.py", - "test/manual_test.py", - "test/memory_alignment_test.py", - "test/multi_connection_neighbor_engine_test.py", - "test/neighboring_engine_test.py", - "test/quantization_test.py", - "test/rank_two_test.py", - "test/reshape_transpose_test.py", - "test/vgg_block_nchw_test.py", - "test/vgg_block_test.py", - ], - additional_deps = [ - ":tf_trt_integration_test_base", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_test_lib", - ], - tags = [ - "no_cuda_on_cpu_tap", - "no_windows", - "nomac", - ], -) - -cuda_py_tests( - name = "tf_trt_integration_test_no_oss", - srcs = [ - "test/unary_test.py", - ], - additional_deps = [ - ":tf_trt_integration_test_base", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_test_lib", - ], - tags = [ - "no_cuda_on_cpu_tap", - "no_oss", # TODO(b/117274186): re-enable in OSS after crash fixed - "no_pip", # TODO(b/117274186): re-enable in OSS after crash fixed - "no_windows", - "nomac", - ], + actual = "//tensorflow/compiler/tf2tensorrt:trt_conversion", ) -cuda_py_test( - name = "quantization_mnist_test", - srcs = ["test/quantization_mnist_test.py"], - additional_deps = [ - ":tf_trt_integration_test_base", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python/keras:keras", - "//tensorflow/python/estimator:estimator", - ], - data = [ - "test/testdata/checkpoint", - "test/testdata/model.ckpt-46900.data-00000-of-00001", - "test/testdata/model.ckpt-46900.index", - ], - tags = [ - "no_cuda_on_cpu_tap", - "no_oss", # TODO(b/121194394): re-enable in OSS after OOM is fixed. - "no_pip", - "no_tap", # It is not able to download the mnist data. - "no_windows", - "nomac", - ], +alias( + name = "trt_op_kernels", + actual = "//tensorflow/compiler/tf2tensorrt:trt_op_kernels", ) -cc_library( - name = "utils", - srcs = ["convert/utils.cc"], - hdrs = ["convert/utils.h"], - copts = tf_copts(), - deps = [ - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "test_utils", - srcs = ["test/utils.cc"], - hdrs = ["test/utils.h"], - deps = [ - "//tensorflow/core:lib", - "@com_googlesource_code_re2//:re2", - ], +alias( + name = "trt_engine_op_op_lib", + actual = "//tensorflow/compiler/tf2tensorrt:trt_engine_op_op_lib", ) diff --git a/tensorflow/contrib/tensorrt/__init__.py b/tensorflow/contrib/tensorrt/__init__.py index 140ad4828208ae4844a49bf664955b50cd9e51cd..fd551d70b4385b14b84b7b98a6d16b0c03733d38 100644 --- a/tensorflow/contrib/tensorrt/__init__.py +++ b/tensorflow/contrib/tensorrt/__init__.py @@ -18,18 +18,6 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.python.framework import errors - -# pylint: disable=unused-import,wildcard-import,g-import-not-at-top -try: - from tensorflow.contrib.tensorrt.python import * -except errors.NotFoundError as e: - no_trt_message = ( - '**** Failed to initialize TensorRT. This is either because the TensorRT' - ' installation path is not in LD_LIBRARY_PATH, or because you do not have' - ' it installed. If not installed, please go to' - ' https://developer.nvidia.com/tensorrt to download and install' - ' TensorRT ****') - print(no_trt_message) - raise e -# pylint: enable=unused-import,wildcard-import,g-import-not-at-top +# pylint: disable=unused-import,wildcard-import +from tensorflow.contrib.tensorrt.python import * +# pylint: enable=unused-import,wildcard-import diff --git a/tensorflow/contrib/tensorrt/custom_plugin_examples/BUILD b/tensorflow/contrib/tensorrt/custom_plugin_examples/BUILD index 69058c5826822c519a69d50860c06b8ab3ec6578..0a2cf105baf5efb62d0c535c1f2d081973ec0ea3 100644 --- a/tensorflow/contrib/tensorrt/custom_plugin_examples/BUILD +++ b/tensorflow/contrib/tensorrt/custom_plugin_examples/BUILD @@ -45,10 +45,10 @@ tf_custom_op_library( "inc_op_kernel.cu.cc", ], deps = [ - "//tensorflow/contrib/tensorrt:trt_plugins", + "//tensorflow/compiler/tf2tensorrt:trt_plugins", "//tensorflow/core:framework_lite", ] + if_tensorrt([ - "@local_config_tensorrt//:nv_infer", + "@local_config_tensorrt//:tensorrt", ]), ) @@ -64,10 +64,10 @@ tf_kernel_library( "inc_op_kernel.cu.cc", ], deps = [ - "//tensorflow/contrib/tensorrt:trt_plugins", + "//tensorflow/compiler/tf2tensorrt:trt_plugins", "//tensorflow/core:stream_executor_headers_lib", ] + if_tensorrt([ - "@local_config_tensorrt//:nv_infer", + "@local_config_tensorrt//:tensorrt", ]) + tf_custom_op_library_additional_deps(), ) diff --git a/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_plugin.cc b/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_plugin.cc index 8d4c893af56689185da72398919e2241d451594b..7c9075142a02546ddd580e861ac87cb86badd739 100644 --- a/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_plugin.cc +++ b/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_plugin.cc @@ -15,8 +15,8 @@ limitations under the License. #include "tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_plugin.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h" #include "tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.h" -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.h" #if GOOGLE_CUDA #if GOOGLE_TENSORRT diff --git a/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_plugin.h b/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_plugin.h index 189e9c939b9ffd4450f7ba95fe1abdbbc049b430..fb048d7b19da0f010ed918b147013b20d37ed0dd 100644 --- a/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_plugin.h +++ b/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_plugin.h @@ -19,7 +19,7 @@ limitations under the License. #include #include -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin.h" #if GOOGLE_CUDA #if GOOGLE_TENSORRT diff --git a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h deleted file mode 100644 index b545f497f32d5a1a6960b748467ca189b7debf6c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h +++ /dev/null @@ -1,145 +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_CONTRIB_TENSORRT_KERNELS_TRT_ENGINE_OP_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_KERNELS_TRT_ENGINE_OP_H_ - -#include -#include - -#include "tensorflow/contrib/tensorrt/convert/utils.h" -#include "tensorflow/contrib/tensorrt/log/trt_logger.h" -#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h" -#include "tensorflow/core/framework/function.h" -#include "tensorflow/core/framework/graph.pb.h" -#include "tensorflow/core/framework/op.h" -#include "tensorflow/core/framework/op_kernel.h" -#include "tensorflow/core/platform/mutex.h" - -#if GOOGLE_CUDA -#if GOOGLE_TENSORRT -#include "cuda/include/cuda_runtime_api.h" -#include "tensorrt/include/NvInfer.h" - -namespace tensorflow { -namespace tensorrt { -struct TRTInt8Calibrator; -class TRTCalibrationResource; -class AsyncHelper; -// TODO(Sami): Remove this file? - -// This OP can construct TRTEngine on the fly and if construction of engine -// fails, executes equivalent subgraph as a TensorFlow function. -class TRTEngineOp : public AsyncOpKernel { - public: - explicit TRTEngineOp(OpKernelConstruction* context); - - void ComputeAsync(OpKernelContext* context, - AsyncOpKernel::DoneCallback done) override; - ~TRTEngineOp(); - - private: - // Execute calibration - void ExecuteCalibration(OpKernelContext* ctx, AsyncHelper* helper); - - // Construct a function handle for executing native funcdef graph - Status ConstructFunctionHandle(OpKernelContext* ctx); - - // Execute replaced native segment as function Op. - void ExecuteNativeSegment(OpKernelContext* ctx, AsyncHelper* helper); - - // Execute the tensorrt engine. Returns whether we need to retry by running - // the native segment. - bool ExecuteTrtEngine(OpKernelContext* ctx, const int num_batch, - nvinfer1::ICudaEngine* trt_engine_ptr, - nvinfer1::IExecutionContext* trt_execution_context_ptr); - - // Allocate necessary resources for calibration - Status AllocateCalibrationResources(OpKernelContext* ctx, - TRTCalibrationResource** cr); - - // TODO(samikama): context should go to a resource manager! - typedef std::pair, - TrtUniquePtrType> - EngineCtxPair; - EngineCtxPair& GetEngine(int batch_size, OpKernelContext* ctx); - - // Return engine batch closest to input batch. - int GetEngineBatch(OpKernelContext* ctx); - - nvinfer1::IGpuAllocator* GetAllocator(OpKernelContext* ctx); - - // map to keep engines and their execution context for given batch size. - std::unordered_map engine_map_; - std::vector input_nodes_; - std::vector output_nodes_; - - // keep device allocator for TRT. - std::unique_ptr allocator_; - - // serialized protobuf segment or trt engine depending on static_engine_ flag. - string serialized_segment_; - - // Name of the function for TF native execution of the segment. - string funcdef_name_; - - // GraphDef representation of the segment. - GraphDef segment_graph_; - - // Lookup table for temporary staging areas of input tensors for calibration. - std::unordered_map> device_buffers_; - - // Temporary staging areas for calibration inputs. - std::vector dev_tensors_; - - // Engine Precision mode. - int precision_mode_; - - // Whether engine is constructed during the conversion or needs to be - // constructed from protobuf segment. - bool static_engine_; - - // Whether to calibrate INT8 engine. - bool calibration_mode_; - - // Whether non-batch ranks of the inputs are assumed to be fixed or not for - // engine construction. - bool fixed_input_size_; - - // Batches of the cached engines - std::vector cached_engine_batches_; - - // Maximum number of cached engines - int max_cached_engines_; - - int64 workspace_size_; - mutex engine_mutex_; - FunctionLibraryRuntime::Handle native_func_; - - // The finalized calibrator for inference. - std::unique_ptr calibrator_; - - // If true, create calibration graph for INT8 mode. Otherwise, we are using - // user-provided quantization ranges. - bool use_calibration_; -}; - -} // namespace tensorrt -} // namespace tensorflow - -#endif // GOOGLE_TENSORRT -#endif // GOOGLE_CUDA - -#endif // TENSORFLOW_CONTRIB_TENSORRT_KERNELS_TRT_ENGINE_OP_H_ diff --git a/tensorflow/contrib/tensorrt/python/__init__.py b/tensorflow/contrib/tensorrt/python/__init__.py index 7cdfe2b1a612be2eec473d806d0eb44b611ca68a..0cae401023e7d3e3780b9dd2e2a92c9fd0e92db8 100644 --- a/tensorflow/contrib/tensorrt/python/__init__.py +++ b/tensorflow/contrib/tensorrt/python/__init__.py @@ -19,12 +19,6 @@ from __future__ import division from __future__ import print_function # pylint: disable=unused-import,line-too-long -from tensorflow.contrib.tensorrt.python.ops import trt_engine_op -from tensorflow.contrib.tensorrt.python.trt_convert import add_test_value from tensorflow.contrib.tensorrt.python.trt_convert import calib_graph_to_infer_graph -from tensorflow.contrib.tensorrt.python.trt_convert import clear_test_values from tensorflow.contrib.tensorrt.python.trt_convert import create_inference_graph -from tensorflow.contrib.tensorrt.python.trt_convert import enable_test_value -from tensorflow.contrib.tensorrt.python.trt_convert import get_test_value -from tensorflow.contrib.tensorrt.python.trt_convert import is_tensorrt_enabled # pylint: enable=unused-import,line-too-long diff --git a/tensorflow/contrib/tensorrt/python/trt_convert.py b/tensorflow/contrib/tensorrt/python/trt_convert.py index 203b2697babe32b45523109708cbf062dceee33b..4a959378138dec6f1c1a3f490704d7aebeae9b47 100644 --- a/tensorflow/contrib/tensorrt/python/trt_convert.py +++ b/tensorflow/contrib/tensorrt/python/trt_convert.py @@ -18,404 +18,41 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import six as _six -# pylint: disable=unused-import,line-too-long -from tensorflow.contrib.tensorrt.wrap_conversion import add_test_value -from tensorflow.contrib.tensorrt.wrap_conversion import calib_convert -from tensorflow.contrib.tensorrt.wrap_conversion import clear_test_values -from tensorflow.contrib.tensorrt.wrap_conversion import enable_test_value -from tensorflow.contrib.tensorrt.wrap_conversion import get_linked_tensorrt_version -from tensorflow.contrib.tensorrt.wrap_conversion import get_loaded_tensorrt_version -from tensorflow.contrib.tensorrt.wrap_conversion import get_test_value -from tensorflow.contrib.tensorrt.wrap_conversion import is_tensorrt_enabled -# pylint: enable=unused-import,line-too-long -from tensorflow.core.framework import graph_pb2 -from tensorflow.core.protobuf import config_pb2 -from tensorflow.core.protobuf import meta_graph_pb2 -from tensorflow.core.protobuf import rewriter_config_pb2 -from tensorflow.python.client import session -from tensorflow.python.framework import errors_impl as _impl -from tensorflow.python.framework import graph_util -from tensorflow.python.framework import importer -from tensorflow.python.framework import ops -from tensorflow.python.grappler import tf_optimizer -from tensorflow.python.platform import tf_logging -from tensorflow.python.saved_model import builder -from tensorflow.python.saved_model import loader_impl -from tensorflow.python.saved_model import tag_constants -from tensorflow.python.training import saver - -if _six.PY2: - _to_bytes = lambda s: s - _to_string = lambda s: s -else: - _to_bytes = lambda s: s.encode("utf-8", errors="surrogateescape") - _to_string = lambda s: s.decode("utf-8") - - -class TrtPrecisionMode(object): - FP32 = "FP32" - FP16 = "FP16" - INT8 = "INT8" - - @staticmethod - def supported_precision_modes(): - return [TrtPrecisionMode.FP32, TrtPrecisionMode.FP16, TrtPrecisionMode.INT8] - - -def get_tensorrt_rewriter_config(rewriter_config=None, - max_batch_size=1, - max_workspace_size_bytes=2 << 20, - precision_mode=TrtPrecisionMode.FP32, - minimum_segment_size=3, - is_dynamic_op=False, - maximum_cached_engines=1, - cached_engine_batch_sizes=None, - use_calibration=True): - """Returns a RewriterConfig proto for TRT transformation. - - Args: - rewriter_config: a template RewriterConfig proto used to create a - TRT-enabled RewriterConfig. If None, it will use a default one. - max_batch_size: max size for the input batch - max_workspace_size_bytes: the maximum GPU temporary memory which the TRT - engine can use at execution time. This corresponds to the 'workspaceSize' - parameter of nvinfer1::IBuilder::setMaxWorkspaceSize(). - precision_mode: one of TrtPrecisionMode.supported_precision_modes(). - minimum_segment_size: the minimum number of nodes required for a subgraph to - be replaced by TRTEngineOp. - is_dynamic_op: whether to generate dynamic TRT ops which will build the TRT - network and engine at run time. - maximum_cached_engines: max number of cached TRT engines in dynamic TRT ops. - If the number of cached engines is already at max but none of them can - serve the input, the TRTEngineOp will fall back to run the TF function - based on which the TRTEngineOp is created. - cached_engine_batch_sizes: a list of batch sizes used to create cached - engines, only used when is_dynamic_op is True. The length of the list - should be smaller than maximum_cached_engines, and the dynamic TRT op will - use this list to determine the batch sizes of the cached engines, instead - of making the decision on the fly. This is useful when we know the most - common batch size(s) the application is going to generate. - use_calibration: this argument is ignored if precision_mode is not INT8. If - set to True, a calibration graph will be created to calibrate the missing - ranges. The calibration graph must be converted to an inference graph - using calib_graph_to_infer_graph() after running calibration. if set to - False, quantization nodes will be expected for every tensor in the graph - (exlcuding those which will be fused). If a range is missing, an error - will occur. Please note that accuracy may be negatively affected if there - is a mismatch between which tensors TRT quantizes and which tensors were - trained with fake quantization. - - Returns: - A RewriterConfig proto which sets a TensorRTOptimizer to run Grappler. - - Raises: - TypeError: if any of the parameters are of unexpected type. - ValueError: if any of the parameters are of unexpected value. - """ - if rewriter_config is not None and not isinstance( - rewriter_config, rewriter_config_pb2.RewriterConfig): - raise TypeError("rewriter_config should be a RewriterConfig proto.") - - rewriter_config_with_trt = rewriter_config_pb2.RewriterConfig() - if rewriter_config is None: - # Layout optimizer may add Const nodes followed by Reshape nodes, thus we - # need to run constant folding again. - rewriter_config_with_trt.optimizers.extend( - ["constfold", "layout", "constfold"]) - rewriter_config_with_trt.meta_optimizer_iterations = ( - rewriter_config_pb2.RewriterConfig.ONE) - else: - rewriter_config_with_trt.CopyFrom(rewriter_config) - - if precision_mode.upper() not in TrtPrecisionMode.supported_precision_modes(): - raise ValueError(("precision mode '{}' is not supported." - "It should be one of {}").format( - precision_mode, - TrtPrecisionMode.supported_precision_modes)) - - optimizer = rewriter_config_with_trt.custom_optimizers.add() - optimizer.name = "TensorRTOptimizer" - optimizer.parameter_map["minimum_segment_size"].i = minimum_segment_size - optimizer.parameter_map["max_batch_size"].i = max_batch_size - optimizer.parameter_map["is_dynamic_op"].b = is_dynamic_op - optimizer.parameter_map[ - "max_workspace_size_bytes"].i = max_workspace_size_bytes - optimizer.parameter_map["precision_mode"].s = _to_bytes(precision_mode) - optimizer.parameter_map["maximum_cached_engines"].i = maximum_cached_engines - if cached_engine_batch_sizes: - if not isinstance(cached_engine_batch_sizes, list): - raise TypeError("cached_engine_batch_sizes should be a list.") - if len(cached_engine_batch_sizes) > maximum_cached_engines: - raise ValueError("cached_engine_batch_sizes should not contain more than " - "maximum_cached_engines items.") - optimizer.parameter_map["cached_engine_batches"].list.i.extend( - cached_engine_batch_sizes) - optimizer.parameter_map["use_calibration"].b = use_calibration - return rewriter_config_with_trt - - -def create_inference_graph(input_graph_def, - outputs, - max_batch_size=1, - max_workspace_size_bytes=2 << 20, - precision_mode=TrtPrecisionMode.FP32, - minimum_segment_size=3, - is_dynamic_op=False, - maximum_cached_engines=1, - cached_engine_batch_sizes=None, - use_calibration=True, - input_saved_model_dir=None, - input_saved_model_tags=None, - output_saved_model_dir=None, - session_config=None): - """Python wrapper for the TRT transformation. - - Args: - input_graph_def: a GraphDef object containing a model to be transformed. If - set to None, the graph will be read from the SavedModel loaded from - input_saved_model_dir. - outputs: list of tensors or node names for the model outputs. Only used when - input_graph_def is not None. - max_batch_size: max size for the input batch. - max_workspace_size_bytes: the maximum GPU temporary memory which the TRT - engine can use at execution time. This corresponds to the 'workspaceSize' - parameter of nvinfer1::IBuilder::setMaxWorkspaceSize(). - precision_mode: one of TrtPrecisionMode.supported_precision_modes(). - minimum_segment_size: the minimum number of nodes required for a subgraph to - be replaced by TRTEngineOp. - is_dynamic_op: whether to generate dynamic TRT ops which will build the TRT - network and engine at run time. - maximum_cached_engines: max number of cached TRT engines in dynamic TRT ops. - If the number of cached engines is already at max but none of them can - serve the input, the TRTEngineOp will fall back to run the TF function - based on which the TRTEngineOp is created. - cached_engine_batch_sizes: a list of batch sizes used to create cached - engines, only used when is_dynamic_op is True. The length of the list - should be smaller than maximum_cached_engines, and the dynamic TRT op will - use this list to determine the batch sizes of the cached engines, instead - of making the decision on the fly. This is useful when we know the most - common batch size(s) the application is going to generate. - use_calibration: this argument is ignored if precision_mode is not INT8. If - set to True, a calibration graph will be created to calibrate the missing - ranges. The calibration graph must be converted to an inference graph - using calib_graph_to_infer_graph() after running calibration. if set to - False, quantization nodes will be expected for every tensor in the graph - (exlcuding those which will be fused). If a range is missing, an error - will occur. Please note that accuracy may be negatively affected if there - is a mismatch between which tensors TRT quantizes and which tensors were - trained with fake quantization. - input_saved_model_dir: the directory to load the SavedModel which contains - the input graph to transforms. Used only when input_graph_def is None. - input_saved_model_tags: list of tags to load the SavedModel. - output_saved_model_dir: if not None, construct a SavedModel using the - returned GraphDef and save it to the specified directory. This option only - works when the input graph is loaded from a SavedModel, i.e. when - input_saved_model_dir is specified and input_graph_def is None. - session_config: the ConfigProto used to create a Session. It's also used as - a template to create a TRT-enabled ConfigProto for conversion. If not - specified, a default ConfigProto will be used. - - Returns: - A GraphDef transformed from input_graph_def (or the SavedModel graph def - loaded from input_saved_model_dir, if input_graph_def is not present), where - all TRT compatible subgraphs are replaced with TRTEngineOps, and a TF - function is added for each of the subgraphs. - - If is_dynamic_op is True, each TRTEngineOp will contain a serialized - subgraph GraphDef, which will be converted to a TRT engine at execution time - and the TRT engine will be cached for future usage. A new TRT engine will be - created each time when none of the cached engines match the input shapes. If - it fails to execute the TRT engine or the number of cached engines reaches - maximum_cached_engines, the op will fall back to call the corresponding TF - function. - - If is_dynamic_op is False, each TRTEngineOp will contain a serialized TRT - engine created from the corresponding subgraph. No more engines will be - created on the fly, and the op will fall back to call the corresponding TF - function when it fails to execute the engine. - - Raises: - ValueError: if the combination of the parameters is invalid. - RuntimeError: if the TensorRT library version is incompatible. - """ - compiled_version = get_linked_tensorrt_version() - loaded_version = get_loaded_tensorrt_version() - version_mismatch = False - if loaded_version[0] < compiled_version[0]: - tf_logging.error( - "TensorRT version mismatch. Tensorflow was compiled against " + - "TensorRT %s but library loaded from environment is TensorRT %s" % - (".".join([str(x) for x in compiled_version]), - ".".join([str(x) for x in loaded_version])) + - ". Please make sure that correct version of TensorRT " + - "is available in the system and added to ldconfig or LD_LIBRARY_PATH") - raise RuntimeError("Incompatible TensorRT library version") - for i in zip(loaded_version, compiled_version): - if i[0] != i[1]: - tf_logging.warn("TensorRT mismatch. Compiled against version " + - "%s, but loaded %s. Things may not work" % - (".".join([str(x) for x in compiled_version]), - ".".join([str(x) for x in loaded_version]))) - version_mismatch = True - break - if not version_mismatch: - tf_logging.info("Running against TensorRT version %s" % ".".join( - [str(x) for x in loaded_version])) - - if session_config is None: - session_config = config_pb2.ConfigProto() - - if input_saved_model_tags is None: - input_saved_model_tags = [tag_constants.SERVING] - saved_model_loader = None - grappler_meta_graph_def = None - - if input_graph_def is None: - # Read from SavedModel and freeze the graph if necessary. - if input_saved_model_dir is None: - raise ValueError("input_graph_def and input_saved_model_dir cannot be " - "both None") - with ops.Graph().as_default(): - with session.Session(config=session_config) as sess: - saved_model_loader = loader_impl.SavedModelLoader(input_saved_model_dir) - input_meta_graph_def = saved_model_loader.load(sess, - input_saved_model_tags) - output_node_names = set() - - def _gather_names(tensor_info): - """Get the node names from a TensorInfo.""" - return set( - [tensor_info[key].name.split(":")[0] for key in tensor_info]) - - # Get input and outputs from all SignatureDef. - for key in input_meta_graph_def.signature_def: - signature_def = input_meta_graph_def.signature_def[key] - output_node_names.update(_gather_names(signature_def.inputs)) - output_node_names.update(_gather_names(signature_def.outputs)) - - # Freeze the variables in the SavedModel graph and copy the frozen - # graph over. - frozen_graph_def = graph_util.convert_variables_to_constants( - sess, sess.graph.as_graph_def(add_shapes=True), - list(output_node_names)) - grappler_meta_graph_def = meta_graph_pb2.MetaGraphDef() - grappler_meta_graph_def.graph_def.CopyFrom(frozen_graph_def) - - # Copy the collections that are not variables. - for key in input_meta_graph_def.collection_def: - # TODO(laigd): currently we use the collection key to filter out - # collections that depend on variable ops, but this may miss some - # other user-defined collections. A better way would be to use - # CollectionDef::NodeList for the filtering. - if key not in [ - "variables", "local_variables", "model_variables", - "trainable_variables", "train_op", "table_initializer" - ]: - grappler_meta_graph_def.collection_def[key].CopyFrom( - input_meta_graph_def.collection_def[key]) - - # Copy other information. - grappler_meta_graph_def.meta_info_def.CopyFrom( - input_meta_graph_def.meta_info_def) - for key in input_meta_graph_def.signature_def: - grappler_meta_graph_def.signature_def[key].CopyFrom( - input_meta_graph_def.signature_def[key]) - # TODO(laigd): maybe add back AssetFileDef. - else: - if output_saved_model_dir is not None: - raise ValueError("output_saved_model_dir cannot be set when " - "input_graph_def is set") - # Create MetaGraphDef from input graph. - graph = ops.Graph() - with graph.as_default(): - importer.import_graph_def(input_graph_def, name="") - grappler_meta_graph_def = saver.export_meta_graph( - graph_def=graph.as_graph_def(add_shapes=True), graph=graph) - if outputs: - output_collection = meta_graph_pb2.CollectionDef() - output_list = output_collection.node_list.value - for i in outputs: - if isinstance(i, ops.Tensor): - output_list.append(_to_bytes(i.name)) - else: - output_list.append(_to_bytes(i)) - # TODO(laigd): use another key as the outputs are really not train_op. - grappler_meta_graph_def.collection_def["train_op"].CopyFrom( - output_collection) - - # Create TRT-enabled ConfigProto. - session_config_with_trt = config_pb2.ConfigProto() - session_config_with_trt.CopyFrom(session_config) - rewriter_config = None - if (session_config_with_trt.HasField("graph_options") and - session_config_with_trt.graph_options.HasField("rewrite_options")): - rewriter_config = session_config_with_trt.graph_options.rewrite_options - rewriter_config_with_trt = get_tensorrt_rewriter_config( - rewriter_config, max_batch_size, max_workspace_size_bytes, precision_mode, - minimum_segment_size, is_dynamic_op, maximum_cached_engines, - cached_engine_batch_sizes, use_calibration) - session_config_with_trt.graph_options.rewrite_options.CopyFrom( - rewriter_config_with_trt) - - # Run Grappler. - transformed_graph_def = tf_optimizer.OptimizeGraph( - session_config_with_trt, grappler_meta_graph_def, graph_id=b"tf_graph") - - # Optionally write the transformed graphdef as SavedModel. - if output_saved_model_dir is not None: - saved_model_builder = builder.SavedModelBuilder(output_saved_model_dir) - with ops.Graph().as_default(): - importer.import_graph_def(transformed_graph_def, name="") - # We don't use TRT here. - with session.Session(config=session_config) as sess: - saved_model_builder.add_meta_graph_and_variables( - sess, - input_saved_model_tags, - signature_def_map=grappler_meta_graph_def.signature_def) - # Ignore other meta graphs from the input SavedModel. - saved_model_builder.save() - - return transformed_graph_def +from tensorflow.python.compiler.tensorrt import trt_convert + + +def create_inference_graph( + input_graph_def, + outputs, + max_batch_size=1, + max_workspace_size_bytes=trt_convert.DEFAULT_TRT_MAX_WORKSPACE_SIZE_BYTES, + precision_mode=trt_convert.TrtPrecisionMode.FP32, + minimum_segment_size=3, + is_dynamic_op=False, + maximum_cached_engines=1, + cached_engine_batches=None, + use_calibration=True, + input_saved_model_dir=None, + input_saved_model_tags=None, + output_saved_model_dir=None, + session_config=None): + return trt_convert.create_inference_graph( + input_graph_def=input_graph_def, + outputs=outputs, + max_batch_size=max_batch_size, + max_workspace_size_bytes=max_workspace_size_bytes, + precision_mode=precision_mode, + minimum_segment_size=minimum_segment_size, + is_dynamic_op=is_dynamic_op, + maximum_cached_engines=maximum_cached_engines, + cached_engine_batches=cached_engine_batches, + use_calibration=use_calibration, + input_saved_model_dir=input_saved_model_dir, + input_saved_model_tags=input_saved_model_tags, + output_saved_model_dir=output_saved_model_dir, + session_config=session_config) def calib_graph_to_infer_graph(calibration_graph_def, is_dynamic_op=False): - """Convert an existing calibration graph to inference graph. - - Args: - calibration_graph_def: the calibration GraphDef object with calibration data - is_dynamic_op: whether to create dynamic static engines from calibration - - Returns: - New GraphDef with TRTEngineOps placed in graph replacing calibration nodes. - Raises: - RuntimeError: if the returned status message is malformed. - """ - - is_calib_graph = False - for n in calibration_graph_def.node: - if n.op == "TRTEngineOp": - is_calib_graph = is_calib_graph or not n.attr["calibration_data"].s - if not is_calib_graph: - tf_logging.error( - "Not a calib graph. Doesn't seem to contain any calibration nodes.") - return None - graph_str = calibration_graph_def.SerializeToString() - out = calib_convert(graph_str, is_dynamic_op) - status = _to_string(out[0]) - output_graph_def_string = out[1] - del graph_str # Save some memory - if len(status) < 2: - raise _impl.UnknownError(None, None, status) - if status[:2] != "OK": - msg = status.split(";") - if len(msg) == 1: - raise RuntimeError("Status message is malformed {}".format(status)) - # pylint: disable=protected-access - raise _impl._make_specific_exception(None, None, ";".join(msg[1:]), - int(msg[0])) - # pylint: enable=protected-access - output_graph_def = graph_pb2.GraphDef() - output_graph_def.ParseFromString(output_graph_def_string) - del output_graph_def_string # Save some memory - return output_graph_def + return trt_convert.calib_graph_to_infer_graph( + calibration_graph_def=calibration_graph_def, is_dynamic_op=is_dynamic_op) diff --git a/tensorflow/contrib/tensorrt/resources/trt_resources.h b/tensorflow/contrib/tensorrt/resources/trt_resources.h deleted file mode 100644 index aac9e5c7bd725fc10bcaa04536ebc7be071b4d4c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/resources/trt_resources.h +++ /dev/null @@ -1,79 +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_CONTRIB_TENSORRT_RESOURCES_TRT_RESOURCES_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_RESOURCES_H_ - -#include -#include -#include -#include -#include - -#include "tensorflow/contrib/tensorrt/convert/utils.h" -#include "tensorflow/contrib/tensorrt/log/trt_logger.h" -#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h" -#include "tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h" -#include "tensorflow/core/framework/resource_mgr.h" - -#if GOOGLE_CUDA -#if GOOGLE_TENSORRT - -#include "tensorrt/include/NvInfer.h" - -namespace tensorflow { -namespace tensorrt { - -class TRTCalibrationResource : public tensorflow::ResourceBase { - public: - ~TRTCalibrationResource() { - LOG(INFO) << "Destroying Calibration Resource " << std::endl - << DebugString(); - builder_.reset(); - engine_.reset(); - // We need to manually destroy the builder and engine before the allocator - // is destroyed. - allocator_.reset(); - } - - string DebugString() override { - std::stringstream oss; - using std::dec; - using std::endl; - using std::hex; - oss << " Calibrator = " << hex << calibrator_.get() << dec << endl - << " Builder = " << hex << builder_.get() << dec << endl - << " Engine = " << hex << engine_.get() << dec << endl - << " Logger = " << hex << &logger_ << dec << endl - << " Allocator = " << hex << allocator_.get() << dec << endl - << " Thread = " << hex << thr_.get() << dec << endl; - return oss.str(); - } - - std::unique_ptr calibrator_; - TrtUniquePtrType builder_; - TrtUniquePtrType engine_; - std::unique_ptr allocator_; - tensorflow::tensorrt::Logger logger_; - // TODO(sami): Use threadpool threads! - std::unique_ptr thr_; -}; - -} // namespace tensorrt -} // namespace tensorflow - -#endif -#endif -#endif // TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_RESOURCES_H_ diff --git a/tensorflow/contrib/tensorrt/shape_fn/trt_shfn.cc b/tensorflow/contrib/tensorrt/shape_fn/trt_shfn.cc index f30dba59ad55317d7ad7730e4dc66c9aba4e6a6b..5c60d6b589ed6a16276226726d989e949bcbf9d7 100644 --- a/tensorflow/contrib/tensorrt/shape_fn/trt_shfn.cc +++ b/tensorflow/contrib/tensorrt/shape_fn/trt_shfn.cc @@ -14,14 +14,14 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/contrib/tensorrt/shape_fn/trt_shfn.h" -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.h" #include #include #if GOOGLE_CUDA #if GOOGLE_TENSORRT -#include "tensorflow/contrib/tensorrt/log/trt_logger.h" +#include "tensorflow/compiler/tf2tensorrt/plugin/trt_plugin_factory.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorrt/include/NvInfer.h" diff --git a/tensorflow/contrib/tensorrt/test/manual_test.py b/tensorflow/contrib/tensorrt/test/manual_test.py deleted file mode 100644 index 1187c759b4b5483cbf5afe136401abe86d6ef989..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/test/manual_test.py +++ /dev/null @@ -1,114 +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. -# ============================================================================== -"""Basic tests for TF-TensorRT integration.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import ast -import os - -from tensorflow.contrib.tensorrt.test import tf_trt_integration_test_base as trt_test -from tensorflow.core.framework import graph_pb2 -from tensorflow.python.platform import gfile -from tensorflow.python.platform import test - - -class ManualTest(trt_test.TfTrtIntegrationTestBase): - - def __init__(self, methodName='runTest'): # pylint: disable=invalid-name - super(ManualTest, self).__init__(methodName) - self._params_map = None - - def _GetEnv(self): - """Get an environment variable specifying the manual test parameters. - - The value of the environment variable is the string representation of a dict - which should contain the following keys: - - 'graph_path': the file path to the serialized frozen graphdef - - 'input_names': TfTrtIntegrationTestParams.input_names - - 'input_dims': TfTrtIntegrationTestParams.input_dims - - 'expected_output_dims': TfTrtIntegrationTestParams.expected_output_dims - - 'output_name': the name of op to fetch - - 'expected_engines_to_run': ExpectedEnginesToRun() will return this - - 'expected_engines_to_build': ExpectedEnginesToBuild() will return this - - 'max_batch_size': ConversionParams.max_batch_size - - Returns: - The value of the environment variable. - """ - return os.getenv('TRT_MANUAL_TEST_PARAMS', '') - - def _GetParamsMap(self): - """Parse the environment variable as a dict and return it.""" - if self._params_map is None: - self._params_map = ast.literal_eval(self._GetEnv()) - return self._params_map - - def GetParams(self): - """Testing conversion of manually provided frozen graph.""" - params_map = self._GetParamsMap() - gdef = graph_pb2.GraphDef() - with gfile.Open(params_map['graph_path'], 'rb') as f: - gdef.ParseFromString(f.read()) - return trt_test.TfTrtIntegrationTestParams( - gdef=gdef, - input_names=params_map['input_names'], - input_dims=params_map['input_dims'], - output_names=params_map['output_names'], - expected_output_dims=params_map['expected_output_dims']) - - def GetConversionParams(self, run_params): - """Return a ConversionParams for test.""" - conversion_params = super(ManualTest, self).GetConversionParams(run_params) - params_map = self._GetParamsMap() - if 'max_batch_size' in params_map: - conversion_params = conversion_params._replace( - max_batch_size=params_map['max_batch_size']) - return conversion_params - - def ExpectedEnginesToBuild(self, run_params): - """Return the expected engines to build.""" - return self._GetParamsMap()['expected_engines_to_build'] - - def ExpectedEnginesToRun(self, run_params): - """Return the expected engines to run.""" - params_map = self._GetParamsMap() - if 'expected_engines_to_run' in params_map: - return params_map['expected_engines_to_run'] - return self.ExpectedEnginesToBuild(run_params) - - def ExpectedAbsoluteTolerance(self, run_params): - """The absolute tolerance to compare floating point results.""" - params_map = self._GetParamsMap() - if 'atol' in params_map: - return params_map['atol'] - return 1.e-3 - - def ExpectedRelativeTolerance(self, run_params): - """The relative tolerance to compare floating point results.""" - params_map = self._GetParamsMap() - if 'rtol' in params_map: - return params_map['rtol'] - return 1.e-3 - - def ShouldRunTest(self, run_params): - """Whether to run the test.""" - return len(self._GetEnv()) - - -if __name__ == '__main__': - test.main() diff --git a/tensorflow/contrib/tensorrt/test/test_tftrt.py b/tensorflow/contrib/tensorrt/test/test_tftrt.py deleted file mode 100644 index d26f26008635733c6c364a98b72b88c1e552f5fe..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/test/test_tftrt.py +++ /dev/null @@ -1,287 +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. -# ============================================================================== -"""Script to test TF-TensorRT integration.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import argparse -import numpy as np -import six as _six - -# normally we should do import tensorflow as tf and then -# tf.placeholder, tf.constant, tf.nn.conv2d etc but -# it looks like internal builds don't like it so -# importing every module individually - -from tensorflow.contrib import tensorrt as trt -from tensorflow.core.protobuf import config_pb2 as cpb2 -from tensorflow.core.protobuf import rewriter_config_pb2 as rwpb2 -from tensorflow.python.client import session as csess -from tensorflow.python.framework import constant_op as cop -from tensorflow.python.framework import dtypes as dtypes -from tensorflow.python.framework import importer as importer -from tensorflow.python.framework import ops as ops -from tensorflow.python.ops import array_ops as aops -from tensorflow.python.ops import math_ops as mops -from tensorflow.python.ops import nn as nn -from tensorflow.python.ops import nn_ops as nn_ops - - -def py2bytes(inp): - return inp - - -def py3bytes(inp): - return inp.encode("utf-8", errors="surrogateescape") - - -def py2string(inp): - return inp - - -def py3string(inp): - return inp.decode("utf-8") - - -if _six.PY2: - to_bytes = py2bytes - to_string = py2string -else: - to_bytes = py3bytes - to_string = py3string - - -def get_multi_engine_graph_def(mode="FP32"): - """Create a simple graph and return its graph_def.""" - dtype = dtypes.float32 - if mode.upper() == "FP16": - dtype = dtypes.float16 - else: - pass - - g = ops.Graph() - with g.as_default(): - x = aops.placeholder(shape=[None, 3, 7, 5], name="input", dtype=dtype) - with g.name_scope("Global_scope"): - with g.name_scope("first_scope"): - e = cop.constant( - np.random.randn(3, 2, 3, 4), name="weights", dtype=dtype) - conv = nn.conv2d( - input=x, - filter=e, - data_format="NCHW", - strides=[1, 1, 1, 1], - padding="VALID", - name="conv") - b = cop.constant(np.random.randn(1, 4, 1, 1), name="bias1", dtype=dtype) - t = conv * b - - b = cop.constant(np.random.randn(1, 4, 1, 1), name="bias2", dtype=dtype) - q = conv / b - edge = mops.sin(q) - edge1 = mops.cos(conv) - with g.name_scope("test_scope"): - de = edge + edge1 - t -= edge1 - q *= edge - t += q - t -= de - k = aops.squeeze(t, name="output") - print(k.dtype) - return g.as_graph_def() - - -def get_simple_graph_def(): - """Create a simple graph and return its graph_def.""" - g = ops.Graph() - with g.as_default(): - a = aops.placeholder( - dtype=dtypes.float32, shape=(None, 24, 24, 2), name="input") - e = cop.constant( - [[[[1., 0.5, 4., 6., 0.5, 1.], [1., 0.5, 1., 1., 0.5, 1.]]]], - name="weights", - dtype=dtypes.float32) - conv = nn.conv2d( - input=a, filter=e, strides=[1, 2, 2, 1], padding="SAME", name="conv") - b = cop.constant( - [4., 1.5, 2., 3., 5., 7.], name="bias", dtype=dtypes.float32) - t = nn.bias_add(conv, b, name="biasAdd") - relu = nn.relu(t, "relu") - idty = aops.identity(relu, "ID") - v = nn_ops.max_pool( - idty, [1, 2, 2, 1], [1, 2, 2, 1], "VALID", name="max_pool") - aops.squeeze(v, name="output") - return g.as_graph_def() - - -def execute_graph(gdef, dumm_inp): - """Run given graphdef once.""" - print("executing") - gpu_options = None - if trt.trt_convert.get_linked_tensorrt_version()[0] == 3: - gpu_options = cpb2.GPUOptions(per_process_gpu_memory_fraction=0.50) - sessconfig = cpb2.ConfigProto(gpu_options=gpu_options) - ops.reset_default_graph() - g = ops.Graph() - with g.as_default(): - inp, out = importer.import_graph_def( - graph_def=gdef, return_elements=["input", "output"]) - inp = inp.outputs[0] - out = out.outputs[0] - with csess.Session(config=sessconfig, graph=g) as sess: - val = sess.run(out, {inp: dumm_inp}) - return val - - -# Use real data that is representative of the inference dataset -# for calibration. For this test script it is random data. -def execute_calibration(gdef, dumm_inp): - """Run given calibration graph multiple times.""" - gpu_options = None - if trt.trt_convert.get_linked_tensorrt_version()[0] == 3: - gpu_options = cpb2.GPUOptions(per_process_gpu_memory_fraction=0.50) - ops.reset_default_graph() - g = ops.Graph() - with g.as_default(): - inp, out = importer.import_graph_def( - graph_def=gdef, return_elements=["input", "output"]) - inp = inp.outputs[0] - out = out.outputs[0] - with csess.Session( - config=cpb2.ConfigProto(gpu_options=gpu_options), graph=g) as sess: - # run over real calibration data here, we are mimicking a calibration set of - # 30 different batches. Use as much calibration data as you want - for _ in range(30): - val = sess.run(out, {inp: dumm_inp}) - return val - - -def user(multi_engine, - run_graph=execute_graph, - run_calibration=execute_calibration): - """Example function that converts a graph to TFTRT graph.""" - if multi_engine: - inp_dims = (2, 3, 7, 5) - orig_graph = get_multi_engine_graph_def() - else: - inp_dims = (100, 24, 24, 2) - orig_graph = get_simple_graph_def() # use a frozen graph for inference - dummy_input = np.random.random_sample(inp_dims) - # Get optimized graph - trt_graph = trt.create_inference_graph( - input_graph_def=orig_graph, - outputs=["output"], - max_batch_size=inp_dims[0], - max_workspace_size_bytes=1 << 25, - precision_mode="FP32", # TRT Engine precision "FP32","FP16" or "INT8" - minimum_segment_size=2, # minimum number of nodes in an engine - is_dynamic_op=False, - maximum_cached_engines=1, - cached_engine_batch_sizes=[]) - o1 = run_graph(orig_graph, dummy_input) - o2 = run_graph(trt_graph, dummy_input) - o3 = run_graph(trt_graph, dummy_input) - assert np.array_equal(o1, o2) - assert np.array_equal(o3, o2) # sanity check - fp16_graph = trt.create_inference_graph( - input_graph_def=orig_graph, - outputs=["output"], - max_batch_size=inp_dims[0], - max_workspace_size_bytes=1 << 25, - precision_mode="FP16", # TRT Engine precision "FP32","FP16" or "INT8" - minimum_segment_size=2, # minimum number of nodes in an engine - is_dynamic_op=False, - maximum_cached_engines=1, - cached_engine_batch_sizes=[]) - int8_calib_gdef = trt.create_inference_graph( - input_graph_def=orig_graph, - outputs=["output"], - max_batch_size=inp_dims[0], - max_workspace_size_bytes=1 << 25, - precision_mode="INT8", # TRT Engine precision "FP32","FP16" or "INT8" - minimum_segment_size=2, # minimum number of nodes in an engine - is_dynamic_op=False, - maximum_cached_engines=1, - cached_engine_batch_sizes=[]) - o4 = run_graph(fp16_graph, dummy_input) - _ = run_calibration(int8_calib_gdef, dummy_input) - int8_graph = trt.calib_graph_to_infer_graph(int8_calib_gdef) - o5 = run_graph(int8_graph, dummy_input) - print("Is FP32 == FP16? %s (False is possible)" % np.allclose(o1, o4)) - print("Is FP32 == INT8? %s (False is possible)" % np.allclose(o1, o5)) - print("Pass") - - -def auto(multi_engine): - """Run the conversion as an optimization pass.""" - if multi_engine: - inp_dims = (2, 3, 7, 5) - orig_graph = get_multi_engine_graph_def() - else: - inp_dims = (100, 24, 24, 2) - orig_graph = get_simple_graph_def() # use a frozen graph for inference - dummy_input = np.random.random_sample(inp_dims) - opt_config = rwpb2.RewriterConfig() - opt_config.meta_optimizer_iterations = opt_config.ONE - opt_config.optimizers.extend(["constfold", "layout"]) - custom_op = opt_config.custom_optimizers.add() - custom_op.name = "TensorRTOptimizer" - custom_op.parameter_map["minimum_segment_size"].i = 3 - custom_op.parameter_map["precision_mode"].s = to_bytes("FP32") - custom_op.parameter_map["max_batch_size"].i = inp_dims[0] - custom_op.parameter_map["max_workspace_size_bytes"].i = 1 << 25 - print(custom_op) - gpu_options = None - if trt.trt_convert.get_linked_tensorrt_version()[0] == 3: - gpu_options = cpb2.GPUOptions(per_process_gpu_memory_fraction=0.50) - graph_options = cpb2.GraphOptions(rewrite_options=opt_config) - sessconfig = cpb2.ConfigProto( - gpu_options=gpu_options, graph_options=graph_options) - print(sessconfig) - g = ops.Graph() - ops.reset_default_graph() - with g.as_default(): - inp, out = importer.import_graph_def( - graph_def=orig_graph, return_elements=["input", "output"], name="") - inp = inp.outputs[0] - out = out.outputs[0] - with csess.Session(config=sessconfig, graph=g) as sess: - val = sess.run(out, {inp: dummy_input}) - print(val.shape) - - -if "__main__" in __name__: - P = argparse.ArgumentParser( - prog="tftrt_test", - description="Example utilization of TensorFlow-TensorRT integration") - P.add_argument( - "--automatic", - "-a", - action="store_true", - help="Do TRT conversion automatically", - default=False) - P.add_argument( - "--multi-engine", - "-m", - action="store_true", - help="Use a graph that will result in 2 engines", - default=False) - flags, unparsed = P.parse_known_args() - if flags.automatic: - auto(flags.multi_engine) - else: - user(flags.multi_engine) diff --git a/tensorflow/contrib/timeseries/examples/BUILD b/tensorflow/contrib/timeseries/examples/BUILD index 57797214d1684550aa7ad2664b71d22b504f70ed..e10be88ece8ebba9635af955b3c3410f29e5503c 100644 --- a/tensorflow/contrib/timeseries/examples/BUILD +++ b/tensorflow/contrib/timeseries/examples/BUILD @@ -105,6 +105,7 @@ py_binary( data = ["data/multivariate_periods.csv"], srcs_version = "PY2AND3", tags = ["no_pip"], + visibility = ["//visibility:public"], deps = select({ ":empty_condition": [], "//conditions:default": [], @@ -113,6 +114,7 @@ py_binary( "//tensorflow:tensorflow_py", "//tensorflow/contrib/timeseries/python/timeseries:estimators", "//tensorflow/contrib/timeseries/python/timeseries:model", + "//tensorflow/contrib/timeseries/python/timeseries:state_management", ], ) diff --git a/tensorflow/contrib/timeseries/examples/predict_test.py b/tensorflow/contrib/timeseries/examples/predict_test.py index 678fd71cd8b94ee0be46e10a9a673de55bd44215..b353f85cb5df0cf961d1900b241e4fa1a84a24b4 100644 --- a/tensorflow/contrib/timeseries/examples/predict_test.py +++ b/tensorflow/contrib/timeseries/examples/predict_test.py @@ -43,10 +43,6 @@ class PeriodTrendExampleTest(test.TestCase): self.assertAllEqual([700], mean.shape) self.assertAllEqual([700], upper_limit.shape) self.assertAllEqual([700], lower_limit.shape) - # Check that variance hasn't blown up too much. This is a relatively good - # indication that training was successful. - self.assertLess(upper_limit[-1] - lower_limit[-1], - 1.5 * (upper_limit[0] - lower_limit[0])) def test_ar(self): (times, observed, all_times, mean, @@ -55,7 +51,6 @@ class PeriodTrendExampleTest(test.TestCase): self.assertAllEqual(all_times.shape, mean.shape) self.assertAllEqual(all_times.shape, upper_limit.shape) self.assertAllEqual(all_times.shape, lower_limit.shape) - self.assertLess((upper_limit - lower_limit).mean(), 4.) if __name__ == "__main__": diff --git a/tensorflow/contrib/timeseries/python/timeseries/BUILD b/tensorflow/contrib/timeseries/python/timeseries/BUILD index 4b90b596b28efec83aa349782c4874d79b6817c7..d1be31ddc799ce4c4ef9baa15729fde7925f2f6c 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/BUILD +++ b/tensorflow/contrib/timeseries/python/timeseries/BUILD @@ -155,11 +155,11 @@ py_library( py_test( name = "head_test", - size = "large", + size = "medium", srcs = [ "head_test.py", ], - shard_count = 4, + shard_count = 10, srcs_version = "PY2AND3", tags = ["no_pip_gpu"], # b/63391119 deps = [ @@ -281,6 +281,7 @@ py_library( "input_pipeline.py", ], srcs_version = "PY2AND3", + visibility = ["//visibility:public"], deps = [ ":feature_keys", ":model_utils", @@ -361,9 +362,10 @@ py_library( srcs_version = "PY2AND3", deps = [ ":feature_keys", + ":math_utils", ":model", ":model_utils", - "//tensorflow/contrib/distributions:distributions_py", + "//tensorflow/contrib/rnn:rnn_py", "//tensorflow/python:array_ops", "//tensorflow/python:check_ops", "//tensorflow/python:constant_op", diff --git a/tensorflow/contrib/timeseries/python/timeseries/ar_model.py b/tensorflow/contrib/timeseries/python/timeseries/ar_model.py index bcadf4094e1e79fff1685515f2bde0b88f717cac..3626701d24163ef52564b42d8a630bd9c5a788eb 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/ar_model.py +++ b/tensorflow/contrib/timeseries/python/timeseries/ar_model.py @@ -18,9 +18,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib import distributions - from tensorflow.contrib.rnn.python.ops import lstm_ops +from tensorflow.contrib.timeseries.python.timeseries import math_utils from tensorflow.contrib.timeseries.python.timeseries import model from tensorflow.contrib.timeseries.python.timeseries import model_utils from tensorflow.contrib.timeseries.python.timeseries.feature_keys import PredictionFeatures @@ -462,11 +461,12 @@ class ARModel(model.TimeSeriesModel): if self.loss == ARModel.NORMAL_LIKELIHOOD_LOSS: covariance = prediction_ops["covariance"] sigma = math_ops.sqrt(gen_math_ops.maximum(covariance, 1e-5)) - normal = distributions.Normal(loc=targets, scale=sigma) - loss_op = -math_ops.reduce_sum(normal.log_prob(prediction)) + loss_op = -math_ops.reduce_sum( + math_utils.normal_log_prob(targets, sigma, prediction)) else: assert self.loss == ARModel.SQUARED_LOSS, self.loss - loss_op = math_ops.reduce_sum(math_ops.square(prediction - targets)) + loss_op = math_ops.reduce_sum( + math_ops.squared_difference(prediction, targets)) loss_op /= math_ops.cast( math_ops.reduce_prod(array_ops.shape(targets)), loss_op.dtype) return loss_op @@ -965,16 +965,11 @@ class AnomalyMixtureARModel(ARModel): anomaly_variance = prediction_ops["anomaly_params"] anomaly_sigma = math_ops.sqrt( gen_math_ops.maximum(anomaly_variance, 1e-5)) - normal = distributions.Normal(loc=targets, scale=anomaly_sigma) - log_prob = normal.log_prob(prediction) + log_prob = math_utils.normal_log_prob(targets, anomaly_sigma, prediction) else: assert self._anomaly_distribution == AnomalyMixtureARModel.CAUCHY_ANOMALY anomaly_scale = prediction_ops["anomaly_params"] - cauchy = distributions.StudentT( - df=array_ops.ones([], dtype=anomaly_scale.dtype), - loc=targets, - scale=anomaly_scale) - log_prob = cauchy.log_prob(prediction) + log_prob = math_utils.cauchy_log_prob(targets, anomaly_scale, prediction) return log_prob def loss_op(self, targets, prediction_ops): @@ -983,8 +978,7 @@ class AnomalyMixtureARModel(ARModel): covariance = prediction_ops["covariance"] # Normal data log probability. sigma = math_ops.sqrt(gen_math_ops.maximum(covariance, 1e-5)) - normal1 = distributions.Normal(loc=targets, scale=sigma) - log_prob1 = normal1.log_prob(prediction) + log_prob1 = math_utils.normal_log_prob(targets, sigma, prediction) log_prob1 += math_ops.log(1 - self._anomaly_prior_probability) # Anomaly log probability. log_prob2 = self._anomaly_log_prob(targets, prediction_ops) diff --git a/tensorflow/contrib/timeseries/python/timeseries/math_utils.py b/tensorflow/contrib/timeseries/python/timeseries/math_utils.py index aab330643862c1ccf073d2a0e34e1c475b1ec15f..b7375e5055e29efea3f23c3b9b9f3af59f45495b 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/math_utils.py +++ b/tensorflow/contrib/timeseries/python/timeseries/math_utils.py @@ -21,6 +21,8 @@ from __future__ import print_function import collections import math +import numpy as np + from tensorflow.contrib import lookup from tensorflow.contrib.layers.python.layers import layers @@ -43,6 +45,32 @@ from tensorflow.python.ops import variable_scope from tensorflow.python.util import nest +def normal_log_prob(loc, scale, x): + """Computes the Normal log pdf.""" + z = (x - loc) / scale + return -0.5 * (math_ops.square(z) + + np.log(2. * np.pi) + math_ops.log(scale)) + + +def cauchy_log_prob(loc, scale, x): + """Computes the Cauchy log pdf.""" + z = (x - loc) / scale + return (-np.log(np.pi) - math_ops.log(scale) - + math_ops.log1p(math_ops.square(z))) + + +def mvn_tril_log_prob(loc, scale_tril, x): + """Computes the MVN log pdf under tril scale. Doesn't handle batches.""" + x0 = x - loc + z = linalg_ops.matrix_triangular_solve( + scale_tril, x0[..., array_ops.newaxis])[..., 0] + log_det_cov = 2. * math_ops.reduce_sum(math_ops.log( + array_ops.matrix_diag_part(scale_tril)), axis=-1) + d = math_ops.cast(array_ops.shape(scale_tril)[-1], log_det_cov.dtype) + return -0.5 * (math_ops.reduce_sum(math_ops.square(z), axis=-1) + + d * np.log(2. * np.pi) + log_det_cov) + + def clip_covariance( covariance_matrix, maximum_variance_ratio, minimum_variance): """Enforce constraints on a covariance matrix to improve numerical stability. diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD index 125750e7639ad40c481472a93353e6fb7055be96..cf5e749042afd83f927a3d22edfd3a9538ab2ffd 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD +++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD @@ -78,7 +78,6 @@ py_library( srcs = ["kalman_filter.py"], srcs_version = "PY2AND3", deps = [ - "//tensorflow/contrib/distributions:distributions_py", "//tensorflow/contrib/timeseries/python/timeseries:math_utils", "//tensorflow/python:array_ops", "//tensorflow/python:control_flow_ops", @@ -235,7 +234,6 @@ py_library( srcs = ["filtering_postprocessor.py"], srcs_version = "PY2AND3", deps = [ - "//tensorflow/contrib/distributions:distributions_py", "//tensorflow/contrib/timeseries/python/timeseries:math_utils", "//tensorflow/python:array_ops", "//tensorflow/python:check_ops", diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/filtering_postprocessor.py b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/filtering_postprocessor.py index e9e2ac0aaf4c4d6c41f5007662f261af3de9bbd1..3fa2fbd9f77cb887c30fde264815728ca345f45a 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/filtering_postprocessor.py +++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/filtering_postprocessor.py @@ -22,8 +22,6 @@ import abc import six -from tensorflow.contrib import distributions - from tensorflow.contrib.timeseries.python.timeseries import math_utils from tensorflow.python.framework import dtypes @@ -91,10 +89,10 @@ def cauchy_alternative_to_gaussian(current_times, current_values, outputs): """ del current_times # unused cauchy_scale = math_utils.entropy_matched_cauchy_scale(outputs["covariance"]) - individual_log_pdfs = distributions.StudentT( - df=array_ops.ones([], dtype=current_values.dtype), + individual_log_pdfs = math_utils.cauchy_log_prob( loc=outputs["mean"], - scale=cauchy_scale).log_prob(current_values) + scale=cauchy_scale, + x=current_values) return math_ops.reduce_sum(individual_log_pdfs, axis=1) diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/kalman_filter.py b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/kalman_filter.py index a614386121e000961bf8b32625a28e1251654320..c0ec797bc5b7c41ca996c807840ce38311201f87 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/kalman_filter.py +++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/kalman_filter.py @@ -18,8 +18,6 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib import distributions - from tensorflow.contrib.timeseries.python.timeseries import math_utils from tensorflow.python.framework import dtypes @@ -137,9 +135,10 @@ class KalmanFilter(object): with ops.control_dependencies([non_negative_assert]): observation_covariance_cholesky = linalg_ops.cholesky( symmetrized_observation_covariance) - log_prediction_prob = distributions.MultivariateNormalTriL( - predicted_observation, observation_covariance_cholesky).log_prob( - observation) + log_prediction_prob = math_utils.mvn_tril_log_prob( + loc=predicted_observation, + scale_tril=observation_covariance_cholesky, + x=observation) (posterior_state, posterior_state_var) = self.posterior_from_prior_state( prior_state=estimated_state, diff --git a/tensorflow/contrib/tpu/BUILD b/tensorflow/contrib/tpu/BUILD index b9c225b44a2f03c8f6d1ef23f75cb5df5667a002..294dbddcb5ee1b1758182c10e2816f353d989084 100644 --- a/tensorflow/contrib/tpu/BUILD +++ b/tensorflow/contrib/tpu/BUILD @@ -61,6 +61,7 @@ py_library( py_library( name = "tpu_estimator", srcs = [ + "python/tpu/_tpu_estimator_embedding.py", "python/tpu/error_handling.py", "python/tpu/tpu_config.py", "python/tpu/tpu_context.py", @@ -70,7 +71,9 @@ py_library( srcs_version = "PY2AND3", deps = [ ":async_checkpoint", + ":feature_column", ":functional", + ":tpu_embedding", ":tpu_lib", ":tpu_ordinal_selector_py", "//tensorflow/contrib/training:training_py", @@ -109,12 +112,12 @@ tf_gen_op_libs( "functional_ops", ], deps = [ - "//tensorflow/contrib/tpu/proto:tpu_embedding_configuration_proto_cc", "//tensorflow/contrib/tpu/utils:tpu_embedding_optimization_parameters_utils", "//tensorflow/contrib/tpu/utils:tpu_embedding_output_layout_utils", "//tensorflow/core:lib", "//tensorflow/core:lib_proto_parsing", "//tensorflow/core:protos_all_cc", + "//tensorflow/core/protobuf/tpu:tpu_embedding_configuration_proto_cc", ], ) @@ -131,10 +134,10 @@ tf_custom_op_library( "ops/tpu_embedding_ops.cc", ], deps = [ - "//tensorflow/contrib/tpu/proto:tpu_embedding_configuration_proto_cc", "//tensorflow/contrib/tpu/utils:tpu_embedding_optimization_parameters_utils", "//tensorflow/contrib/tpu/utils:tpu_embedding_output_layout_utils", "//tensorflow/core:lib_proto_parsing", + "//tensorflow/core/protobuf/tpu:tpu_embedding_configuration_proto_cc", ], ) @@ -159,18 +162,19 @@ tf_gen_op_wrapper_py( ) tf_custom_op_library( - name = "python/ops/_tpu_ordinal_selector.so", + name = "python/ops/_tpu_ordinal_selector_op.so", srcs = ["ops/tpu_ordinal_selector_op.cc"], ) tf_custom_op_py_library( name = "tpu_ordinal_selector_py", - srcs = ["ops/gen_tpu_ordinal_selector_op.py"], - dso = [":python/ops/_tpu_ordinal_selector.so"], + srcs = ["python/ops/tpu_ordinal_selector_op.py"], + dso = [":python/ops/_tpu_ordinal_selector_op.so"], kernels = [ ":tpu_ordinal_selector_op_op_lib", ], srcs_version = "PY2AND3", + visibility = ["//visibility:public"], deps = [ ":tpu_ordinal_selector_op", ], @@ -183,6 +187,11 @@ tf_gen_op_wrapper_py( ], ) +tf_custom_op_library( + name = "python/ops/_functional_ops.so", + srcs = ["ops/functional_ops.cc"], +) + tf_gen_op_wrapper_py( name = "gen_functional_ops", out = "python/tpu/gen_functional_ops.py", @@ -192,9 +201,14 @@ tf_gen_op_wrapper_py( deps = [":functional_ops_op_lib"], ) -py_library( +tf_custom_op_py_library( name = "functional", srcs = ["python/tpu/functional.py"], + dso = [":python/ops/_functional_ops.so"], + kernels = [ + ":functional_ops_op_lib", + ], + srcs_version = "PY2AND3", visibility = [ "//visibility:public", ], @@ -217,7 +231,7 @@ py_library( tf_custom_op_py_library( name = "tpu_py", - srcs = glob(["python/ops/*.py"]), + srcs = ["python/ops/tpu_ops.py"], dso = [":python/ops/_tpu_ops.so"], kernels = [ ":all_ops", @@ -263,7 +277,6 @@ py_library( "//learning/brain:__subpackages__", "//tensorflow:__subpackages__", "//third_party/cloud_tpu/models/keras_colab:__subpackages__", - "//third_party/cloud_tpu/models/mnist_keras:__subpackages__", "//third_party/cloud_tpu/models/resnet50_keras:__subpackages__", ], deps = [ @@ -271,8 +284,8 @@ py_library( "//tensorflow/contrib/cluster_resolver:cluster_resolver_py", "//tensorflow/contrib/distribute", "//tensorflow/contrib/framework:framework_py", - "//tensorflow/contrib/tpu/proto:compilation_result_proto_py", "//tensorflow/core:protos_all_py", + "//tensorflow/core/protobuf/tpu:compilation_result_proto_py", "//tensorflow/python:array_ops", "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", @@ -312,18 +325,21 @@ py_library( srcs_version = "PY2AND3", deps = [ ":datasets", + ":functional", ":profiler", + ":tpu_ordinal_selector_py", ":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/contrib/tpu/proto:compilation_result_proto_py", - "//tensorflow/contrib/tpu/proto:optimization_parameters_proto_py", - "//tensorflow/contrib/tpu/proto:topology_proto_py", - "//tensorflow/contrib/tpu/proto:tpu_embedding_configuration_proto_py", - "//tensorflow/contrib/tpu/proto:tpu_embedding_output_layout_proto_py", "//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", @@ -453,8 +469,9 @@ py_library( srcs = ["python/tpu/tpu_embedding.py"], srcs_version = "PY2AND3", deps = [ - "//tensorflow/contrib/tpu:tpu_ops", - "//tensorflow/contrib/tpu/proto:tpu_embedding_configuration_proto_py", + ":tpu_lib", + ":tpu_ops", + "//tensorflow/core/protobuf/tpu:tpu_embedding_configuration_proto_py", "//tensorflow/python:array_ops", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:init_ops", diff --git a/tensorflow/contrib/tpu/ops/replication_ops.cc b/tensorflow/contrib/tpu/ops/replication_ops.cc index 285e11d92de7a684ed87974414ec73c274cc7aa5..d4180d1a20bc59f3fbb37b2dbc67790ded9d2d90 100644 --- a/tensorflow/contrib/tpu/ops/replication_ops.cc +++ b/tensorflow/contrib/tpu/ops/replication_ops.cc @@ -31,6 +31,7 @@ REGISTER_OP("TPUReplicateMetadata") // Deprecated. Use num_cores_per_replica instead. .Attr("computation_shape: list(int) = []") .Attr("host_compute_core: list(string) = []") + .Attr("padding_map: list(string) = []") .SetShapeFn(shape_inference::UnknownShape); REGISTER_OP("TPUReplicatedInput") @@ -105,6 +106,7 @@ REGISTER_OP("TPUReplicate") .Attr("NumVariables: int >= 0") .Attr("Tguaranteed_constants: list(type) >= 0") .Attr("output_types: list(type) >= 0") + .Attr("padding_map: list(string) = []") .Input("inputs: Tinputs") .Input("broadcast_inputs: Tbroadcast_inputs") .Input("variables: NumVariables * resource") diff --git a/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc b/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc index 676aed0b7b651494eda80ff2d7c7c31097529590..b991698359ebf713e388ea0000ca4118b020e5f9 100644 --- a/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc +++ b/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc @@ -13,7 +13,6 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/tpu/proto/tpu_embedding_configuration.pb.h" #include "tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.h" #include "tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.h" #include "tensorflow/core/framework/attr_value.pb.h" @@ -23,6 +22,7 @@ limitations under the License. #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/lib/strings/stringprintf.h" +#include "tensorflow/core/protobuf/tpu/tpu_embedding_configuration.pb.h" namespace tensorflow { @@ -466,7 +466,7 @@ inputs: A TensorList of gradients with which to update embedding tables. configuration given to tpu.initialize_system. learning_rates: A TensorList of float32 scalars, one for each dynamic learning rate tag: see the comments in - //third_party/tensorflow/contrib/tpu/proto/optimization_parameters.proto. + //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. diff --git a/tensorflow/contrib/tpu/profiler/pip_package/setup.py b/tensorflow/contrib/tpu/profiler/pip_package/setup.py index f27ae38e0434991da7475e631be1c6cb4a463118..807cf26fe983b4ebe17695d6f4f90ecfc0e0cbf5 100644 --- a/tensorflow/contrib/tpu/profiler/pip_package/setup.py +++ b/tensorflow/contrib/tpu/profiler/pip_package/setup.py @@ -33,7 +33,7 @@ setup( long_description='Tools for capture TPU profile', url='https://www.tensorflow.org/tfrc/', author='Google Inc.', - author_email='opensource@google.com', + author_email='packages@tensorflow.org', packages=['cloud_tpu_profiler'], package_data={ 'cloud_tpu_profiler': ['data/*'], diff --git a/tensorflow/contrib/tpu/profiler/trace_events.proto b/tensorflow/contrib/tpu/profiler/trace_events.proto index cb2b9162677a0ebe8240a98671b1cabc1cee0c9f..96c4784c691d8f34cf8715cdc0ed9886412f5f90 100644 --- a/tensorflow/contrib/tpu/profiler/trace_events.proto +++ b/tensorflow/contrib/tpu/profiler/trace_events.proto @@ -56,4 +56,7 @@ message TraceEvent { // The duration of the event in picoseconds if applicable. // Events without duration are called instant events. uint64 duration_ps = 10; + + // Extra arguments that will be displayed in trace view. + map args = 11; } diff --git a/tensorflow/contrib/tpu/python/ops/tpu_ops.py b/tensorflow/contrib/tpu/python/ops/tpu_ops.py index 6a6eba282a12d68cc3cd4e46a46a1b4190fb737b..55f7c6bcbc11b3a11bb3372aa4f26d3c8a87ff3c 100644 --- a/tensorflow/contrib/tpu/python/ops/tpu_ops.py +++ b/tensorflow/contrib/tpu/python/ops/tpu_ops.py @@ -157,7 +157,7 @@ if platform.system() != "Windows": _SUPPORTED_INFEED_DTYPES = set([ dtypes.bool, dtypes.int32, dtypes.int64, dtypes.bfloat16, dtypes.float32, - dtypes.complex64 + dtypes.complex64, dtypes.uint32 ]) def infeed_dequeue(dtype, shape, name=None): @@ -217,13 +217,19 @@ if platform.system() != "Windows": Args: inputs: A TensorList of gradients with which to update embedding tables. - Contains one tensor per embedding table in the model. + 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 embedding - table, containing the learning rates for each table when dynamic - learning rate is enabled through the OptimizationParameters in - TPUEmbeddingConfiguration. When the learning rate is constant, the list - should be empty (optional). + 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: @@ -337,9 +343,8 @@ if platform.system() != "Windows": 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(). + 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 diff --git a/tensorflow/contrib/tensorrt/python/ops/trt_engine_op.py b/tensorflow/contrib/tpu/python/ops/tpu_ordinal_selector_op.py similarity index 70% rename from tensorflow/contrib/tensorrt/python/ops/trt_engine_op.py rename to tensorflow/contrib/tpu/python/ops/tpu_ordinal_selector_op.py index 31a313182be9a2fca7457a539670dbc911ccabb1..5ca38cd1bae5753a7398834bd96d3b26e66b4941 100644 --- a/tensorflow/contrib/tensorrt/python/ops/trt_engine_op.py +++ b/tensorflow/contrib/tpu/python/ops/tpu_ordinal_selector_op.py @@ -1,4 +1,4 @@ -# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# 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. @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================= -"""Exposes the Python wrapper of TRTEngineOp.""" + +"""Operations to select TPU core to run.""" from __future__ import absolute_import from __future__ import division @@ -22,13 +23,16 @@ import platform if platform.system() != "Windows": # pylint: disable=wildcard-import,unused-import,g-import-not-at-top - from tensorflow.contrib.tensorrt.ops.gen_trt_engine_op import * + from tensorflow.contrib.tpu.ops.gen_tpu_ordinal_selector_op import * from tensorflow.contrib.util import loader from tensorflow.python.platform import resource_loader # pylint: enable=wildcard-import,unused-import,g-import-not-at-top - _trt_engine_op = loader.load_op_library( - resource_loader.get_path_to_datafile("_trt_engine_op.so")) + _tpu_ordinal_selector_op = loader.load_op_library( + resource_loader.get_path_to_datafile("_tpu_ordinal_selector_op.so")) + else: - raise RuntimeError("Windows platforms are not supported") + # We have already built the appropriate libraries into the binary via CMake + # if we have built contrib, so we don't need this + pass diff --git a/tensorflow/contrib/tpu/python/tpu/_tpu_estimator_embedding.py b/tensorflow/contrib/tpu/python/tpu/_tpu_estimator_embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..dd239d5d78fbdc012566398b3a5bec89eeaf4ed2 --- /dev/null +++ b/tensorflow/contrib/tpu/python/tpu/_tpu_estimator_embedding.py @@ -0,0 +1,333 @@ +# 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. +# =================================================================== +"""Tooling for support TPU embedding in TPUEstimator.""" + +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 = {} + + 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 diff --git a/tensorflow/contrib/tpu/python/tpu/device_assignment.py b/tensorflow/contrib/tpu/python/tpu/device_assignment.py index 6906501ecf90c8e577aa0becf2dba818deb19df4..3313dc749c2c7606101b2dc96614df2d052dfed1 100644 --- a/tensorflow/contrib/tpu/python/tpu/device_assignment.py +++ b/tensorflow/contrib/tpu/python/tpu/device_assignment.py @@ -25,6 +25,9 @@ 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 = {} diff --git a/tensorflow/contrib/tpu/python/tpu/feature_column.py b/tensorflow/contrib/tpu/python/tpu/feature_column.py index 8edf131bc24fd003806263570b63ee8514c49896..e80d0de70548b4b628d8ac7ebd6c9c1e5ded26ed 100644 --- a/tensorflow/contrib/tpu/python/tpu/feature_column.py +++ b/tensorflow/contrib/tpu/python/tpu/feature_column.py @@ -408,19 +408,25 @@ def _record_variable_scope_and_name(embedding_var_name, var_def_dict = collection[0] - captured_scope = None + captured_scope = variable_scope.get_variable_scope() + captured_scope_name = captured_scope.name - if is_shared_embedding and (embedding_var_name in var_def_dict): + 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 shared embedding name is different, ' + '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: - # scope contains var_scope_name. - captured_scope = variable_scope.get_variable_scope() - var_def_dict[embedding_var_name] = (captured_scope, + var_def_dict[embedding_var_name] = (captured_scope_name, embedding_var_name_in_fc) diff --git a/tensorflow/contrib/tpu/python/tpu/functional.py b/tensorflow/contrib/tpu/python/tpu/functional.py index 1ec9b5b33d007eb2eaa557438f32ea69053261c6..24c85156e53a9b770f811c4cf3b903eab6553c76 100644 --- a/tensorflow/contrib/tpu/python/tpu/functional.py +++ b/tensorflow/contrib/tpu/python/tpu/functional.py @@ -18,8 +18,22 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import platform + from tensorflow.contrib.tpu.python.tpu import gen_functional_ops TPUPartitionedCall = gen_functional_ops._tpu_partitioned_call # pylint: disable=invalid-name,protected-access + +if platform.system() != "Windows": + # pylint: disable=wildcard-import,unused-import,g-import-not-at-top + from tensorflow.contrib.tpu.ops.gen_tpu_ordinal_selector_op import * + + from tensorflow.contrib.util import loader + from tensorflow.python.platform import resource_loader + # pylint: enable=wildcard-import,unused-import,g-import-not-at-top + + _tpu_partitioned_call_op = loader.load_op_library( + resource_loader.get_path_to_datafile("../ops/_functional_ops.so") + ) diff --git a/tensorflow/contrib/tpu/python/tpu/keras_support.py b/tensorflow/contrib/tpu/python/tpu/keras_support.py index 37fe9af8c4b154a2e20a957f6ca5d97df3d413be..4322d1708e909340f603595735ca7043f36c3b10 100644 --- a/tensorflow/contrib/tpu/python/tpu/keras_support.py +++ b/tensorflow/contrib/tpu/python/tpu/keras_support.py @@ -56,7 +56,6 @@ import six from tensorflow.contrib.cluster_resolver.python.training import tpu_cluster_resolver as tpu_cluster_resolver_lib from tensorflow.contrib.framework.python.framework import experimental -from tensorflow.contrib.tpu.proto import compilation_result_pb2 as tpu_compilation_result from tensorflow.contrib.tpu.python.ops import tpu_ops from tensorflow.contrib.tpu.python.tpu import keras_tpu_variables from tensorflow.contrib.tpu.python.tpu import tpu @@ -64,6 +63,7 @@ from tensorflow.contrib.tpu.python.tpu import tpu_function 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.client import session as tf_session from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import iterator_ops diff --git a/tensorflow/contrib/tpu/python/tpu/session_support.py b/tensorflow/contrib/tpu/python/tpu/session_support.py index f5735cecc38b7033f21fc4d4105cfead233379fa..5cb2ca6478a1d7589cd2aa2d52c82306b3fd11f4 100644 --- a/tensorflow/contrib/tpu/python/tpu/session_support.py +++ b/tensorflow/contrib/tpu/python/tpu/session_support.py @@ -172,7 +172,8 @@ class WorkerHeartbeatManager(object): """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)) + 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 diff --git a/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py b/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py index 66689cde2fb53c1d9909224fcf38e5e9a7dd729b..43b9168eccbec4cd8ce874beff7b0f1d8e09e812 100644 --- a/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py +++ b/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py @@ -27,14 +27,19 @@ 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 @@ -51,6 +56,7 @@ _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' @@ -59,21 +65,27 @@ _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*--([^=]+)='([^']*)'") @@ -82,21 +94,31 @@ _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_FILE = 'trace_file_path' -_FLAG_NAME_REPORT_FILE = 'report_file_path' +_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): @@ -151,6 +173,54 @@ def keras_layer_tracepoint(layer, checkpoint_name): 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. @@ -202,15 +272,18 @@ class TensorTracer(object): 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_FILE, _FLAG_NAME_REPORT_FILE, + _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_NAME_OP_RANGE, + _FLAG_DUMP_BEFORE_AFTER_GRAPHS] tensor_tracer_flags = os.environ.get(_FLAGS_ENV_VAR) if not tensor_tracer_flags: return @@ -337,6 +410,10 @@ class TensorTracer(object): 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): @@ -369,6 +446,29 @@ class TensorTracer(object): '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.""" @@ -377,10 +477,7 @@ class TensorTracer(object): return True # Reasons for not including following op types: # Assign: cause incorrect result with CPU tracing. - # others: compilation problems. - # TODO(deveci): Check if 'Pack', 'Shape', 'Reshape', 'ArgMin', 'ArgMax' - # are still unsafe now that we have handled int64 tensors. - if op.type in ['Assign', 'Pack', 'Shape', 'Reshape', 'ArgMin', 'ArgMax']: + if op.type in ['Assign']: return True return False @@ -471,7 +568,7 @@ class TensorTracer(object): temporarily_marked_ops, sorted_ops) # pylint: disable=protected-access for ctrl_output_op in op._control_outputs: - # pylint: enable=protected-access + # pylint: enable=protected-access visit(ctrl_output_op, cycle, permanently_marked_ops, temporarily_marked_ops, sorted_ops) temporarily_marked_ops.remove(op) @@ -537,7 +634,7 @@ class TensorTracer(object): TensorTracer.check_submode(self._submode) self._part_tensor_size = _TRACE_MODE_PART_TENSOR_SIZE self._instrument_records = {} - self._set_trace_file_path() + self._set_trace_dir() self._set_report_file() self._set_op_range() self._set_excluded_opnames() @@ -545,42 +642,33 @@ class TensorTracer(object): 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, num_replicas, result_tensor): + def _add_replica_id_to_graph(self): """Adds nodes for computing the replica ID to the graph.""" - if not num_replicas: + 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' - return result_tensor - - self._num_replicas = 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') - use_replica_id = array_ops.identity(self._replica_id).op - with ops.control_dependencies([use_replica_id]): - # Adds a control dependency from the result_tensor to - # the replica_id to ensure that replica_id will be added to the graph. - return array_ops.identity(result_tensor) - - def _set_trace_file_path(self): - """Sets the path of the output trace file.""" - - found, self._trace_file_path = TensorTracer.get_flag_value( - _FLAG_NAME_TRACE_FILE) - if found and self._trace_file_path \ + 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(): - if os.path.isabs(self._trace_file_path): - raise ValueError('If use_test_undeclared_outputs_dir is set,' - 'trace_file_path cannot be an absolute path (%s)' - %self._trace_file_path) - outputs_dir = os.environ.get(_TEST_UNDECLARED_OUTPUTS_DIR_ENV_VAR) - self._trace_file_path = os.path.join(outputs_dir, - self._trace_file_path) + 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.""" @@ -662,6 +750,25 @@ class TensorTracer(object): 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.""" @@ -680,6 +787,9 @@ class TensorTracer(object): 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): @@ -726,6 +836,20 @@ class TensorTracer(object): 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.""" @@ -737,6 +861,62 @@ class TensorTracer(object): 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. @@ -775,41 +955,16 @@ class TensorTracer(object): else: msg = '"%s"'%tensor_name - if self._trace_file_path: - output_stream = _OUTPUT_STREAM_ESCAPE + self._trace_file_path + 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 - print_op = logging_ops.print_v2(msg, array_ops.shape(output_tensor), - '@', self._replica_id, - '\n', output_tensor, '\n', - summarize=num_elements, - output_stream=output_stream) - with ops.control_dependencies([print_op]): - return array_ops.identity(tensor).op - - def _detect_nan_inf(tensor): - """Trace function for detecting any NaN/Inf in the tensor.""" - - if tensor.dtype.is_floating: - output_tensor = math_ops.reduce_any( - gen_math_ops.logical_or(gen_math_ops.is_nan(tensor), - gen_math_ops.is_inf(tensor))) - else: - output_tensor = constant_op.constant(False) - - return _print_tensor(tensor_name, -1, tensor, output_tensor) - - def _show_norm(tensor): - tensor = math_ops.cast(tensor, dtypes.float64) - output_tensor = linalg_ops.norm(tensor) - return _print_tensor(tensor_name, -1, tensor, output_tensor) - - def _show_max_abs(tensor): - output_tensor = math_ops.cast(math_ops.reduce_max(math_ops.abs(tensor)), - dtypes.float64) - zero = constant_op.constant(0, dtypes.float64) - output_tensor = gen_math_ops.maximum(zero, output_tensor) - return _print_tensor(tensor_name, -1, tensor, output_tensor) + 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.""" @@ -822,34 +977,28 @@ class TensorTracer(object): return _print_tensor(tensor_name, -1, tensor, tensor) - if self._trace_mode == _TRACE_MODE_NAN_INF: - return _detect_nan_inf if self._trace_mode == _TRACE_MODE_PART_TENSOR: return _show_part_tensor - if self._trace_mode == _TRACE_MODE_FULL_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 - if self._trace_mode == _TRACE_MODE_NORM: - return _show_norm - if self._trace_mode == _TRACE_MODE_MAX_ABS: - return _show_max_abs 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): + def _skip_op(self, op_id, op, user_included, user_excluded, + in_exec_path=True): """Returns True if we should not trace Op.""" - 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 - if not self._inside_op_range(op_id): + if TensorTracer.while_loop_op(op): self._instrument_records[op.name] = TensorTracer.reason( - op_id, _REASON_OUTSIDE_OP_RANGE) + op_id, _REASON_WHILELOOP_OP) return True if TensorTracer.unsafe_op(op): self._instrument_records[op.name] = TensorTracer.reason( @@ -859,10 +1008,27 @@ class TensorTracer(object): 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, @@ -879,7 +1045,12 @@ class TensorTracer(object): 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) @@ -889,9 +1060,18 @@ class TensorTracer(object): op_id, _REASON_USER_EXCLUDED) return True if not out_tensor.get_shape().is_fully_defined(): - self._instrument_records[out_tensor.name] = TensorTracer.reason( - op_id, _REASON_DYNAMIC_SHAPE) - return True + # 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 @@ -909,18 +1089,125 @@ class TensorTracer(object): op_id, _REASON_TENSOR_GET_TRACED) return False - def _pre_tracing(self, graph): + 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) - return (operations, succeed, sorted_or_cycle) + 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.""" @@ -950,24 +1237,200 @@ class TensorTracer(object): _TENSOR_TRACER_CHECKPOINT)) return checkpoint_operations - def trace_tpu(self, graph, result_tensor, num_replicas=None): - """Traces the tensors generated by TPU Ops in a TF graph. + 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 executed on the TPU. - result_tensor: a result tensor of evaluating the graph. - num_replicas: number of replicas used on the TPU. + 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: - A tuple (result_tensor_copy, tracing_ops), where: - result_tensor_copy: an exact copy of result_tensor - tracing_ops: a list of tracing ops. If this list - is non empty, the caller of this function - should pose control dependencies upon these - Ops so that they will be executed when the - graph is evaluated. + 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.""" @@ -980,79 +1443,181 @@ class TensorTracer(object): return math_ops.cast(tensor, dtypes.float32) return tensor - self._device_type = _DEVICE_TYPE_TPU TensorTracer.check_device_type(self._device_type) - result_tensor_copy = self._add_replica_id_to_graph(num_replicas, - result_tensor) - (operations, succeed, sorted_or_cycle) = self._pre_tracing(graph) + # 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 = [] - 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) - if self._skip_op(op_id, op, user_included, user_excluded): - continue + # 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] - if self._skip_tensor(op_id, out_tensor, user_included, - user_excluded): + 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() - tensor_name = out_tensor.name - out_tensor = _cast_unsupported_dtypes(out_tensor) - trace_op = tpu.outside_compilation( - self._make_tensor_trace_fun(tensor_name), out_tensor) - if consumers: - for consumer_op in consumers: - # pylint: disable=protected-access - consumer_op._add_control_input(trace_op) - # pylint: enable=protected-access + # 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: - # if there is no consumer, we will add the control dependence later - # when we add the control dependency to the output operations. + 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) - return (result_tensor_copy, tracing_ops) + # 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. - def trace_cpu(self, 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. + """ + + 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: - tracing_calls: a map from keys to trace calls. - A key is constructed from an Op's name. - A trace call consists of a function and a tensor ( - the function will be invoked with the tensor). + tensor_fetches: an exact copy of tensor_fetches that has additional + dependencies. + Raises: + RuntimeError: If tensor_fetches is None or empty. """ - self._device_type = _DEVICE_TYPE_CPU - TensorTracer.check_device_type(self._device_type) self._num_replicas = 1 + self._num_replicas_per_host = 1 + self._num_hosts = 1 self._replica_id = 0 - (operations, succeed, sorted_or_cycle) = self._pre_tracing(graph) - tracing_calls = {} - checkpoint_operations = self._get_checkpoints(graph) + 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 + - 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) - if self._skip_op(op_id, op, user_included, user_excluded): - 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 - trace_fun = self._make_tensor_trace_fun(out_tensor.name) - trace_call = (trace_fun, [out_tensor]) - trace_call_key = 'tensor_tracing_cpu-%s:%d'%(op.name, i) - tracing_calls[trace_call_key] = trace_call - self._post_tracing(succeed, sorted_or_cycle) - return tracing_calls diff --git a/tensorflow/contrib/tpu/python/tpu/topology.py b/tensorflow/contrib/tpu/python/tpu/topology.py index 6ae718cc2c9716587849aeee8abcd0a1de82a9ae..00ee21e694d15d2e795b0b35289e1e116b9e76cf 100644 --- a/tensorflow/contrib/tpu/python/tpu/topology.py +++ b/tensorflow/contrib/tpu/python/tpu/topology.py @@ -21,7 +21,7 @@ from __future__ import print_function import numpy as np from six.moves import xrange # pylint: disable=redefined-builtin -from tensorflow.contrib.tpu.proto import topology_pb2 +from tensorflow.core.protobuf.tpu import topology_pb2 def _tpu_device_name(job, task, device): diff --git a/tensorflow/contrib/tpu/python/tpu/tpu.py b/tensorflow/contrib/tpu/python/tpu/tpu.py index a267dd435c04c21bbeb3dc09a325267e7b22286e..673129b4bef8a7470192a5d7650a858257f653bb 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu.py @@ -19,20 +19,24 @@ 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 @@ -323,6 +327,30 @@ class TPUReplicateContext(control_flow_ops.XLAControlFlowContext): 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: @@ -373,11 +401,14 @@ class TPUReplicateContext(control_flow_ops.XLAControlFlowContext): 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. - external_control_inputs = [ - array_ops.identity(x.outputs[0]).op - for x in external_control_inputs - if x.outputs - ] + 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 @@ -481,7 +512,8 @@ def replicate(computation, inputs=None, infeed_queue=None, device_assignment=None, - name=None): + name=None, + maximum_shapes=None): """Builds a graph operator that runs a replicated TPU computation. Args: @@ -502,15 +534,125 @@ def replicate(computation, 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 lists of output tensors, indexed by `[replica_num][output_num]`. + 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)[1] + 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, @@ -518,7 +660,8 @@ def split_compile_and_replicate(computation, infeed_queue=None, device_assignment=None, name=None, - use_tpu=True): + 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 @@ -548,6 +691,15 @@ def split_compile_and_replicate(computation, 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]`. @@ -555,6 +707,10 @@ def split_compile_and_replicate(computation, 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 @@ -633,13 +789,34 @@ def split_compile_and_replicate(computation, 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 + graph = ops.get_default_graph() # Fan-in: Builds a TPUReplicatedInput node for each input. - computation_inputs = [] - for i in range(0, flat_input_arity): + flat_replicated_inputs = [] + for i in range(0, len(flat_inputs[0])): replicas = [flat_inputs[replica][i] for replica in xrange(num_replicas)] - computation_inputs.append( + flat_replicated_inputs.append( tpu_ops.tpu_replicated_input(replicas, name="input{}".format(i))) cluster_name = graph.unique_name("cluster") @@ -659,18 +836,26 @@ def split_compile_and_replicate(computation, # computation. This is to avoid orphaned TPUReplicatedInput nodes. # TODO(phawkins): consider instead pruning unused TPUReplicatedInput # and eliding trivial TPUReplicatedInput/TPUReplicatedOutput pairs. - computation_inputs = [ + flat_replicated_inputs = [ array_ops.identity(x, name="replicated_input_{}".format(i)) - for i, x in enumerate(computation_inputs) + for i, x in enumerate(flat_replicated_inputs) ] - for i in computation_inputs: + for i in flat_replicated_inputs: # pylint: disable=protected-access - i.op._set_attr("_tpu_input_identity", attr_value_pb2.AttrValue(b=True)) + # 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=computation_inputs) + 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. @@ -712,51 +897,12 @@ def split_compile_and_replicate(computation, vscope.set_use_resource(saved_use_resource) vscope.set_custom_getter(saved_custom_getter) - # 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, (list, tuple)): - 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.") - output_arity = len(output_tensors) + 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) - # 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) - output_tensors = new_output_tensors context.ExitResult(output_tensors) finally: context.report_unsupported_operations() @@ -768,11 +914,6 @@ def split_compile_and_replicate(computation, 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 - # Fan-out: Builds a TPUReplicatedOutput node for each output. - outputs = [tpu_ops.tpu_replicated_output(output_tensors[i], num_replicas, - name="output{}".format(i)) - for i in xrange(output_arity)] - with ops.control_dependencies([metadata]): if use_tpu: compile_status = tpu_ops.tpu_compilation_result() @@ -782,28 +923,146 @@ def split_compile_and_replicate(computation, else: compile_status = control_flow_ops.no_op(name="compilation_status") - with ops.control_dependencies(output_operations): - if output_arity == 0: - # Returns a list of NoOps dependent on the replication Op, indexed by - # [replica_num]. - return [ - compile_status, [ - control_flow_ops.no_op(name="shard_%d" % i) - for i in range(num_replicas) - ] - ] - else: - # Wraps the outputs in identity operators so the names of any possible - # `fetch` nodes are preserved by the replication rewrite. - return [ - compile_status, [[ - array_ops.identity( - outputs[out][replica], - name="output_%d_shard_%d" % (out, replica)) - for out in xrange(output_arity) - ] - for replica in xrange(num_replicas)] + 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, @@ -830,9 +1089,6 @@ def split_compile_and_shard(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. @@ -874,6 +1130,8 @@ def split_compile_and_shard(computation, 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.") @@ -1126,11 +1384,15 @@ def rewrite(computation, case the core attached to task 0, TPU device 0 is used. name: (Deprecated) Does nothing. Returns: - A list of output tensors. + 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. """ - if inputs is not None and not isinstance(inputs, (list, tuple)): - raise TypeError("tpu.rewrite() inputs must be a list or tuple") - # TODO(b/36647078) remove disable when pylint bug is fixed. # pylint: disable=indexing-exception return replicate( diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_config.py b/tensorflow/contrib/tpu/python/tpu/tpu_config.py index 9f8d14706845baa1ed45c84b2c15d372915a0eb4..82ebb9a37f22db0c150481d3e0ce3877f882d245 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_config.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_config.py @@ -223,7 +223,8 @@ class RunConfig(run_config_lib.RunConfig): # 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) + 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 ' diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_context.py b/tensorflow/contrib/tpu/python/tpu/tpu_context.py index 672462447944b777375331d49727c4d5366cf295..ed1e0f0401a96c34e6ff9323685857b64e10bd14 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_context.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_context.py @@ -21,6 +21,7 @@ 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 @@ -192,8 +193,14 @@ class _InternalTPUContext(object): ``` """ - def __init__(self, config, train_batch_size, eval_batch_size, - predict_batch_size, use_tpu, eval_on_tpu=True): + 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 @@ -208,7 +215,7 @@ class _InternalTPUContext(object): 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 num_cores_per_replica: + if self._model_parallelism_enabled: self._computation_shape = _NUM_CORES_TO_COMPUTATION_SHAPE[ num_cores_per_replica] else: @@ -216,6 +223,8 @@ class _InternalTPUContext(object): 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: @@ -293,6 +302,30 @@ class _InternalTPUContext(object): 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 @@ -710,11 +743,15 @@ class _OneCoreTPUContext(_InternalTPUContext): def _get_tpu_context(config, train_batch_size, eval_batch_size, - predict_batch_size, use_tpu, eval_on_tpu): + 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.)') @@ -722,4 +759,5 @@ def _get_tpu_context(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) + predict_batch_size, use_tpu, eval_on_tpu, + embedding_config_spec) diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_embedding.py b/tensorflow/contrib/tpu/python/tpu/tpu_embedding.py index ccba8a46c7cad0337119672e02314684f4451479..eb99a18d83987b098dcae9d58d9af14deebc4f56 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_embedding.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_embedding.py @@ -26,8 +26,10 @@ import six from tensorflow.contrib.framework.python.framework import experimental from tensorflow.contrib.tpu.ops import gen_tpu_ops -from tensorflow.contrib.tpu.proto import tpu_embedding_configuration_pb2 as elc 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 @@ -43,19 +45,6 @@ from tensorflow.python.ops import variables TRAINING = elc.TPUEmbeddingConfiguration.TRAINING INFERENCE = elc.TPUEmbeddingConfiguration.INFERENCE -# TODO(shizhiw): A better interface is to make `num_hosts` and -# `num_cores_per_host` optional parameters for `TPUEmbedding` -# constructor. Usually they can be automatically detected, but -# user can also specify them for debugging (b/112112496). -# Auto-detection can be done with `tpu_system_metadata.py`. -_MASTER_JOB = 'tpu_worker' -_HOST_PATTERN = '/job:tpu_worker/task:{}/device:CPU:0' -_NUM_CORES_PER_HOST = 8 - -_TEST_MASTER_JOB = None -_TEST_HOST = '/replica:0/task:0/device:CPU:0' -_TEST_NUM_CORES_PER_HOST = 2 - class TableConfig( collections.namedtuple( @@ -112,6 +101,25 @@ class TableConfig( 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): @@ -248,6 +256,7 @@ class TPUEmbedding(object): 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() @@ -273,6 +282,7 @@ class TPUEmbedding(object): embedding.config_proto)) sess.run(variables.global_variables_initializer()) sess.run(embedding.init_ops) + sess.run(embedding_variables_and_ops.load_ops()) sess.run(enqueue_ops) loss_val = sess.run(loss) ``` @@ -301,10 +311,9 @@ class TPUEmbedding(object): table_to_config_dict, feature_to_table_dict, batch_size, - num_hosts, mode, - optimization_parameters=None, - tpu_embedding_test=False): + master, + optimization_parameters=None): """API for using TPU for embedding lookups. Args: @@ -315,12 +324,11 @@ class TPUEmbedding(object): 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. - num_hosts: An `int` representing the number of TPU hosts. 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. - tpu_embedding_test: A `bool`. Only used for testing. Raises: ValueError: if any input is invalid. @@ -337,15 +345,17 @@ class TPUEmbedding(object): self._batch_size = batch_size - if tpu_embedding_test: - self._num_hosts = 1 - self._hosts = [_TEST_HOST] - self._num_cores_per_host = _TEST_NUM_CORES_PER_HOST - else: - self._num_hosts = num_hosts - self._hosts = [_HOST_PATTERN.format(i) for i in range(self._num_hosts)] - self._num_cores_per_host = _NUM_CORES_PER_HOST - self._num_cores = self._num_cores_per_host * self._num_hosts + 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 @@ -379,9 +389,6 @@ class TPUEmbedding(object): self._config_proto = self._create_config_proto() - self._create_variables_and_ops() - self._init_ops.extend(self._load_parameters_ops) - @property def hosts(self): """A list of device names for CPU hosts. @@ -389,7 +396,7 @@ class TPUEmbedding(object): Returns: A list of device names for CPU hosts. """ - return self._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`. @@ -447,23 +454,17 @@ class TPUEmbedding(object): """ return self._init_ops - # TODO(shizhiw): get table variables the same way as getting slot variables. @property - def table_to_table_variables_dict(self): - return copy.copy(self._table_to_table_variables_dict) + def table_to_config_dict(self): + return copy.copy(self._table_to_config_dict) - def get_slot_names(self): - """Return a list of the names of slots created by `TPUEmbedding`.""" - return self._optimizer_handler.get_slot_names() - - def get_slot(self, table, name): - """Return a slot named `name` create for `table` by `TPUEmbedding`.""" - return self._optimizer_handler.get_slot(table, name) + @property + def feature_to_table_dict(self): + return copy.copy(self._feature_to_table_dict) - # TODO(shizhiw): expose load to user too? @property - def retrieve_parameters_ops(self): - return self._retrieve_parameters_ops + def optimization_parameters(self): + return self._optimization_parameters def _create_config_proto(self): """Create `TPUEmbeddingConfiguration`.""" @@ -481,6 +482,11 @@ class TPUEmbedding(object): 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) @@ -495,30 +501,88 @@ class TPUEmbedding(object): return config_proto - def _create_variables_and_ops(self): - """Create embedding variables and return ops to load them into TPU.""" - self._load_parameters_ops = [] - self._retrieve_parameters_ops = [] - self._table_to_table_variables_dict = {} + 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): - # TODO(shizhiw): allow user to specify variable name so that - # they could make the name consistent with CPU etc. - variable_name = table table_variables = _create_partitioned_variables( - name=variable_name, + 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]) - self._table_to_table_variables_dict[table] = table_variables - - self._optimizer_handler.create_variables_and_ops( - table, variable_name, self._num_hosts, - self._table_to_config_dict[table], table_variables, - self._load_parameters_ops, self._retrieve_parameters_ops) + 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 _create_dummy_table_variables(self): """Create dummy embedding table variables. @@ -812,13 +876,11 @@ class _OptimizerHandler(object): def set_optimization_parameters(self, table_descriptor): raise NotImplementedError() - def create_variables_and_ops(self, table, variable_name): - raise NotImplementedError() - - def get_slot_names(self): + def get_default_slot_variable_names(self, table): raise NotImplementedError() - def get_slot(self, table, name): + def create_variables_and_ops(self, table, slot_variable_names, num_hosts, + table_config, table_variables): raise NotImplementedError() @@ -832,51 +894,64 @@ class _AdagradHandler(_OptimizerHandler): def set_optimization_parameters(self, table_descriptor): table_descriptor.optimization_parameters.adagrad.SetInParent() - def create_variables_and_ops(self, table, variable_name, num_hosts, - table_config, table_variables, - load_parameters_ops, retrieve_parameters_ops): - optimizer_name = 'Adagrad' + 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='%s/%s' % (variable_name, optimizer_name), + 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) - - self._table_to_accumulator_variables_dict[table] = accumulator_variables - 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)) - 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)) - - load_parameters_ops.append(load_parameters_op) - retrieve_parameters_ops.append(retrieve_parameters_op) - - def get_slot_names(self): - return ['accumulator'] - - def get_slot(self, table, name): - if name not in self.get_slot_names(): - raise ValueError('Adagrad has {} as slot names; got {}.' - .format(self.get_slot_names(), name)) - return self._table_to_accumulator_variables_dict[table] + 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): @@ -899,13 +974,15 @@ class _AdamHandler(_OptimizerHandler): table_descriptor.optimization_parameters.adam.use_sum_inside_sqrt = ( self._optimization_parameters.sum_inside_sqrt) - def create_variables_and_ops(self, table, variable_name, num_hosts, - table_config, table_variables, - load_parameters_ops, retrieve_parameters_ops): - optimizer_name = 'Adam' + 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='%s/%s/m' % (variable_name, optimizer_name), + name=slot_variable_names.m, num_hosts=num_hosts, vocabulary_size=table_config.vocabulary_size, embedding_dimension=table_config.dimension, @@ -913,52 +990,63 @@ class _AdamHandler(_OptimizerHandler): initializer=m_initializer) v_initializer = init_ops.zeros_initializer() v_variables = _create_partitioned_variables( - name='%s/%s/v' % (variable_name, optimizer_name), + 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) - - self._table_to_m_variables_dict[table] = m_variables - self._table_to_v_variables_dict[table] = v_variables - - 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)) - 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)) - - load_parameters_ops.append(load_parameters_op) - retrieve_parameters_ops.append(retrieve_parameters_op) - - def get_slot_names(self): - return ['m', 'v'] - - def get_slot(self, table, name): - if name == 'm': - return self._table_to_m_variables_dict[table] - elif name == 'v': - return self._table_to_v_variables_dict[table] - else: - raise ValueError('Adam has {} as slot names; got {}.' - .format(self.get_slot_names(), name)) + 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): @@ -968,38 +1056,58 @@ class _StochasticGradientDescentHandler(_OptimizerHandler): (table_descriptor.optimization_parameters.stochastic_gradient_descent .SetInParent()) - def create_variables_and_ops(self, table, variable_name, num_hosts, - table_config, table_variables, - load_parameters_ops, retrieve_parameters_ops): + 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 - 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)) - 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)) - - load_parameters_ops.append(load_parameters_op) - retrieve_parameters_ops.append(retrieve_parameters_op) - - def get_slot_names(self): - return [] - - def get_slot(self, table, name): - raise ValueError('Stochastic gradient descent does not have slot variable.') + 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): @@ -1077,34 +1185,3 @@ def _create_partitioned_variables(name, initializer=initializer, collections=collections, trainable=False)) - - -@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 - ] diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py index 98babbb54385c08c63836cce54dee451ecb73960..4f761e3599bfbd3a9429c8d456ae0b368229904f 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py @@ -31,9 +31,9 @@ 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.ops import gen_tpu_ordinal_selector_op -from tensorflow.contrib.tpu.proto import compilation_result_pb2 as tpu_compilation_result from tensorflow.contrib.tpu.python.ops import tpu_ops +from tensorflow.contrib.tpu.python.ops import tpu_ordinal_selector_op +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 @@ -44,10 +44,13 @@ from tensorflow.contrib.tpu.python.tpu import tpu_context from tensorflow.contrib.tpu.python.tpu import tpu_feed 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 @@ -124,6 +127,16 @@ def _is_iterable(obj): 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: @@ -341,22 +354,18 @@ class TPUEstimatorSpec(model_fn_lib._TPUEstimatorSpec): # pylint: disable=prote hooks = None if self.host_call is not None: hooks = [_OutfeedHostCallHook(host_call_ret['host_call'])] - if tensor_tracer.TensorTracer.is_enabled(): + loss = self.loss + if tensor_tracer.TensorTracer.is_enabled() \ + and self.train_op is not None: tt = tensor_tracer.TensorTracer() - tracing_calls = tt.trace_cpu(ops.get_default_graph()) - tracing_call_ret = _OutfeedHostCall.create_cpu_hostcall(tracing_calls) - tracing_functions = tracing_call_ret.values() - if tracing_functions: - if hooks: - hooks.extend([_OutfeedHostCallHook(tracing_functions)]) - else: - hooks = [_OutfeedHostCallHook(tracing_functions)] + 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=self.loss, + loss=loss, train_op=self.train_op, eval_metric_ops=eval_metric_ops, export_outputs=self.export_outputs, @@ -431,26 +440,36 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): run_infeed_loop_on_coordinator=True, rendezvous=None, master=None, - session_config=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 - self._should_initialize_tpu = True + # 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() - self._init_ops = [] if self._should_initialize_tpu: self._finalize_ops = [tpu.shutdown_system(job=self._master_job)] else: @@ -510,7 +529,10 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): 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)) + 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, @@ -852,6 +874,7 @@ def generate_per_host_v2_enqueue_ops_fn_for_host( """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): @@ -870,6 +893,10 @@ def generate_per_host_v2_enqueue_ops_fn_for_host( 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 = ( @@ -898,6 +925,11 @@ def generate_per_host_v2_enqueue_ops_fn_for_host( 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: @@ -1338,7 +1370,7 @@ def call_computation(computation, return tpu_functional.TPUPartitionedCall( args=tpu_subgraph.captured_inputs, - device_ordinal=gen_tpu_ordinal_selector_op.tpu_ordinal_selector(), + device_ordinal=tpu_ordinal_selector_op.tpu_ordinal_selector(), Tout=[o.type for o in tpu_subgraph.definition.signature.output_arg], f=tpu_subgraph) else: @@ -1364,6 +1396,12 @@ class _ModelFnWrapper(object): 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): + if self._ctx.embedding_config: + tpu_embedding_ = self._ctx.embedding_config.tpu_embedding + embedding_activations = tpu_embedding_.get_activations() + 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. @@ -1396,6 +1434,7 @@ class _ModelFnWrapper(object): del loss # unused; required in function signature. inputs = dequeue_fn() features, labels = inputs.features_and_labels() + self._add_embedding_features(features) estimator_spec = self._verify_estimator_spec( self._call_model_fn(features, labels)) @@ -1408,15 +1447,23 @@ class _ModelFnWrapper(object): captured_training_hooks.capture(estimator_spec.training_hooks) - tracing_ops = [] if tensor_tracer.TensorTracer.is_enabled(): tt = tensor_tracer.TensorTracer() - loss, tracing_ops = tt.trace_tpu(ops.get_default_graph(), loss, - self._ctx.num_replicas) + loss = tt.trace_tpu(ops.get_default_graph(), + loss, train_op, + self._ctx.num_replicas, + self._ctx.num_of_replicas_per_host, + self._ctx.num_hosts) + + if self._ctx.embedding_config is None: + apply_sparse_grads = [] + else: + tpu_embedding_ = self._ctx.embedding_config.tpu_embedding + apply_sparse_grads = [tpu_embedding_.generate_send_gradients_op()] # We must run train_op to update the variables prior to running the # outfeed. - with ops.control_dependencies([train_op]+tracing_ops): + 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): @@ -1462,6 +1509,7 @@ class _ModelFnWrapper(object): """Evaluation step function for use inside a while loop.""" inputs = dequeue_fn() features, labels = inputs.features_and_labels() + self._add_embedding_features(features) tpu_estimator_spec = self._call_model_fn(features, labels) if not isinstance(tpu_estimator_spec, model_fn_lib._TPUEstimatorSpec): # pylint: disable=protected-access @@ -1801,6 +1849,10 @@ class _OutfeedHostCall(object): 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: @@ -1808,6 +1860,14 @@ class _OutfeedHostCall(object): 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. # @@ -1841,7 +1901,7 @@ class _OutfeedHostCall(object): # The user-provided eval_metrics[1] is a dict. dequeue_ops = dict(zip(self._tensor_keys[name], dequeue_ops)) try: - ret[name] = self._host_fns[name](**dequeue_ops) + 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 ' @@ -1849,8 +1909,10 @@ class _OutfeedHostCall(object): 'function\'s arguments', name, e, name) raise else: - ret[name] = self._host_fns[name](*dequeue_ops) + 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 @@ -1957,9 +2019,9 @@ class TPUEstimator(estimator_lib.Estimator): ========== `model_fn` should return `TPUEstimatorSpec`, which expects the `eval_metrics` - for TPU evaluation. However, if eval_on_tpu is False, `model_fn` must return - `EstimatorSpec` and the evaluation will execute on CPU or GPU; in this case - the following discussion on TPU evaluation does not apply. + 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 @@ -2142,8 +2204,11 @@ class TPUEstimator(estimator_lib.Estimator): 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_exported_model_uses_all_cores=False, + experimental_export_device_assignment=False, + experimental_embedding_config_spec=None): """Constructs an `TPUEstimator` instance. Args: @@ -2186,7 +2251,11 @@ class TPUEstimator(estimator_lib.Estimator): 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 besides the one on CPU. + 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 @@ -2198,6 +2267,13 @@ class TPUEstimator(estimator_lib.Estimator): 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. @@ -2259,11 +2335,19 @@ class TPUEstimator(estimator_lib.Estimator): # 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) + 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 = {} @@ -2277,35 +2361,43 @@ class TPUEstimator(estimator_lib.Estimator): export_tags=None, check_variables=True): if self._export_to_tpu and mode != model_fn_lib.ModeKeys.PREDICT: - raise NotImplementedError( - 'TPUEstimator only handles mode PREDICT for exporting ' - 'when `export_to_tpu` is `True`; ' - 'got {}.'.format(mode)) - - (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)) + logging.warning('TPUEstimator only handles mode PREDICT for exporting ' + 'when `export_to_tpu` is `True`; Mode {} will be ignored ' + 'for TPU.'.format(mode)) - if self._export_to_tpu: + 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=False, + save_variables=save_variables, mode=mode, export_tags=export_tags, - check_variables=False)) + check_variables=check_variables)) def _call_model_fn(self, features, labels, mode, config): if mode == _REWRITE_FOR_INFERENCE_MODE: @@ -2365,7 +2457,16 @@ class TPUEstimator(estimator_lib.Estimator): tpu_computation, tpu_capture = self._build_tpu_computation_for_inference( features, labels, mode, config) - tensors_on_cpu = tpu.rewrite_for_inference(tpu_computation) + 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 + 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()) @@ -2654,7 +2755,11 @@ class TPUEstimator(estimator_lib.Estimator): if self._log_every_n_steps is not None: examples_hook = ExamplesPerSecondHook( ctx.global_batch_size, - output_dir=self.model_dir, + # 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): @@ -2671,6 +2776,10 @@ class TPUEstimator(estimator_lib.Estimator): assert callable(features), '`input_fn` is not callable.' input_fn = features + tpu_init_ops = [] + if ctx.embedding_config: + tpu_init_ops.extend(ctx.embedding_config.tpu_embedding.init_ops) + 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()) @@ -2685,6 +2794,24 @@ class TPUEstimator(estimator_lib.Estimator): 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 = [] @@ -2724,7 +2851,7 @@ class TPUEstimator(estimator_lib.Estimator): 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: @@ -2760,6 +2887,8 @@ class TPUEstimator(estimator_lib.Estimator): 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() @@ -2830,7 +2959,8 @@ class TPUEstimator(estimator_lib.Estimator): rendezvous=self._rendezvous[mode], master=self._config.evaluation_master, session_config=self._session_config, - )] + input_hooks + tpu_init_ops=tpu_init_ops) + ] + input_hooks if eval_hooks: hooks.extend(eval_hooks) diff --git a/tensorflow/contrib/tpu/utils/BUILD b/tensorflow/contrib/tpu/utils/BUILD index c27b73728702dcb1c84a82d3a07d15978ed2710f..5cbed402f75bf7ecf67ea06a2ad8d89260d7c1d1 100644 --- a/tensorflow/contrib/tpu/utils/BUILD +++ b/tensorflow/contrib/tpu/utils/BUILD @@ -8,9 +8,9 @@ cc_library( hdrs = ["tpu_embedding_optimization_parameters_utils.h"], visibility = ["//visibility:public"], deps = [ - "//tensorflow/contrib/tpu/proto:optimization_parameters_proto_cc", "//tensorflow/core:framework_headers_lib", "//tensorflow/core:lib_proto_parsing", + "//tensorflow/core/protobuf/tpu:optimization_parameters_proto_cc", "@com_google_absl//absl/base", ], ) @@ -21,10 +21,10 @@ cc_library( hdrs = ["tpu_embedding_output_layout_utils.h"], visibility = ["//visibility:public"], deps = [ - "//tensorflow/contrib/tpu/proto:tpu_embedding_configuration_proto_cc", - "//tensorflow/contrib/tpu/proto:tpu_embedding_output_layout_proto_cc", "//tensorflow/core:framework_headers_lib", "//tensorflow/core:lib_proto_parsing", "//tensorflow/core:protos_all_cc", + "//tensorflow/core/protobuf/tpu:tpu_embedding_configuration_proto_cc", + "//tensorflow/core/protobuf/tpu:tpu_embedding_output_layout_proto_cc", ], ) diff --git a/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.cc b/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.cc index 76cb5531cd0bc3a375d1434c31fa14a9d7f42476..d1df7e78abb166c5a4d7e508a5f1ade16e44854b 100644 --- a/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.cc +++ b/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.cc @@ -134,12 +134,16 @@ Status GetGradientAccumulationSupport(OptimizationAlgorithm alg, } } namespace { -// Make a normal state variable specification. +// Make a normal state variable specification. Please refer to +// //tensorflow/core/protobuf/tpu/optimization_parameters.proto +// (StateVariableSpecification message) for instructions on how to set the +// padding_initial_value field. StateVariableSpecification MakeStandardStateVariableSpecification( - const string& name) { + const string& name, double padding_initial_value) { StateVariableSpecification result; result.set_name(name); - result.mutable_user_defined(); + result.mutable_user_defined()->set_padding_initial_value( + padding_initial_value); return result; } } // namespace @@ -149,14 +153,14 @@ Status GetOptimizationAlgorithmStateVariables( std::vector* state_variables) { // The first parameter set is always the weights themselves. state_variables->push_back( - MakeStandardStateVariableSpecification("parameters")); + MakeStandardStateVariableSpecification("parameters", 0.0)); // The order of the returned parameters needs to match the offsets used by // the algorithm implementations in test_util.cc and // address_handler_program_creator.cc. switch (alg) { case OptimizationAlgorithm::kAdagrad: { state_variables->push_back( - MakeStandardStateVariableSpecification("accumulators")); + MakeStandardStateVariableSpecification("accumulators", 0.1)); break; } case OptimizationAlgorithm::kStochasticGradientDescent: { @@ -165,53 +169,58 @@ Status GetOptimizationAlgorithmStateVariables( } case OptimizationAlgorithm::kFtrl: { state_variables->push_back( - MakeStandardStateVariableSpecification("accumulators")); + MakeStandardStateVariableSpecification("accumulators", 0.1)); state_variables->push_back( - MakeStandardStateVariableSpecification("linears")); + MakeStandardStateVariableSpecification("linears", 0.0)); break; } case OptimizationAlgorithm::kAdam: { state_variables->push_back( - MakeStandardStateVariableSpecification("momenta")); + MakeStandardStateVariableSpecification("momenta", 0.0)); state_variables->push_back( - MakeStandardStateVariableSpecification("velocities")); + MakeStandardStateVariableSpecification("velocities", 0.0)); break; } case OptimizationAlgorithm::kMomentum: { state_variables->push_back( - MakeStandardStateVariableSpecification("momenta")); + MakeStandardStateVariableSpecification("momenta", 0.0)); break; } case OptimizationAlgorithm::kRmsProp: { - state_variables->push_back(MakeStandardStateVariableSpecification("ms")); - state_variables->push_back(MakeStandardStateVariableSpecification("mom")); + state_variables->push_back( + MakeStandardStateVariableSpecification("ms", 1.0)); + state_variables->push_back( + MakeStandardStateVariableSpecification("mom", 0.0)); break; } case OptimizationAlgorithm::kCenteredRmsProp: { - state_variables->push_back(MakeStandardStateVariableSpecification("ms")); - state_variables->push_back(MakeStandardStateVariableSpecification("mom")); - state_variables->push_back(MakeStandardStateVariableSpecification("mg")); + state_variables->push_back( + MakeStandardStateVariableSpecification("ms", 1.0)); + state_variables->push_back( + MakeStandardStateVariableSpecification("mom", 0.0)); + state_variables->push_back( + MakeStandardStateVariableSpecification("mg", 0.0)); break; } case OptimizationAlgorithm::kMdlAdagradLight: { state_variables->push_back( - MakeStandardStateVariableSpecification("accumulators")); + MakeStandardStateVariableSpecification("accumulators", 0.1)); state_variables->push_back( - MakeStandardStateVariableSpecification("weights")); + MakeStandardStateVariableSpecification("weights", 0.0)); state_variables->push_back( - MakeStandardStateVariableSpecification("benefits")); + MakeStandardStateVariableSpecification("benefits", 0.0)); break; } case OptimizationAlgorithm::kAdadelta: { state_variables->push_back( - MakeStandardStateVariableSpecification("accumulators")); + MakeStandardStateVariableSpecification("accumulators", 0.0)); state_variables->push_back( - MakeStandardStateVariableSpecification("updates")); + MakeStandardStateVariableSpecification("updates", 0.0)); break; } case OptimizationAlgorithm::kProximalAdagrad: { state_variables->push_back( - MakeStandardStateVariableSpecification("accumulators")); + MakeStandardStateVariableSpecification("accumulators", 0.1)); break; } case OptimizationAlgorithm::PARAMETERS_NOT_SET: { diff --git a/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.h b/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.h index 81d50264edb93e889d736c62a493b058e2f1bd56..7a7833bf2db2267f9bd05c4cda9baf5d3320ed3b 100644 --- a/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.h +++ b/tensorflow/contrib/tpu/utils/tpu_embedding_optimization_parameters_utils.h @@ -18,8 +18,8 @@ limitations under the License. #include #include "absl/base/casts.h" -#include "tensorflow/contrib/tpu/proto/optimization_parameters.pb.h" #include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/protobuf/tpu/optimization_parameters.pb.h" namespace tensorflow { namespace tpu { diff --git a/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.cc b/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.cc index 8480ec4b8bb98e867db3e4e4ed14d4cc529efe49..e65abe3894ead90f2b11444a8e246cb2f1d6a252 100644 --- a/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.cc +++ b/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.cc @@ -14,8 +14,8 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.h" -#include "tensorflow/contrib/tpu/proto/tpu_embedding_output_layout.pb.h" #include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/protobuf/tpu/tpu_embedding_output_layout.pb.h" namespace tensorflow { namespace tpu { diff --git a/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.h b/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.h index c10fbeeff2b5af93a118902c0afb3b59cc1a9d60..1a04c7bdb4dcdfb2aff053609aebf0c2630e925b 100644 --- a/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.h +++ b/tensorflow/contrib/tpu/utils/tpu_embedding_output_layout_utils.h @@ -16,9 +16,9 @@ limitations under the License. #ifndef TENSORFLOW_CONTRIB_TPU_UTILS_TPU_EMBEDDING_OUTPUT_LAYOUT_UTILS_H_ #define TENSORFLOW_CONTRIB_TPU_UTILS_TPU_EMBEDDING_OUTPUT_LAYOUT_UTILS_H_ -#include "tensorflow/contrib/tpu/proto/tpu_embedding_configuration.pb.h" #include "tensorflow/core/framework/tensor_shape.pb.h" #include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/protobuf/tpu/tpu_embedding_configuration.pb.h" namespace tensorflow { namespace tpu { diff --git a/tensorflow/contrib/training/BUILD b/tensorflow/contrib/training/BUILD index f6427ae05a20f253edf030eff0f860361616042b..5bc4c3b88efd641b6f17a54753a29b0603c2b98c 100644 --- a/tensorflow/contrib/training/BUILD +++ b/tensorflow/contrib/training/BUILD @@ -264,9 +264,9 @@ py_test( py_test( name = "training_test", - size = "large", + size = "medium", srcs = ["python/training/training_test.py"], - shard_count = 3, + shard_count = 8, srcs_version = "PY2AND3", tags = ["notsan"], deps = [ diff --git a/tensorflow/contrib/training/python/training/hparam.py b/tensorflow/contrib/training/python/training/hparam.py index bcc177601b95172b05d327247bd370c2f8b65d59..27f0d9b2e38c433d4fb4573285ecb8c9946112e8 100644 --- a/tensorflow/contrib/training/python/training/hparam.py +++ b/tensorflow/contrib/training/python/training/hparam.py @@ -499,6 +499,7 @@ class HParams(object): value: New value of the hyperparameter. Raises: + KeyError: If the hyperparameter doesn't exist. ValueError: If there is a type mismatch. """ param_type, is_list = self._hparam_types[name] @@ -517,6 +518,8 @@ class HParams(object): def del_hparam(self, name): """Removes the hyperparameter with key 'name'. + Does nothing if it isn't present. + Args: name: Name of the hyperparameter. """ @@ -525,19 +528,20 @@ class HParams(object): del self._hparam_types[name] def parse(self, values): - """Override hyperparameter values, parsing new values from a string. + """Override existing hyperparameter values, parsing new values from a string. See parse_values for more detail on the allowed format for values. Args: - values: String. Comma separated list of `name=value` pairs where - 'value' must follow the syntax described above. + values: String. Comma separated list of `name=value` pairs where 'value' + must follow the syntax described above. Returns: The `HParams` instance. Raises: - ValueError: If `values` cannot be parsed. + ValueError: If `values` cannot be parsed or a hyperparameter in `values` + doesn't exist. """ type_map = dict() for name, t in self._hparam_types.items(): @@ -548,7 +552,7 @@ class HParams(object): return self.override_from_dict(values_map) def override_from_dict(self, values_dict): - """Override hyperparameter values, parsing new values from a dictionary. + """Override existing hyperparameter values, parsing new values from a dictionary. Args: values_dict: Dictionary of name:value pairs. @@ -557,6 +561,7 @@ class HParams(object): The `HParams` instance. Raises: + KeyError: If a hyperparameter in `values_dict` doesn't exist. ValueError: If `values_dict` cannot be parsed. """ for name, value in values_dict.items(): @@ -596,7 +601,7 @@ class HParams(object): sort_keys=sort_keys) def parse_json(self, values_json): - """Override hyperparameter values, parsing new values from a json object. + """Override existing hyperparameter values, parsing new values from a json object. Args: values_json: String containing a json object of name:value pairs. @@ -605,6 +610,7 @@ class HParams(object): The `HParams` instance. Raises: + KeyError: If a hyperparameter in `values_json` doesn't exist. ValueError: If `values_json` cannot be parsed. """ values_map = json.loads(values_json) diff --git a/tensorflow/contrib/training/python/training/training.py b/tensorflow/contrib/training/python/training/training.py index c272a2ac144068cfb7355c2647eebf5bd0ce9d50..4ceb6e9350f5167efc8f7266d4e748cc6fa4ffd6 100644 --- a/tensorflow/contrib/training/python/training/training.py +++ b/tensorflow/contrib/training/python/training/training.py @@ -244,7 +244,6 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.python.framework import constant_op from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import clip_ops @@ -354,11 +353,11 @@ def multiply_gradients(grads_and_vars, gradient_multipliers): raise ValueError('Requested multiple of `None` gradient.') if isinstance(grad, ops.IndexedSlices): - tmp = grad.values * constant_op.constant( + tmp = grad.values * ops.convert_to_tensor( gradient_multipliers[key], dtype=grad.dtype) grad = ops.IndexedSlices(tmp, grad.indices, grad.dense_shape) else: - grad *= constant_op.constant( + grad *= ops.convert_to_tensor( gradient_multipliers[key], dtype=grad.dtype) multiplied_grads_and_vars.append((grad, var)) return multiplied_grads_and_vars @@ -419,7 +418,7 @@ def create_train_op(total_loss, update_ops = set(update_ops) if not global_update_ops.issubset(update_ops): logging.warning('update_ops in create_train_op does not contain all the ' - ' update_ops in GraphKeys.UPDATE_OPS') + 'update_ops in GraphKeys.UPDATE_OPS') # Make sure update_ops are computed before total_loss. if update_ops: @@ -433,7 +432,7 @@ def create_train_op(total_loss, else: # Make sure that variables_to_train are in tf.trainable_variables() for v in variables_to_train: - assert v in tf_variables.trainable_variables() + assert v.trainable or v in tf_variables.trainable_variables() assert variables_to_train diff --git a/tensorflow/contrib/util/BUILD b/tensorflow/contrib/util/BUILD index d9ccda8e89a4c9a1b3f3d24915b9ad3fb4d9be5f..07dbd5ca8d65ec8232d33c016a7369c68a4c9e1f 100644 --- a/tensorflow/contrib/util/BUILD +++ b/tensorflow/contrib/util/BUILD @@ -16,9 +16,12 @@ cc_library( srcs = ["convert_graphdef_memmapped_format_lib.cc"], hdrs = ["convert_graphdef_memmapped_format_lib.h"], deps = [ + "//tensorflow/core:array_ops_op_lib", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", + "//tensorflow/core:functional_ops_op_lib", "//tensorflow/core:lib", + "//tensorflow/core:nn_ops_op_lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:tensorflow", "//tensorflow/core/kernels:immutable_constant_op", diff --git a/tensorflow/contrib/verbs/verbs_server_lib.cc b/tensorflow/contrib/verbs/verbs_server_lib.cc index 19ef109f671ee57ce2aceb55110c50aa44352223..d07fd5ae6e9cc0dbf67c6b6a4e8db086b4c74aa1 100644 --- a/tensorflow/contrib/verbs/verbs_server_lib.cc +++ b/tensorflow/contrib/verbs/verbs_server_lib.cc @@ -81,7 +81,10 @@ Status VerbsServer::ChannelCacheFactory(const ServerDef& server_def, Status VerbsServer::Init(ServiceInitFunction service_func, RendezvousMgrCreationFunction rendezvous_mgr_func) { std::call_once(reg_mem_visitors_call, []() { RdmaMgr::RegMemVisitors(); }); - Status s = GrpcServer::Init(service_func, rendezvous_mgr_func); + GrpcServerOptions opts; + opts.service_func = service_func; + opts.rendezvous_mgr_func = rendezvous_mgr_func; + Status s = GrpcServer::Init(opts); { mutex_lock l(mu_); CHECK_EQ(verbs_state_, DISCONNECTED); diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD index 1afac057d94999165874ee9dc924336a09e012d2..cc242d0e3c9fbf26874d02f3d2ab81fa0dd36584 100644 --- a/tensorflow/core/BUILD +++ b/tensorflow/core/BUILD @@ -70,10 +70,14 @@ package(default_visibility = [ licenses(["notice"]) # Apache 2.0 +# Export the BUILD file so automated tooling can check licenses +exports_files(["BUILD"]) + load( "//tensorflow:tensorflow.bzl", "cc_header_only_library", "if_android", + "if_emscripten", "if_ios", "if_linux_x86_64", "if_mobile", @@ -84,10 +88,12 @@ load( "tf_copts", "tf_cuda_library", "tf_features_nomodules_if_android", + "tf_features_nomodules_if_emscripten", "tf_gen_op_libs", "tf_generate_proto_text_sources", "tf_genrule_cmd_append_to_srcs", "tf_opts_nortti_if_android", + "tf_opts_nortti_if_emscripten", "transitive_hdrs", ) load("//tensorflow:tensorflow.bzl", "tf_cc_test_mkl") @@ -178,7 +184,6 @@ COMMON_PROTO_SRCS = [ "framework/function.proto", "framework/graph.proto", "framework/graph_transfer_info.proto", - "framework/iterator.proto", "framework/kernel_def.proto", "framework/log_memory.proto", "framework/node_def.proto", @@ -199,10 +204,12 @@ COMMON_PROTO_SRCS = [ "protobuf/cluster.proto", "protobuf/debug.proto", "protobuf/device_properties.proto", + "protobuf/graph_debug_info.proto", "protobuf/queue_runner.proto", "protobuf/rewriter_config.proto", "protobuf/tensor_bundle.proto", "protobuf/saver.proto", + "protobuf/verifier_config.proto", "util/event.proto", "util/memmapped_file_system.proto", "util/saved_tensor_slice.proto", @@ -450,10 +457,7 @@ cc_library( hdrs = ["platform/logger.h"], copts = tf_copts(), visibility = ["//visibility:public"], - deps = [ - ":lib_proto_parsing", - "@protobuf_archive//:protobuf", - ], + deps = [":lib_proto_parsing"], ) filegroup( @@ -504,6 +508,7 @@ cc_library( ":platform_port", ":platform_protobuf", "//tensorflow/core/platform/default/build_config:env", + "//tensorflow/core/platform/default/build_config:port", ], ) @@ -1017,6 +1022,7 @@ cc_library( ":lib", ":lib_internal", ":protos_all_cc", + "//tensorflow/core/util/proto:proto_utils", ], ) @@ -1074,6 +1080,7 @@ tf_gen_op_libs( "tensor_forest_ops", "candidate_sampling_ops", "checkpoint_ops", + "clustering_ops", "collective_ops", "control_flow_ops", "ctc_ops", @@ -1099,6 +1106,7 @@ tf_gen_op_libs( "parsing_ops", "random_grad", "random_ops", + "stateful_random_ops", "remote_fused_graph_ops", "rpc_ops", "scoped_allocator_ops", @@ -1133,6 +1141,13 @@ tf_gen_op_libs( deps = [":protos_all_cc"], ) +tf_gen_op_libs( + op_lib_names = [ + "mkl_array_ops", + ], + deps = [":protos_all_cc"], +) + tf_gen_op_libs( op_lib_names = [ "audio_ops", @@ -1228,6 +1243,7 @@ cc_library( ":tensor_forest_ops_op_lib", ":candidate_sampling_ops_op_lib", ":checkpoint_ops_op_lib", + ":clustering_ops_op_lib", ":collective_ops_op_lib", ":control_flow_ops_op_lib", ":ctc_ops_op_lib", @@ -1253,6 +1269,7 @@ cc_library( ":parsing_ops_op_lib", ":ragged_ops", ":random_ops_op_lib", + ":stateful_random_ops_op_lib", ":remote_fused_graph_ops_op_lib", ":resource_variable_ops_op_lib", ":rpc_ops_op_lib", @@ -1270,7 +1287,10 @@ cc_library( ":training_ops_op_lib", ":user_ops_op_lib", ":word2vec_ops", - ] + if_mkl([":mkl_nn_ops_op_lib"]) + tf_additional_cloud_op_deps(), + ] + if_mkl([ + ":mkl_array_ops_op_lib", + ":mkl_nn_ops_op_lib", + ]) + tf_additional_cloud_op_deps(), alwayslink = 1, ) @@ -1382,6 +1402,7 @@ cc_library( "//tensorflow/core/kernels:tensor_forest_ops", "//tensorflow/core/kernels:candidate_sampler_ops", "//tensorflow/core/kernels:checkpoint_ops", + "//tensorflow/core/kernels:clustering_ops", "//tensorflow/core/kernels:collective_ops", "//tensorflow/core/kernels:control_flow_ops", "//tensorflow/core/kernels:ctc_ops", @@ -1402,12 +1423,15 @@ cc_library( "//tensorflow/core/kernels:manip", "//tensorflow/core/kernels:math", "//tensorflow/core/kernels:multinomial_op", + "//tensorflow/core/kernels:mutex_ops", "//tensorflow/core/kernels:nn", "//tensorflow/core/kernels:parameterized_truncated_normal_op", "//tensorflow/core/kernels:parsing", "//tensorflow/core/kernels:partitioned_function_ops", + "//tensorflow/core/kernels:pooling_ops", "//tensorflow/core/kernels:ragged_ops", "//tensorflow/core/kernels:random_ops", + "//tensorflow/core/kernels:stateful_random_ops", "//tensorflow/core/kernels:random_poisson_op", "//tensorflow/core/kernels:remote_fused_graph_ops", "//tensorflow/core/kernels:required", @@ -1438,6 +1462,7 @@ cc_library( "//tensorflow/core/kernels:mkl_identity_op", "//tensorflow/core/kernels:mkl_input_conversion_op", "//tensorflow/core/kernels:mkl_lrn_op", + "//tensorflow/core/kernels:mkl_requantize_ops", "//tensorflow/core/kernels:mkl_pooling_ops", "//tensorflow/core/kernels:mkl_relu_op", "//tensorflow/core/kernels:mkl_reshape_op", @@ -1750,11 +1775,35 @@ cc_library( cc_library( name = "mobile_additional_lib_deps", deps = tf_additional_lib_deps() + [ + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/strings", ], ) +cc_library( + name = "emscripten_tensorflow_lib_lite_nortti_lite_protos_no_runtime", + srcs = if_emscripten(["//tensorflow/core:mobile_srcs_no_runtime"]), + copts = ["-DSUPPORT_SELECTIVE_REGISTRATION"] + tf_opts_nortti_if_emscripten(), + defines = ["TENSORFLOW_LITE_PROTOS"], + linkopts = ["-lz"], + tags = [ + "manual", + "notap", + ], + visibility = ["//visibility:public"], + deps = [ + ":emscripten_proto_lib_no_rtti_lite_runtime", + ":mobile_additional_lib_deps", + ":stats_calculator_portable", + "//third_party/eigen3", + "@double_conversion//:double-conversion", + "@nsync//:nsync_cpp", + "@zlib_archive//:zlib", + ], + alwayslink = 1, +) + # Native library support for iOS applications. # # bazel build --config=ios_x86_64 \ @@ -1950,6 +1999,14 @@ cc_library( ], ) +cc_library( + name = "rocm", + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/core/platform/default/build_config:rocm", + ], +) + # ----------------------------------------------------------------------------- # Clif-related proto libraries. @@ -2193,6 +2250,7 @@ cc_library( ], }), deps = tf_additional_lib_deps() + [ + "@com_google_absl//absl/meta:type_traits", "@com_google_absl//absl/strings", "//third_party/eigen3", "@com_google_absl//absl/base:core_headers", @@ -2207,7 +2265,6 @@ cc_library( "lib/**/*.cc", "platform/*.cc", "platform/profile_utils/**/*.cc", - ] + [ "framework/resource_handle.cc", "util/env_var.cc", ], @@ -2251,6 +2308,8 @@ cc_library( ":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", @@ -2347,7 +2406,12 @@ cc_library( cc_library( name = "tflite_portable_logging", - srcs = [], + srcs = [ + ] + if_ios([ + "platform/default/logging.cc", + "platform/env_time.cc", + "platform/posix/env_time.cc", + ]), hdrs = [ "lib/bfloat16/bfloat16.h", "platform/default/integral_types.h", @@ -2356,7 +2420,7 @@ cc_library( "platform/macros.h", "platform/platform.h", "platform/types.h", - ] + if_windows(["platform/windows/integral_types.h"]), + ] + if_windows(["platform/windows/integral_types.h"]) + if_ios(["platform/env_time.h"]), copts = tf_copts(), linkopts = ["-ldl"], deps = [ @@ -2625,7 +2689,6 @@ tf_cuda_library( "example/**/*.cc", "framework/**/*.cc", "util/**/*.cc", - ] + [ "graph/edgeset.cc", "graph/graph.cc", "graph/graph_def_builder.cc", @@ -2811,6 +2874,9 @@ tf_cuda_library( ":proto_text", ":protos_all_cc", "//third_party/eigen3", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/strings", ], ) @@ -2867,6 +2933,7 @@ tf_cuda_library( CORE_CPU_LIB_HEADERS = CORE_CPU_BASE_HDRS + [ "common_runtime/allocator_retry.h", + "common_runtime/shared_counter.h", "common_runtime/base_collective_executor.h", "common_runtime/bfc_allocator.h", "common_runtime/hierarchical_tree_broadcaster.h", @@ -3016,6 +3083,7 @@ tf_cuda_library( ":framework", ":graph", ":lib", + ":metrics", ":proto_text", ":protos_all_cc", "//tensorflow/core/grappler:grappler_item", @@ -3647,6 +3715,20 @@ tf_cc_test( ], ) +tf_cc_test( + name = "lib_strings_proto_serialization_test", + srcs = ["lib/strings/proto_serialization_test.cc"], + deps = [ + ":lib", + ":lib_internal", + ":lib_test_internal", + ":protos_all_cc", + ":test", + ":test_main", + "@com_google_absl//absl/memory", + ], +) + tf_cc_test( name = "lib_random_weighted_picker_test", size = "medium", @@ -3699,7 +3781,6 @@ tf_cc_tests( srcs = [ "common_runtime/buf_rendezvous_test.cc", "common_runtime/collective_executor_mgr_test.cc", - "common_runtime/collective_param_resolver_local_test.cc", "common_runtime/collective_rma_local_test.cc", "common_runtime/device_resolver_local_test.cc", "common_runtime/device_set_test.cc", @@ -3815,6 +3896,7 @@ tf_cc_tests( name = "higher_level_tests_needing_kernels", size = "small", srcs = [ + "common_runtime/collective_param_resolver_local_test.cc", "graph/graph_constructor_test.cc", ], linkopts = select({ @@ -4452,7 +4534,7 @@ tf_cc_test( "//tensorflow/cc:scope", "//tensorflow/core/kernels:cwise_op", "//third_party/eigen3", - ], + ] + if_mkl([":mkl_array_ops_op_lib"]), ) tf_cc_test( @@ -5005,6 +5087,39 @@ transitive_hdrs( # ----------------------------------------------------------------------------- # Google-internal targets go here (must be at the end). +load("//tensorflow:tensorflow.bzl", "tf_portable_proto_library") + +genrule( + name = "emscripten_proto_config_lite_runtime", + outs = ["emscripten_proto_config_lite_runtime.asciipb"], + cmd = tf_genrule_cmd_append_to_srcs("optimize_mode:LITE_RUNTIME"), + visibility = ["//visibility:private"], +) + +# We are keeping the "android" version of tf_android_core_proto_headers. All it does is +# normalize CORE_PROTO_SRCS to generate valid output file names. +tf_portable_proto_library( + name = "emscripten_proto_lib_no_rtti_lite_runtime", + config = ":emscripten_proto_config_lite_runtime", + copts = tf_opts_nortti_if_emscripten(), + features = tf_features_nomodules_if_emscripten(), + header_outs = tf_android_core_proto_headers(CORE_PROTO_SRCS) + ["//google/protobuf/any.proto.h"], + link_full_protobuf = False, + prefix_dir = "emscripten_proto_no_rtti", + proto_deps = [ + ":protos_all_cc", + "@protobuf_archive//:protobuf", + ], + visibility = ["//visibility:public"], +) + +# There is currently no need for a full proto version of emscripten tf lib lite. +alias( + name = "emscripten_lib_lite_no_runtime", + actual = "//tensorflow/core:emscripten_tensorflow_lib_lite_nortti_lite_protos_no_runtime", + visibility = ["//visibility:public"], +) + alias( name = "android_srcs_no_runtime", actual = ":mobile_srcs_no_runtime", diff --git a/tensorflow/core/api_def/base_api/api_def_Case.pbtxt b/tensorflow/core/api_def/base_api/api_def_Case.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..56fef3ae6d0d452cb2caa57c36f35a04584864ee --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_Case.pbtxt @@ -0,0 +1,43 @@ +op { + graph_op_name: "Case" + in_arg { + name: "branch_index" + description: "The branch selector, an int32 Tensor." + } + in_arg { + name: "input" + description: "A list of input tensors passed to the branch function." + } + out_arg { + name: "output" + description: "A list of return values." + } + attr { name: "Tin" description: "A list of input types." } + attr { name: "Tout" description: "A list of output types." } + attr { + name: "branches" + description: <